From 3bd5d1e3ac98d56020dfeb089f3c9f7691cfd1e1 Mon Sep 17 00:00:00 2001
From: ice1000 <ice1000kotlin@foxmail.com>
Date: Sat, 4 Jan 2025 11:36:14 -0500
Subject: [PATCH] highlight: highlight import as two segments: the prefix and
 the last element

---
 .../java/org/aya/producer/AyaProducer.java    | 21 +++++++++++++++----
 .../org/aya/syntax/concrete/stmt/Command.java | 12 ++++++++---
 .../aya/syntax/concrete/stmt/StmtVisitor.java |  7 +++++--
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/producer/src/main/java/org/aya/producer/AyaProducer.java b/producer/src/main/java/org/aya/producer/AyaProducer.java
index 38908b57c..b70ad0f48 100644
--- 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
index 9d6c083ee..31dcc4e0a 100644
--- 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
index e04fc10f7..a8d44c4cc 100644
--- 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`