PiCloud/dashboard/src/lib/rhai/symbols.ts

// Symbol table built from the parsed AST.
//
// One walk produces everything the editor features need:
//   * declarations (let, const, fn, params, for-loop binders, catch binder)
//   * usages (every Ident reference, resolved by walking the scope chain)
//   * object-literal field maps (so `obj.` can suggest known keys)
//
// Resolution rules (matching Rhai):
//   * `fn` declarations live in the script-root scope regardless of where
//     they appear textually. They form a flat namespace; nested functions
//     are not allowed in standard Rhai.
//   * A function body is a fresh scope that does NOT inherit the enclosing
//     locals — Rhai's `fn` is a pure function, not a closure. It can
//     still see top-level `fn`s (call them) but not top-level `let`s.
//   * Blocks (if/while/loop/for/try) nest within their containing scope.
//   * `let`/`const` are visible only after their declaration site within
//     their scope.
//
// Known limit: object-literal field tracking is best-effort. We only
// record fields at the literal-initialization site — `let o = #{ a: 1 };`.
// Reassignments and member writes (`o.b = 2;`) don't update the field set.

import type {
	BlockExpr,
	Comment,
	Expr,
	ForStmt,
	FnDecl,
	IfExpr,
	ObjectMapExpr,
	ParseResult,
	Range,
	Stmt,
	SwitchExpr,
	TryStmt
} from './ast';

export type DeclKind = 'let' | 'const' | 'fn' | 'param' | 'for' | 'catch';

export interface Decl {
	kind: DeclKind;
	name: string;
	nameRange: Range;
	// For `fn`: rendered signature like `process(order, user)`.
	signature?: string;
	// For `let`/`const` initialized to an object-map literal — the field
	// names at the literal site. Empty otherwise.
	objectFields?: string[];
	// Lexical visibility: the offset at which references can resolve to
	// this declaration. For `let`/`const` it's just past the declaration;
	// for `fn`, parameters, `for`, and `catch` binders it's the start of
	// the scope they belong to.
	visibleFrom: number;
	scope: Scope;
}

export interface Usage {
	name: string;
	range: Range;
	scope: Scope;
	resolved: Decl | null;
}

export interface Scope {
	id: number;
	kind: 'root' | 'fn' | 'block';
	// A scope's range covers the source span where its locals are
	// reachable. For the root scope this is the whole document.
	range: Range;
	parent: Scope | null;
	children: Scope[];
	decls: Decl[];
}

export interface SymbolTable {
	root: Scope;
	allDecls: Decl[];
	usages: Usage[];
	// Public API used by the editor features.
	declAt(pos: number): Decl | null;
	declOfUsageAt(pos: number): Decl | null;
	usagesOf(decl: Decl): Range[];
	objectFieldsOf(name: string, atPos: number): string[];
	scopeCompletions(atPos: number): Decl[];
}

export function buildSymbolTable(result: ParseResult): SymbolTable {
	const builder = new Builder(result);
	builder.walkProgram();
	builder.resolveUsages();
	return builder.finish();
}

// Cosmetic helper: format a function's signature for the completion
// `detail` field. Kept here so the symbol table is the single source of
// truth for "how a `fn` shows up in the UI".
export function renderFnSignature(fn: FnDecl): string {
	return `${fn.name}(${fn.params.map((p) => p.name).join(', ')})`;
}

class Builder {
	allDecls: Decl[] = [];
	usages: Usage[] = [];
	root: Scope;
	private currentScope: Scope;
	private nextScopeId = 0;

	constructor(private result: ParseResult) {
		const span = { start: 0, end: result.source.length };
		this.root = this.makeScope('root', span, null);
		this.currentScope = this.root;
	}

	private makeScope(kind: Scope['kind'], range: Range, parent: Scope | null): Scope {
		const s: Scope = { id: this.nextScopeId++, kind, range, parent, children: [], decls: [] };
		if (parent) parent.children.push(s);
		return s;
	}

	private declare(d: Decl): void {
		this.currentScope.decls.push(d);
		this.allDecls.push(d);
	}

	// ----------------------------------------------------------------- walk

	walkProgram(): void {
		// First pass: hoist `fn` decls into the root scope so calls anywhere
		// in the file can resolve to them regardless of source order.
		for (const stmt of this.result.program.stmts) {
			if (stmt.kind === 'FnDecl') {
				this.declare({
					kind: 'fn',
					name: stmt.name,
					nameRange: stmt.nameRange,
					signature: renderFnSignature(stmt),
					visibleFrom: 0,
					scope: this.root
				});
			}
		}
		// Second pass: walk statements normally. Skip re-declaring the
		// already-hoisted fn names; just descend into their bodies.
		for (const stmt of this.result.program.stmts) this.walkStmt(stmt);
	}

