vscode-extension/.gitignore

@@ -1,4 +1,4 @@
 node_modules
 client/dist
 server/dist
-*.vsix
+*.vsix

vscode-extension/server/src/semanticTokens.ts

@@ -2,7 +2,11 @@ import { parser } from '../../../src/parser/shrimp'
 import * as Terms from '../../../src/parser/shrimp.terms'
 import { SyntaxNode } from '@lezer/common'
 import { TextDocument } from 'vscode-languageserver-textdocument'
-import { SemanticTokensBuilder, SemanticTokenTypes } from 'vscode-languageserver/node'
+import {
+  SemanticTokensBuilder,
+  SemanticTokenTypes,
+  SemanticTokenModifiers,
+} from 'vscode-languageserver/node'
 
 export const TOKEN_TYPES = [
   SemanticTokenTypes.function,
@@ -14,9 +18,14 @@ export const TOKEN_TYPES = [
   SemanticTokenTypes.parameter,
   SemanticTokenTypes.property,
   SemanticTokenTypes.regexp,
+  SemanticTokenTypes.comment,
 ]
 
-export const TOKEN_MODIFIERS: string[] = []
+export const TOKEN_MODIFIERS = [
+  SemanticTokenModifiers.declaration,
+  SemanticTokenModifiers.modification,
+  SemanticTokenModifiers.readonly,
+]
 
 export function buildSemanticTokens(document: TextDocument): number[] {
   const text = document.getText()
@@ -30,12 +39,12 @@ export function buildSemanticTokens(document: TextDocument): number[] {
 
 // Walk the tree and collect tokens
 function walkTree(node: SyntaxNode, document: TextDocument, builder: SemanticTokensBuilder) {
-  const tokenType = getTokenType(node.type.id)
+  const tokenInfo = getTokenType(node.type.id, node.parent?.type.id)
 
-  if (tokenType !== undefined) {
+  if (tokenInfo !== undefined) {
     const start = document.positionAt(node.from)
     const length = node.to - node.from
-    builder.push(start.line, start.character, length, tokenType, 0)
+    builder.push(start.line, start.character, length, tokenInfo.type, tokenInfo.modifiers)
   }
 
   let child = node.firstChild
@@ -45,25 +54,54 @@ function walkTree(node: SyntaxNode, document: TextDocument, builder: SemanticTok
   }
 }
 
-// Map Lezer node IDs to semantic token type indices
-function getTokenType(nodeTypeId: number): number | undefined {
+// Map Lezer node IDs to semantic token type indices and modifiers
+function getTokenType(nodeTypeId: number, parentTypeId?: number): { type: number; modifiers: number } | undefined {
   switch (nodeTypeId) {
-    case Terms.FunctionCall:
-    case Terms.FunctionDef:
-      return TOKEN_TYPES.indexOf(SemanticTokenTypes.function)
-
     case Terms.Identifier:
+      // Check parent to determine if this identifier is a function call or variable
+      if (parentTypeId === Terms.FunctionCall) {
+        return {
+          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.function),
+          modifiers: 0,
+        }
+      }
+      if (parentTypeId === Terms.FunctionDef) {
+        return {
+          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.function),
+          modifiers: getModifierBits(SemanticTokenModifiers.declaration),
+        }
+      }
+      if (parentTypeId === Terms.FunctionCallOrIdentifier) {
+        return {
+          type: TOKEN_TYPES.indexOf(SemanticTokenTypes.function),
+          modifiers: 0,
+        }
+      }
+      // Otherwise it's a regular variable
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.variable),
+        modifiers: 0,
+      }
+
     case Terms.AssignableIdentifier:
-    case Terms.FunctionCallOrIdentifier:
-      return TOKEN_TYPES.indexOf(SemanticTokenTypes.variable)
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.variable),
+        modifiers: getModifierBits(SemanticTokenModifiers.modification),
+      }
 
     case Terms.String:
     case Terms.StringFragment:
     case Terms.Word:
-      return TOKEN_TYPES.indexOf(SemanticTokenTypes.string)
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.string),
+        modifiers: 0,
+      }
 
     case Terms.Number:
-      return TOKEN_TYPES.indexOf(SemanticTokenTypes.number)
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.number),
+        modifiers: 0,
+      }
 
     case Terms.Plus:
     case Terms.Minus:
@@ -79,23 +117,59 @@ function getTokenType(nodeTypeId: number): number | undefined {
     case Terms.Modulo:
     case Terms.And:
     case Terms.Or:
-      return TOKEN_TYPES.indexOf(SemanticTokenTypes.operator)
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.operator),
+        modifiers: 0,
+      }
 
     case Terms.keyword:
     case Terms.Do:
-      return TOKEN_TYPES.indexOf(SemanticTokenTypes.keyword)
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.keyword),
+        modifiers: 0,
+      }
 
     case Terms.Params:
     case Terms.NamedParam:
-      return TOKEN_TYPES.indexOf(SemanticTokenTypes.parameter)
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.parameter),
+        modifiers: 0,
+      }
 
     case Terms.DotGet:
-      return TOKEN_TYPES.indexOf(SemanticTokenTypes.property)
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.property),
+        modifiers: 0,
+      }
 
     case Terms.Regex:
-      return TOKEN_TYPES.indexOf(SemanticTokenTypes.regexp)
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.regexp),
+        modifiers: 0,
+      }
+
+    case Terms.Comment:
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.comment),
+        modifiers: 0,
+      }
+
+    case Terms.NamedArg:
+      return {
+        type: TOKEN_TYPES.indexOf(SemanticTokenTypes.property),
+        modifiers: 0,
+      }
 
     default:
       return undefined
   }
 }
+
+const getModifierBits = (...modifiers: SemanticTokenModifiers[]): number => {
+  let bits = 0
+  for (const modifier of modifiers) {
+    const index = TOKEN_MODIFIERS.indexOf(modifier)
+    if (index !== -1) bits |= 1 << index
+  }
+  return bits
+}

vscode-extension/server/src/server.ts

@@ -86,27 +86,38 @@ function handleParseTree(params: { uri: string }) {
 
   const text = document.getText()
   const tree = parser.parse(text)
-  const treeString = tree.toString()
+  const cursor = tree.cursor()
 
-  // Format with indentation, without parentheses
   let formatted = ''
-  let indent = 0
-  for (let i = 0; i < treeString.length; i++) {
-    const char = treeString[i]
-    if (char === '(') {
-      formatted += '\n'
-      indent++
-      formatted += '  '.repeat(indent)
-    } else if (char === ')') {
-      indent--
-    } else if (char === ',') {
-      formatted += '\n'
-      formatted += '  '.repeat(indent)
-    } else {
-      formatted += char
+  let depth = 0
+
+  const printNode = () => {
+    const nodeName = cursor.name
+    const nodeText = text.slice(cursor.from, cursor.to)
+    const indent = '  '.repeat(depth)
+
+    formatted += `${indent}${nodeName}`
+    if (nodeText) {
+      const escapedText = nodeText.replace(/\n/g, '\\n').replace(/\r/g, '\\r')
+      formatted += ` "${escapedText}"`
+    }
+    formatted += '\n'
+  }
+
+  const traverse = (): void => {
+    printNode()
+
+    if (cursor.firstChild()) {
+      depth++
+      do {
+        traverse()
+      } while (cursor.nextSibling())
+      cursor.parent()
+      depth--
     }
   }
 
+  traverse()
   return formatted
 }