Track that scope!

vscode-extension/server/src/scopeTracker.test.ts

@@ -0,0 +1,145 @@
+import { test, expect, describe } from 'bun:test'
+import { ScopeTracker } from './scopeTracker'
+import { TextDocument } from 'vscode-languageserver-textdocument'
+import { parser } from '../../../src/parser/shrimp'
+import * as Terms from '../../../src/parser/shrimp.terms'
+
+describe('ScopeTracker', () => {
+  test('top-level assignment is in scope', () => {
+    const code = 'x = 5\necho x'
+    const { tree, tracker } = parseAndGetScope(code)
+
+    // Find the 'x' identifier in 'echo x'
+    const identifiers: any[] = []
+    tree.topNode.cursor().iterate((node: any) => {
+      if (node.type.id === Terms.Identifier) {
+        identifiers.push(node.node)
+      }
+    })
+
+    // Second identifier should be the 'x' in 'echo x'
+    const xInEcho = identifiers[1]
+    expect(xInEcho).toBeDefined()
+    expect(tracker.isInScope('x', xInEcho)).toBe(true)
+  })
+
+  test('undeclared variable is not in scope', () => {
+    const code = 'echo x'
+    const { tree, tracker } = parseAndGetScope(code)
+
+    // Find the 'x' identifier
+    let xNode: any = null
+    tree.topNode.cursor().iterate((node: any) => {
+      if (node.type.id === Terms.Identifier) {
+        xNode = node.node
+      }
+    })
+
+    expect(xNode).toBeDefined()
+    expect(tracker.isInScope('x', xNode)).toBe(false)
+  })
+
+  test('function parameter is in scope inside function', () => {
+    const code = `greet = do name:
+  echo name
+end`
+    const { tree, tracker } = parseAndGetScope(code)
+
+    // Find all identifiers
+    const identifiers: any[] = []
+    tree.topNode.cursor().iterate((node: any) => {
+      if (node.type.id === Terms.Identifier) {
+        identifiers.push(node.node)
+      }
+    })
+
+    // Find the 'name' in 'echo name' (should be last identifier)
+    const nameInEcho = identifiers[identifiers.length - 1]
+    expect(tracker.isInScope('name', nameInEcho)).toBe(true)
+  })
+
+  test('assignment before usage is in scope', () => {
+    const code = `x = 5
+y = 10
+echo x y`
+    const { tree, tracker } = parseAndGetScope(code)
+
+    // Find identifiers
+    const identifiers: any[] = []
+    tree.topNode.cursor().iterate((node: any) => {
+      if (node.type.id === Terms.Identifier) {
+        identifiers.push(node.node)
+      }
+    })
+
+    // Last two identifiers should be 'x' and 'y' in 'echo x y'
+    const xInEcho = identifiers[identifiers.length - 2]
+    const yInEcho = identifiers[identifiers.length - 1]
+
+    expect(tracker.isInScope('x', xInEcho)).toBe(true)
+    expect(tracker.isInScope('y', yInEcho)).toBe(true)
+  })
+
+  test('assignment after usage is not in scope', () => {
+    const code = `echo x
+x = 5`
+    const { tree, tracker } = parseAndGetScope(code)
+
+    // Find the first 'x' identifier (in echo)
+    let xNode: any = null
+    tree.topNode.cursor().iterate((node: any) => {
+      if (node.type.id === Terms.Identifier && !xNode) {
+        xNode = node.node
+      }
+    })
+
+    expect(tracker.isInScope('x', xNode)).toBe(false)
+  })
+
+  test('nested function has access to outer scope', () => {
+    const code = `x = 5
+greet = do:
+  echo x
+end`
+    const { tree, tracker } = parseAndGetScope(code)
+
+    // Find all identifiers
+    const identifiers: any[] = []
+    tree.topNode.cursor().iterate((node: any) => {
+      if (node.type.id === Terms.Identifier) {
+        identifiers.push(node.node)
+      }
+    })
+
+    // Find the 'x' in 'echo x' (should be last identifier)
+    const xInEcho = identifiers[identifiers.length - 1]
+    expect(tracker.isInScope('x', xInEcho)).toBe(true)
+  })
+
+  test('inner function parameter shadows outer variable', () => {
+    const code = `x = 5
+greet = do x:
+  echo x
+end`
+    const { tree, tracker } = parseAndGetScope(code)
+
+    // Find all identifiers
+    const identifiers: any[] = []
+    tree.topNode.cursor().iterate((node: any) => {
+      if (node.type.id === Terms.Identifier) {
+        identifiers.push(node.node)
+      }
+    })
+
+    // The 'x' in 'echo x' should have 'x' in scope (from parameter)
+    const xInEcho = identifiers[identifiers.length - 1]
+    expect(tracker.isInScope('x', xInEcho)).toBe(true)
+  })
+})
+
+const parseAndGetScope = (code: string) => {
+  const document = TextDocument.create('test://test.sh', 'shrimp', 1, code)
+  const tree = parser.parse(code)
+  const tracker = new ScopeTracker(document)
+  return { document, tree, tracker }
+}