	private walkStmt(stmt: Stmt): void {
		switch (stmt.kind) {
			case 'Let':
			case 'Const': {
				if (stmt.init) this.walkExpr(stmt.init);
				const objectFields =
					stmt.init && stmt.init.kind === 'ObjectMap'
						? (stmt.init as ObjectMapExpr).entries.map((e) => e.key)
						: undefined;
				this.declare({
					kind: stmt.kind === 'Let' ? 'let' : 'const',
					name: stmt.name,
					nameRange: stmt.nameRange,
					objectFields,
					visibleFrom: stmt.nameRange.end,
					scope: this.currentScope
				});
				return;
			}
			case 'FnDecl': {
				const prev = this.currentScope;
				const fnScope = this.makeScope('fn', stmt.body, prev);
				this.currentScope = fnScope;
				for (const p of stmt.params) {
					this.declare({
						kind: 'param',
						name: p.name,
						nameRange: { start: p.start, end: p.end },
						visibleFrom: stmt.body.start,
						scope: fnScope
					});
				}
				for (const s of stmt.body.stmts) this.walkStmt(s);
				this.currentScope = prev;
				return;
			}
			case 'ExprStmt':
				this.walkExpr(stmt.expr);
				return;
			case 'Return':
				if (stmt.value) this.walkExpr(stmt.value);
				return;
			case 'While':
				this.walkExpr(stmt.cond);
				this.walkBlock(stmt.body);
				return;
			case 'Loop':
				this.walkBlock(stmt.body);
				return;
			case 'For':
				this.walkFor(stmt);
				return;
			case 'Try':
				this.walkTry(stmt);
				return;
			case 'Break':
			case 'Continue':
				return;
		}
	}

	private walkFor(stmt: ForStmt): void {
		this.walkExpr(stmt.iter);
		const prev = this.currentScope;
		const blockScope = this.makeScope('block', stmt.body, prev);
		this.currentScope = blockScope;
		this.declare({
			kind: 'for',
			name: stmt.varName,
			nameRange: stmt.varRange,
			visibleFrom: stmt.body.start,
			scope: blockScope
		});
		for (const s of stmt.body.stmts) this.walkStmt(s);
		this.currentScope = prev;
	}

	private walkTry(stmt: TryStmt): void {
		this.walkBlock(stmt.body);
		const prev = this.currentScope;
		const handlerScope = this.makeScope('block', stmt.handler, prev);
		this.currentScope = handlerScope;
		if (stmt.catchVar && stmt.catchVarRange) {
			this.declare({
				kind: 'catch',
				name: stmt.catchVar,
				nameRange: stmt.catchVarRange,
				visibleFrom: stmt.handler.start,
				scope: handlerScope
			});
		}
		for (const s of stmt.handler.stmts) this.walkStmt(s);
		this.currentScope = prev;
	}

	private walkBlock(block: BlockExpr): void {
		const prev = this.currentScope;
		const blockScope = this.makeScope('block', block, prev);
		this.currentScope = blockScope;
		for (const s of block.stmts) this.walkStmt(s);
		this.currentScope = prev;
	}

	private walkExpr(expr: Expr): void {
		switch (expr.kind) {
			case 'Ident':
				if (expr.name) {
					this.usages.push({
						name: expr.name,
						range: { start: expr.start, end: expr.end },
						scope: this.currentScope,
						resolved: null
					});
				}
				return;
			case 'Number':
			case 'String':
			case 'Bool':
			case 'Null':
				return;
			case 'Call':
				this.walkExpr(expr.callee);
				for (const a of expr.args) this.walkExpr(a);
				return;
			case 'Member':
				this.walkExpr(expr.object);
				// We don't record the property as a usage — it's resolved
				// against the object's shape, not the lexical scope.
				return;
			case 'Index':
				this.walkExpr(expr.object);
				this.walkExpr(expr.index);
				return;
			case 'Unary':
				this.walkExpr(expr.operand);
				return;
			case 'Binary':
				this.walkExpr(expr.left);
				this.walkExpr(expr.right);
				return;
			case 'Assign':
				this.walkExpr(expr.target);
				this.walkExpr(expr.value);
				return;
			case 'Paren':
				this.walkExpr(expr.expr);
				return;
			case 'ObjectMap':
				for (const e of expr.entries) this.walkExpr(e.value);
				return;
			case 'Array':
				for (const e of expr.elements) this.walkExpr(e);
				return;
			case 'FnExpr': {
				const prev = this.currentScope;
				const fnScope = this.makeScope('fn', expr.body, prev);
				this.currentScope = fnScope;
				for (const p of expr.params) {
					this.declare({
						kind: 'param',
						name: p.name,
						nameRange: { start: p.start, end: p.end },
						visibleFrom: expr.body.start,
						scope: fnScope
					});
				}
				for (const s of expr.body.stmts) this.walkStmt(s);
				this.currentScope = prev;
				return;
			}
			case 'IfExpr':
				this.walkIf(expr);
				return;
			case 'SwitchExpr':
				this.walkSwitch(expr);
				return;
			case 'BlockExpr':
				this.walkBlock(expr);
				return;
		}
	}

