highlight: highlight import as two segments: the prefix and the last …

…element
aya-prover · Jan 4, 2025 · 3bd5d1e · 3bd5d1e
1 parent df382d5
commit 3bd5d1e
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 9 deletions.
diff --git a/producer/src/main/java/org/aya/producer/AyaProducer.java b/producer/src/main/java/org/aya/producer/AyaProducer.java
@@ -35,6 +35,7 @@
 import org.aya.syntax.ref.LocalVar;
 import org.aya.syntax.ref.ModulePath;
 import org.aya.util.Arg;
+import org.aya.util.Pair;
 import org.aya.util.binop.Assoc;
 import org.aya.util.binop.OpDecl;
 import org.aya.util.error.Panic;
@@ -113,12 +114,21 @@ public record AyaProducer(
       type(node.child(TYPE)));
   }
 
+  public @NotNull Pair<SourcePos, SourcePos> importQualifiedPos(@NotNull GenericNode<?> qualifiedId) {
+    var ids = qualifiedId.childrenOfType(WEAK_ID);
+    var prefix = ids.dropLast(1).toImmutableSeq();
+    var prefixPos = prefix.isEmpty() ? SourcePos.NONE : sourcePosOf(prefix.get(0)).union(sourcePosOf(prefix.getLast()));
+    return new Pair<>(prefixPos, sourcePosOf(ids.getLast()));
+  }
+
   public @NotNull Command.Import importCmd(@NotNull GenericNode<?> node) {
     var acc = node.peekChild(KW_PUBLIC);
     var asId = node.peekChild(WEAK_ID);
     var importMod = node.child(QUALIFIED_ID);
+    var importModPos = importQualifiedPos(importMod);
     return new Command.Import(
-      sourcePosOf(importMod),
+      importModPos.component1(),
+      importModPos.component2(),
       modulePath(importMod),
       asId == null ? null : weakId(asId),
       acc == null ? Stmt.Accessibility.Private : Stmt.Accessibility.Public
@@ -142,9 +152,12 @@ public record AyaProducer(
       false,
       openImport
     );
-    return openImport
-      ? ImmutableSeq.of(new Command.Import(namePos, modName, null, accessibility), open)
-      : ImmutableSeq.of(open);
+    if (openImport) {
+      var importPos = importQualifiedPos(modNameNode);
+      return ImmutableSeq.of(new Command.Import(importPos.component1(), importPos.component2(),
+        modName, null, accessibility), open);
+    }
+    return ImmutableSeq.of(open);
   }
 
   public UseHide hideList(SeqView<? extends GenericNode<?>> hideLists, UseHide.Strategy strategy) {

diff --git a/syntax/src/main/java/org/aya/syntax/concrete/stmt/Command.java b/syntax/src/main/java/org/aya/syntax/concrete/stmt/Command.java
@@ -14,13 +14,19 @@
 
 public sealed interface Command extends Stmt {
   default @Override void descentInPlace(@NotNull PosedUnaryOperator<Expr> f, @NotNull PosedUnaryOperator<Pattern> p) { }
-  /** @author re-xyr */
+  /// @param sourcePosExceptLast can be NONE if the entire import is one `weakId`
   record Import(
-    @Override @NotNull SourcePos sourcePos,
+    @NotNull SourcePos sourcePosExceptLast,
+    @NotNull SourcePos sourcePosLast,
     @NotNull ModulePath path,
     @Nullable WithPos<String> asName,
     @Override @NotNull Accessibility accessibility
-  ) implements Command { }
+  ) implements Command {
+    @Override public @NotNull SourcePos sourcePos() {
+      if (sourcePosExceptLast == SourcePos.NONE) return sourcePosLast;
+      return sourcePosExceptLast.union(sourcePosLast);
+    }
+  }
 
   /** @author re-xyr */
   record Open(

diff --git a/syntax/src/main/java/org/aya/syntax/concrete/stmt/StmtVisitor.java b/syntax/src/main/java/org/aya/syntax/concrete/stmt/StmtVisitor.java
@@ -77,11 +77,14 @@ private void visitVars(@NotNull Stmt stmt) {
       case Generalize g -> g.variables.forEach(v -> visitVarDecl(v.sourcePos, v, noType));
       case Command.Module m -> visitModuleDecl(m.sourcePos(), ModuleName.of(m.name()));
       case Command.Import i -> {
-        visitModuleRef(i.sourcePos(), i.path());
+        // Essentially `i.asName() != null` but fancier
+        var path = i.path();
         if (i.asName() instanceof WithPos(var pos, var asName)) {
+          visitModuleRef(i.sourcePos(), path);
           visitModuleDecl(pos, ModuleName.of(asName));
         } else {
-          // TODO: visitModuleDecl on the last element of i.path
+          if (i.sourcePosExceptLast() != SourcePos.NONE) visitModuleRef(i.sourcePosExceptLast(), path.dropLast(1));
+          visitModuleDecl(i.sourcePosLast(), ModuleName.of(path.last()));
         }
       }
       case Command.Open o when o.fromSugar() -> { }  // handled in `case Decl` or `case Command.Import`