vscode-extension/server/src/scopeTracker.ts

@@ -0,0 +1,135 @@
+import { SyntaxNode } from '@lezer/common'
+import { TextDocument } from 'vscode-languageserver-textdocument'
+import * as Terms from '../../../src/parser/shrimp.terms'
+
+/**
+ * Tracks variables in scope at a given position in the parse tree.
+ * Used to distinguish identifiers (in scope) from words (not in scope).
+ */
+export class ScopeTracker {
+  private document: TextDocument
+  private scopeCache = new Map<number, Set<string>>()
+
+  constructor(document: TextDocument) {
+    this.document = document
+  }
+
+  /**
+   * Check if a name is in scope at the given node's position.
+   */
+  isInScope(name: string, node: SyntaxNode): boolean {
+    const scope = this.getScopeAt(node)
+    return scope.has(name)
+  }
+
+  /**
+   * Get all variables in scope at the given node's position.
+   */
+  private getScopeAt(node: SyntaxNode): Set<string> {
+    const position = node.from
+
+    // Check cache first
+    if (this.scopeCache.has(position)) {
+      return this.scopeCache.get(position)!
+    }
+
+    const scope = new Set<string>()
+
+    // Find all containing function definitions
+    const containingFunctions = this.findContainingFunctions(node)
+
+    // Collect scope from each containing function (inner to outer)
+    for (const fnNode of containingFunctions) {
+      this.collectParams(fnNode, scope)
+      this.collectAssignments(fnNode, position, scope)
+    }
+
+    // Collect top-level assignments
+    const root = this.getRoot(node)
+    this.collectAssignments(root, position, scope)
+
+    this.scopeCache.set(position, scope)
+    return scope
+  }
+
+  /**
+   * Find all function definitions that contain the given node.
+   */
+  private findContainingFunctions(node: SyntaxNode): SyntaxNode[] {
+    const functions: SyntaxNode[] = []
+    let current = node.parent
+
+    while (current) {
+      if (current.type.id === Terms.FunctionDef) {
+        functions.unshift(current) // Add to beginning for outer-to-inner order
+      }
+      current = current.parent
+    }
+
+    return functions
+  }
+
+  /**
+   * Get the root node of the tree.
+   */
+  private getRoot(node: SyntaxNode): SyntaxNode {
+    let current = node
+    while (current.parent) {
+      current = current.parent
+    }
+    return current
+  }
+
+  /**
+   * Collect parameter names from a function definition.
+   */
+  private collectParams(fnNode: SyntaxNode, scope: Set<string>) {
+    let child = fnNode.firstChild
+    while (child) {
+      if (child.type.id === Terms.Params) {
+        let param = child.firstChild
+        while (param) {
+          if (param.type.id === Terms.Identifier) {
+            const text = this.document.getText({
+              start: this.document.positionAt(param.from),
+              end: this.document.positionAt(param.to),
+            })
+            scope.add(text)
+          }
+          param = param.nextSibling
+        }
+        break
+      }
+      child = child.nextSibling
+    }
+  }
+
+  /**
+   * Collect assignment names from a scope node that occur before the given position.
+   */
+  private collectAssignments(scopeNode: SyntaxNode, beforePosition: number, scope: Set<string>) {
+    const cursor = scopeNode.cursor()
+
+    cursor.iterate((node) => {
+      // Stop if we've passed the position we're checking
+      if (node.from >= beforePosition) return false
+
+      if (node.type.id === Terms.Assign) {
+        const assignNode = node.node
+        const child = assignNode.firstChild
+        if (child?.type.id === Terms.AssignableIdentifier) {
+          const text = this.document.getText({
+            start: this.document.positionAt(child.from),
+            end: this.document.positionAt(child.to),
+          })
+          scope.add(text)
+        }
+      }
+
+      // Don't descend into nested functions unless it's the current scope
+      if (node.type.id === Terms.FunctionDef && node.node !== scopeNode) {
+        return false
+      }
+    })
+  }
+}