	private walkIf(expr: IfExpr): void {
		this.walkExpr(expr.cond);
		this.walkBlock(expr.then);
		if (expr.else_) {
			if (expr.else_.kind === 'IfExpr') this.walkIf(expr.else_);
			else this.walkBlock(expr.else_);
		}
	}

	private walkSwitch(expr: SwitchExpr): void {
		this.walkExpr(expr.subject);
		for (const arm of expr.arms) {
			if (arm.pattern) this.walkExpr(arm.pattern);
			if (arm.guard) this.walkExpr(arm.guard);
			this.walkExpr(arm.value);
		}
	}

	// ----------------------------------------------------------- resolution

	resolveUsages(): void {
		for (const u of this.usages) {
			u.resolved = this.resolveName(u.name, u.range.start, u.scope);
		}
	}

	private resolveName(name: string, atPos: number, fromScope: Scope): Decl | null {
		let scope: Scope | null = fromScope;
		let crossedFn = false;
		while (scope) {
			for (const d of scope.decls) {
				if (d.name !== name) continue;
				// Function scopes don't see outer locals — only the root's
				// `fn` decls. So once we've crossed a fn boundary, accept
				// only `fn` declarations.
				if (crossedFn && d.kind !== 'fn') continue;
				if (d.visibleFrom <= atPos) return d;
			}
			if (scope.kind === 'fn') crossedFn = true;
			scope = scope.parent;
		}
		return null;
	}

	// --------------------------------------------------------------- finish

	finish(): SymbolTable {
		const { allDecls, usages, root } = this;

		const declAt = (pos: number): Decl | null => {
			let best: Decl | null = null;
			for (const d of allDecls) {
				if (pos >= d.nameRange.start && pos <= d.nameRange.end) {
					best = d;
				}
			}
			return best;
		};

		const findScopeAt = (pos: number, scope: Scope = root): Scope => {
			for (const c of scope.children) {
				if (pos >= c.range.start && pos <= c.range.end) {
					return findScopeAt(pos, c);
				}
			}
			return scope;
		};

		const declOfUsageAt = (pos: number): Decl | null => {
			const direct = declAt(pos);
			if (direct) return direct;
			for (const u of usages) {
				if (pos >= u.range.start && pos <= u.range.end) return u.resolved;
			}
			return null;
		};

		const usagesOf = (decl: Decl): Range[] => {
			const out: Range[] = [{ start: decl.nameRange.start, end: decl.nameRange.end }];
			for (const u of usages) {
				if (u.resolved === decl) out.push({ start: u.range.start, end: u.range.end });
			}
			out.sort((a, b) => a.start - b.start);
			return out;
		};

		const objectFieldsOf = (name: string, atPos: number): string[] => {
			const fromScope = findScopeAt(atPos);
			const d = (this as Builder).resolveName(name, atPos, fromScope);
			return d?.objectFields ?? [];
		};

		const scopeCompletions = (atPos: number): Decl[] => {
			const seen = new Set<string>();
			const out: Decl[] = [];
			let scope: Scope | null = findScopeAt(atPos);
			let crossedFn = false;
			while (scope) {
				for (const d of scope.decls) {
					if (seen.has(d.name)) continue;
					if (crossedFn && d.kind !== 'fn') continue;
					if (d.visibleFrom > atPos) continue;
					seen.add(d.name);
					out.push(d);
				}
				if (scope.kind === 'fn') crossedFn = true;
				scope = scope.parent;
			}
			return out;
		};

		return { root, allDecls, usages, declAt, declOfUsageAt, usagesOf, objectFieldsOf, scopeCompletions };
	}
}

// Used by symbols.test.ts; harmless to export.
export function commentsAt(comments: Comment[], pos: number): Comment | undefined {
	return comments.find((c) => pos >= c.start && pos < c.end);
}