rework semantic highlighting to use cursor finders

server/src/parser/slangHelpers.ts

@@ -1,5 +1,4 @@
 import { NodeType, RuleNode, TokenNode } from "@nomicfoundation/slang/cst";
-import { Cursor } from "@nomicfoundation/slang/cursor";
 import { RuleKind, TokenKind } from "@nomicfoundation/slang/kinds";
 import { TextRange } from "@nomicfoundation/slang/text_index";
 import _ from "lodash";
@@ -19,33 +18,6 @@ export interface SlangNodeWrapper {
   pathRuleNodes: SlangNode[];
 }
 
-export function walk(
-  cursor: Cursor,
-  onEnter: NodeCallback,
-  onExit: NodeCallback
-) {
-  const node = cursor.node;
-
-  const nodeWrapper: SlangNodeWrapper = {
-    textRange: cursor.textRange,
-    type: node.type,
-    kind: node.kind,
-    text: node.text,
-    pathRuleNodes: cursor.pathRuleNodes,
-  };
-
-  onEnter(nodeWrapper);
-
-  if (nodeWrapper.type === NodeType.Rule) {
-    for (let i = 0; i < node.children.length; i++) {
-      cursor.goToNthChild(i);
-      walk(cursor, onEnter, onExit);
-    }
-  }
-  onExit(nodeWrapper);
-  cursor.goToParent();
-}
-
 export function slangToVSCodeRange(
   doc: TextDocument,
   slangRange: TextRange

server/src/services/semanticHighlight/HighlightVisitor.ts

@@ -1,16 +1,18 @@
 /* eslint-disable @typescript-eslint/no-empty-function */
 /* eslint-disable @typescript-eslint/no-unused-vars */
 import { TextDocument } from "vscode-languageserver-textdocument";
+import { TokenKind } from "@nomicfoundation/slang/kinds";
 import { SlangNodeWrapper } from "../../parser/slangHelpers";
 import { SemanticTokensBuilder } from "./SemanticTokensBuilder";
 
 // Abstraction for a visitor that wants to highlight tokens
 export abstract class HighlightVisitor {
+  public abstract tokenKinds: Set<TokenKind>;
+
   constructor(
     public document: TextDocument,
     public tokenBuilder: SemanticTokensBuilder
   ) {}
 
   public enter(nodeWrapper: SlangNodeWrapper): void {}
-  public exit(nodeWrapper: SlangNodeWrapper): void {}
 }

server/src/services/semanticHighlight/highlighters/ContractDefinitionHighlighter.ts

@@ -1,11 +1,14 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
 import { RuleKind, TokenKind } from "@nomicfoundation/slang/kinds";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
 // Highlights contract definitions
 export class ContractDefinitionHighlighter extends HighlightVisitor {
+  public tokenKinds = new Set([TokenKind.Identifier]);
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     const ancestors = nodeWrapper.pathRuleNodes;
     if (

server/src/services/semanticHighlight/highlighters/CustomTypeHighlighter.ts

@@ -1,11 +1,14 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
 import { RuleKind, TokenKind } from "@nomicfoundation/slang/kinds";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
 // Highlights custom type names
 export class CustomTypeHighlighter extends HighlightVisitor {
+  public tokenKinds = new Set([TokenKind.Identifier]);
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     const ancestors = nodeWrapper.pathRuleNodes;
     if (

server/src/services/semanticHighlight/highlighters/EventDefinitionHighlighter.ts

@@ -1,11 +1,14 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
 import { RuleKind, TokenKind } from "@nomicfoundation/slang/kinds";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
 // Highlights event definitions
 export class EventDefinitionHighlighter extends HighlightVisitor {
+  public tokenKinds = new Set([TokenKind.Identifier]);
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     const ancestors = nodeWrapper.pathRuleNodes;
     if (

server/src/services/semanticHighlight/highlighters/EventEmissionHighlighter.ts

@@ -1,11 +1,14 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
 import { RuleKind, TokenKind } from "@nomicfoundation/slang/kinds";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
 // Highlights event emissions
 export class EventEmissionHighlighter extends HighlightVisitor {
+  public tokenKinds = new Set([TokenKind.Identifier]);
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     const ancestors = nodeWrapper.pathRuleNodes;
     if (

server/src/services/semanticHighlight/highlighters/FunctionCallHighlighter.ts

@@ -1,11 +1,14 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
 import { RuleKind, TokenKind } from "@nomicfoundation/slang/kinds";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
 // Highlights function calls
 export class FunctionCallHighlighter extends HighlightVisitor {
+  public tokenKinds = new Set([TokenKind.Identifier]);
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     const ancestors = nodeWrapper.pathRuleNodes;
     if (

server/src/services/semanticHighlight/highlighters/FunctionDefinitionHighlighter.ts

@@ -1,11 +1,14 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
 import { RuleKind, TokenKind } from "@nomicfoundation/slang/kinds";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
 // Highlights function definitions
 export class FunctionDefinitionHighlighter extends HighlightVisitor {
+  public tokenKinds = new Set([TokenKind.Identifier]);
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     const ancestors = nodeWrapper.pathRuleNodes;
     if (

server/src/services/semanticHighlight/highlighters/InterfaceDefinitionHighlighter.ts

@@ -1,11 +1,14 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
 import { RuleKind, TokenKind } from "@nomicfoundation/slang/kinds";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
 // Highlights interface definitions
 export class InterfaceDefinitionHighlighter extends HighlightVisitor {
+  public tokenKinds = new Set([TokenKind.Identifier]);
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     const ancestors = nodeWrapper.pathRuleNodes;
     if (

server/src/services/semanticHighlight/highlighters/KeywordHighlighter.ts

@@ -1,6 +1,6 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
 import { TokenKind } from "@nomicfoundation/slang/kinds";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
 import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
@@ -91,6 +91,8 @@ const keywordKinds = new Set([
 
 // Highlights keywords
 export class KeywordHighlighter extends HighlightVisitor {
+  public tokenKinds = keywordKinds;
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     if (
       nodeWrapper.type === NodeType.Token &&

server/src/services/semanticHighlight/highlighters/NumberHighlighter.ts

@@ -1,6 +1,7 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
 import { TokenKind } from "@nomicfoundation/slang/kinds";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
@@ -12,6 +13,8 @@ const numberKinds = new Set([
 
 // Highlights numbers
 export class NumberHighlighter extends HighlightVisitor {
+  public tokenKinds = numberKinds;
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     if (
       nodeWrapper.type === NodeType.Token &&

server/src/services/semanticHighlight/highlighters/StringHighlighter.ts

@@ -1,6 +1,7 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
 import { TokenKind } from "@nomicfoundation/slang/kinds";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
@@ -12,6 +13,8 @@ const stringKinds = new Set([
 
 // Highlights strings
 export class StringHighlighter extends HighlightVisitor {
+  public tokenKinds = stringKinds;
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     if (
       nodeWrapper.type === NodeType.Token &&

server/src/services/semanticHighlight/highlighters/StructDefinitionHighlighter.ts

@@ -1,11 +1,14 @@
 import { SemanticTokenTypes } from "vscode-languageserver-protocol";
-import { NodeType } from "@nomicfoundation/slang/cst";
+import { NodeType, TokenNode } from "@nomicfoundation/slang/cst";
 import { RuleKind, TokenKind } from "@nomicfoundation/slang/kinds";
+import { Cursor } from "@nomicfoundation/slang/cursor";
 import { HighlightVisitor } from "../HighlightVisitor";
 import { SlangNodeWrapper } from "../../../parser/slangHelpers";
 
 // Highlights struct definitions
 export class StructDefinitionHighlighter extends HighlightVisitor {
+  public tokenKinds = new Set([TokenKind.Identifier]);
+
   public enter(nodeWrapper: SlangNodeWrapper): void {
     const ancestors = nodeWrapper.pathRuleNodes;
     if (

server/src/services/semanticHighlight/onSemanticTokensFull.ts

@@ -5,13 +5,14 @@ import {
   SemanticTokens,
   SemanticTokensParams,
 } from "vscode-languageserver-protocol";
-import _ from "lodash";
+import _, { Dictionary } from "lodash";
 import { analyze } from "@nomicfoundation/solidity-analyzer";
 import semver from "semver";
 import { Language } from "@nomicfoundation/slang/language";
-import { ProductionKind } from "@nomicfoundation/slang/kinds";
+import { ProductionKind, TokenKind } from "@nomicfoundation/slang/kinds";
+import { Cursor } from "@nomicfoundation/slang/cursor";
+import { TokenNode } from "@nomicfoundation/slang/cst";
 import { ServerState } from "../../types";
-import { walk } from "../../parser/slangHelpers";
 import { CustomTypeHighlighter } from "./highlighters/CustomTypeHighlighter";
 import { SemanticTokensBuilder } from "./SemanticTokensBuilder";
 import { KeywordHighlighter } from "./highlighters/KeywordHighlighter";
@@ -24,6 +25,7 @@ import { EventDefinitionHighlighter } from "./highlighters/EventDefinitionHighli
 import { ContractDefinitionHighlighter } from "./highlighters/ContractDefinitionHighlighter";
 import { InterfaceDefinitionHighlighter } from "./highlighters/InterfaceDefinitionHighlighter";
 import { StructDefinitionHighlighter } from "./highlighters/StructDefinitionHighlighter";
+import { HighlightVisitor } from "./HighlightVisitor";
 
 const emptyResponse: SemanticTokens = { data: [] };
 
@@ -103,20 +105,42 @@ export function onSemanticTokensFull(serverState: ServerState) {
           ];
 
           // Visit the CST
-          span = transaction.startChild({ op: "walk-highlight-tokens" });
-          walk(
-            parseTree.cursor,
-            (nodeWrapper) => {
-              for (const visitor of visitors) {
-                visitor.enter(nodeWrapper);
-              }
-            },
-            (nodeWrapper) => {
-              for (const visitor of visitors) {
-                visitor.exit(nodeWrapper);
+          const indexedVisitors: Dictionary<HighlightVisitor[]> = {};
+          const registeredTokenKinds: TokenKind[] = [];
+
+          for (const visitor of visitors) {
+            for (const tokenKind of visitor.tokenKinds) {
+              indexedVisitors[tokenKind] ||= [];
+              indexedVisitors[tokenKind].push(visitor);
+
+              if (!registeredTokenKinds.includes(tokenKind)) {
+                registeredTokenKinds.push(tokenKind);
               }
             }
-          );
+          }
+
+          const cursor: Cursor = parseTree.cursor;
+          let node: TokenNode;
+
+          span = transaction.startChild({ op: "walk-highlight-tokens" });
+          while (
+            (node = cursor.findTokenWithKind(registeredTokenKinds)) !== null
+          ) {
+            const nodeWrapper = {
+              kind: node.kind,
+              pathRuleNodes: cursor.pathRuleNodes,
+              text: node.text,
+              textRange: cursor.textRange,
+              type: node.type,
+            };
+
+            const registeredVisitors = indexedVisitors[nodeWrapper.kind];
+            for (const visitor of registeredVisitors) {
+              visitor.enter(nodeWrapper);
+            }
+
+            cursor.goToNext();
+          }
           span.finish();
 
           return { status: "ok", result: { data: builder.getTokenData() } };