vscode-extension/server/src/semanticTokens.ts

@@ -7,6 +7,7 @@ import {
   SemanticTokenTypes,
   SemanticTokenModifiers,
 } from 'vscode-languageserver/node'
+import { ScopeTracker } from './scopeTracker'
 
 export const TOKEN_TYPES = [
   SemanticTokenTypes.function,
@@ -31,15 +32,21 @@ export function buildSemanticTokens(document: TextDocument): number[] {
   const text = document.getText()
   const tree = parser.parse(text)
   const builder = new SemanticTokensBuilder()
+  const scopeTracker = new ScopeTracker(document)
 
-  walkTree(tree.topNode, document, builder)
+  walkTree(tree.topNode, document, builder, scopeTracker)
 
   return builder.build().data
 }
 
 // Walk the tree and collect tokens
-function walkTree(node: SyntaxNode, document: TextDocument, builder: SemanticTokensBuilder) {
+function walkTree(
-  const tokenInfo = getTokenType(node.type.id, node.parent?.type.id)
+  node: SyntaxNode,
+  document: TextDocument,
+  builder: SemanticTokensBuilder,
+  scopeTracker: ScopeTracker
+) {
+  const tokenInfo = getTokenType(node, document, scopeTracker)
 
   if (tokenInfo !== undefined) {
     const start = document.positionAt(node.from)
@@ -49,16 +56,29 @@ function walkTree(node: SyntaxNode, document: TextDocument, builder: SemanticTok
 
   let child = node.firstChild
   while (child) {
-    walkTree(child, document, builder)
+    walkTree(child, document, builder, scopeTracker)
     child = child.nextSibling
   }
 }
 
 // Map Lezer node IDs to semantic token type indices and modifiers
+type TokenInfo = { type: number; modifiers: number } | undefined
 function getTokenType(
-  nodeTypeId: number,
+  node: SyntaxNode,
-  parentTypeId?: number
+  document: TextDocument,
-): { type: number; modifiers: number } | undefined {
+  scopeTracker: ScopeTracker
+): TokenInfo {
+  const nodeTypeId = node.type.id
+  const parentTypeId = node.parent?.type.id
+
+  // Special case for now, eventually keywords will go away
+  if (node.type.name === 'keyword') {
+    return {
+      type: TOKEN_TYPES.indexOf(SemanticTokenTypes.keyword),
+      modifiers: 0,
+    }
+  }
+
   switch (nodeTypeId) {
     case Terms.Identifier:
       // Check parent to determine context
@@ -92,6 +112,24 @@ function getTokenType(
           modifiers: 0,
         }
       }
+
+      // Special case: Identifier in PositionalArg - check scope
+      if (parentTypeId === Terms.PositionalArg) {
+        const identifierText = document.getText({
+          start: document.positionAt(node.from),
+          end: document.positionAt(node.to),
+        })
+
+        // If not in scope, treat as string (like a Word)
+        if (!scopeTracker.isInScope(identifierText, node)) {
+          return {
+            type: TOKEN_TYPES.indexOf(SemanticTokenTypes.string),
+            modifiers: 0,
+          }
+        }
+        // If in scope, fall through to treat as variable
+      }
+
       // Otherwise it's a regular variable
       return {
         type: TOKEN_TYPES.indexOf(SemanticTokenTypes.variable),
@@ -149,7 +187,6 @@ function getTokenType(
         modifiers: 0,
       }
 
-    case Terms.keyword:
     case Terms.Do:
     case Terms.colon:
       return {