shrimp/vscode-extension/server/src/scopeTracker.ts
2025-11-05 14:48:12 -08:00

136 lines
3.7 KiB
TypeScript

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
}
})
}
}