@ugo-studio/jspp 0.3.0 → 0.3.2

package/dist/{analysis → interpreter/analysis}/typeAnalyzer.js

@@ -41,12 +41,511 @@ function getParameterType(node) {
 export class TypeAnalyzer {
     traverser = new Traverser();
     scopeManager = new ScopeManager();
+    callSites = new Map();
     functionTypeInfo = new Map();
     functionStack = [];
     nodeToScope = new Map();
     labelStack = [];
     loopDepth = 0;
     switchDepth = 0;
+    /**
+     * Resolve the scope for a node by walking up the parent chain to find
+     * the nearest ancestor that has a mapped scope entry.
+     */
+    resolveScope(node) {
+        let current = node;
+        while (current) {
+            const mapped = this.nodeToScope.get(current);
+            if (mapped)
+                return mapped;
+            current = current.parent;
+        }
+        return this.scopeManager.currentScope;
+    }
+    /**
+     * Walk up the AST to find the containing function-like declaration.
+     */
+    findContainingFunction(node) {
+        let current = node.parent;
+        while (current) {
+            if (ts.isFunctionDeclaration(current) ||
+                ts.isFunctionExpression(current) ||
+                ts.isArrowFunction(current) ||
+                ts.isMethodDeclaration(current) ||
+                ts.isConstructorDeclaration(current)) {
+                return current;
+            }
+            current = current.parent;
+        }
+        return null;
+    }
+    /**
+     * Check if a call expression is inside a given function (i.e. it's a recursive/internal call).
+     */
+    isCallInsideFunction(call, func) {
+        let current = call.parent;
+        while (current) {
+            if (current === func)
+                return true;
+            current = current.parent;
+        }
+        return false;
+    }
+    /**
+     * Infer the type of a parameter by looking at external (non-recursive) call sites.
+     * Returns null if no external call sites found or types are indeterminate.
+     */
+    inferParameterTypeFromCallSites(paramDecl, scope, visited) {
+        const containingFunc = this.findContainingFunction(paramDecl);
+        if (!containingFunc)
+            return null;
+        // Find the parameter index
+        const paramIndex = containingFunc
+            .parameters.indexOf(paramDecl);
+        if (paramIndex < 0)
+            return null;
+        // Get call sites for this function
+        const sites = this.callSites.get(containingFunc);
+        if (!sites || sites.length === 0)
+            return null;
+        // Filter to external call sites only (not recursive calls from within the function)
+        const externalSites = sites.filter((site) => !this.isCallInsideFunction(site, containingFunc));
+        if (externalSites.length === 0)
+            return null;
+        const types = [];
+        for (const site of externalSites) {
+            const arg = site.arguments[paramIndex];
+            if (arg) {
+                types.push(this.inferNodeReturnType(arg, scope, new Set(visited)));
+            }
+        }
+        const uniqueTypes = new Set(types.filter((t) => t !== "any"));
+        if (uniqueTypes.size === 1) {
+            return uniqueTypes.values().next().value;
+        }
+        if (uniqueTypes.size > 1)
+            return "any";
+        return null; // all were "any" or no arguments
+    }
+    inferFunctionReturnType(node, scope, visited = new Set(), typeEnv) {
+        if (!ts.isFunctionDeclaration(node) && !ts.isFunctionExpression(node) &&
+            !ts.isArrowFunction(node) && !ts.isMethodDeclaration(node) &&
+            !ts.isConstructorDeclaration(node)) {
+            return "any";
+        }
+        // Create a unique key for the visit based on node and environment to allow context-sensitive recursion
+        const envKey = typeEnv
+            ? Array.from(typeEnv.entries()).map(([p, e]) => `${p.pos}:${e.pos}`)
+                .join(",")
+            : "no-env";
+        const visitKey = `func:${node.pos}:${envKey}`;
+        if (visited.has(visitKey))
+            return "any";
+        const newVisited = new Set(visited);
+        newVisited.add(visitKey);
+        const info = this.functionTypeInfo.get(node);
+        const internalScope = this.nodeToScope.get(node) || scope ||
+            this.scopeManager.currentScope;
+        // When called with a specific type environment (from a call site), use assignments analysis
+        // to determine the return type for this particular instantiation.
+        if (typeEnv) {
+            if (info && info.assignments && info.assignments.length > 0) {
+                const types = info.assignments.map((expr) => this.inferNodeReturnType(expr, internalScope, newVisited, typeEnv));
+                const uniqueTypes = new Set(types.filter((t) => t !== "any"));
+                if (uniqueTypes.size === 1) {
+                    return uniqueTypes.values().next().value;
+                }
+                if (uniqueTypes.size > 1)
+                    return "any";
+            }
+        }
+        // When called without typeEnv (top-level inference), prefer call-site analysis over
+        // raw assignments. Call-site analysis instantiates the function with actual argument
+        // types and is more precise — raw assignment analysis can't resolve parameter types
+        // and produces misleading "any" from recursion guards.
+        if (!typeEnv &&
+            (ts.isFunctionDeclaration(node) || ts.isFunctionExpression(node) ||
+                ts.isArrowFunction(node))) {
+            const sites = this.callSites.get(node);
+            if (sites && sites.length > 0) {
+                // Only use external call sites — recursive calls create circular reasoning
+                const externalSites = sites.filter((site) => !this.isCallInsideFunction(site, node));
+                if (externalSites.length > 0) {
+                    const callTypes = externalSites.map((site) => {
+                        const env = new Map();
+                        const params = node.parameters;
+                        params.forEach((p, i) => {
+                            if (site.arguments[i]) {
+                                env.set(p, site.arguments[i]);
+                            }
+                        });
+                        return this.inferFunctionReturnType(node, internalScope, new Set(newVisited), env);
+                    });
+                    // Don't filter "any" here — each external call site was instantiated with
+                    // real arguments, so "any" means that instantiation genuinely returns mixed types.
+                    const uniqueCallTypes = new Set(callTypes);
+                    if (uniqueCallTypes.size === 1) {
+                        return uniqueCallTypes.values().next().value;
+                    }
+                    if (uniqueCallTypes.size > 1)
+                        return "any";
+                }
+            }
+        }
+        // Fallback: use assignments analysis without typeEnv (e.g. no call sites exist)
+        if (!typeEnv && info && info.assignments && info.assignments.length > 0) {
+            const types = info.assignments.map((expr) => this.inferNodeReturnType(expr, internalScope, newVisited, typeEnv));
+            const uniqueTypes = new Set(types.filter((t) => t !== "any"));
+            if (uniqueTypes.size === 1) {
+                return uniqueTypes.values().next().value;
+            }
+            if (uniqueTypes.size > 1)
+                return "any";
+        }
+        // Type checking TS signature
+        const signature = node;
+        if (signature.type) {
+            switch (signature.type.kind) {
+                case ts.SyntaxKind.StringKeyword:
+                    return "string";
+                case ts.SyntaxKind.NumberKeyword:
+                    return "number";
+                case ts.SyntaxKind.BooleanKeyword:
+                    return "boolean";
+                case ts.SyntaxKind.ObjectKeyword:
+                    return "object";
+                case ts.SyntaxKind.VoidKeyword:
+                    return "any";
+            }
+        }
+        return "any";
+    }
+    inferNodeReturnType(node, scope, visited = new Set(), typeEnv) {
+        const envKey = typeEnv
+            ? Array.from(typeEnv.entries()).map(([p, e]) => `${p.pos}:${e.pos}`)
+                .join(",")
+            : "";
+        const visitKey = `node:${node.pos}:${node.kind}:${envKey}`;
+        if (visited.has(visitKey))
+            return "any";
+        const newVisited = new Set(visited);
+        newVisited.add(visitKey);
+        if (!scope) {
+            scope = this.resolveScope(node);
+        }
+        // 0a. Handle ParameterDeclaration nodes directly
+        if (ts.isParameter(node)) {
+            if (typeEnv) {
+                const arg = typeEnv.get(node);
+                if (arg) {
+                    return this.inferNodeReturnType(arg, scope, newVisited, typeEnv);
+                }
+            }
+            // Infer from call sites
+            const callSiteType = this.inferParameterTypeFromCallSites(node, scope, newVisited);
+            if (callSiteType)
+                return callSiteType;
+            // Fall back to TS type annotation
+            if (node.type) {
+                switch (node.type.kind) {
+                    case ts.SyntaxKind.StringKeyword:
+                        return "string";
+                    case ts.SyntaxKind.NumberKeyword:
+                        return "number";
+                    case ts.SyntaxKind.BooleanKeyword:
+                        return "boolean";
+                    case ts.SyntaxKind.ObjectKeyword:
+                        return "object";
+                    case ts.SyntaxKind.FunctionType:
+                        return "function";
+                }
+            }
+            return "any";
+        }
+        // 0b. Check parameter environment for identifiers
+        if (ts.isIdentifier(node)) {
+            const definingScope = scope.findScopeFor(node.text);
+            const typeInfo = definingScope
+                ? definingScope.symbols.get(node.text)
+                : null;
+            if (typeInfo && typeInfo.isParameter && typeInfo.declaration &&
+                ts.isParameter(typeInfo.declaration)) {
+                const arg = typeEnv?.get(typeInfo.declaration);
+                if (arg) {
+                    return this.inferNodeReturnType(arg, scope, newVisited, typeEnv);
+                }
+                // No typeEnv for this parameter — try call-site inference
+                const callSiteType = this.inferParameterTypeFromCallSites(typeInfo.declaration, scope, newVisited);
+                if (callSiteType)
+                    return callSiteType;
+            }
+        }
+        if (ts.isVariableDeclarationList(node)) {
+            if (node.declarations.length > 0) {
+                return this.inferNodeReturnType(node.declarations[0], scope, newVisited, typeEnv);
+            }
+            return "any";
+        }
+        if (ts.isVariableDeclaration(node)) {
+            if (node.initializer) {
+                return this.inferNodeReturnType(node.initializer, scope, newVisited, typeEnv);
+            }
+            return "any";
+        }
+        // 1. Literal types
+        if (ts.isNumericLiteral(node))
+            return "number";
+        if (ts.isStringLiteral(node) || ts.isNoSubstitutionTemplateLiteral(node))
+            return "string";
+        if (node.kind === ts.SyntaxKind.TrueKeyword ||
+            node.kind === ts.SyntaxKind.FalseKeyword)
+            return "boolean";
+        if (node.kind === ts.SyntaxKind.NullKeyword ||
+            node.kind === ts.SyntaxKind.UndefinedKeyword)
+            return "any";
+        // 2. Complex literals / expressions
+        if (ts.isObjectLiteralExpression(node) ||
+            ts.isArrayLiteralExpression(node))
+            return "object";
+        if (ts.isFunctionExpression(node) || ts.isArrowFunction(node) ||
+            ts.isClassExpression(node))
+            return "function";
+        if (ts.isClassDeclaration(node) || ts.isFunctionDeclaration(node)) {
+            return "function";
+        }
+        if (ts.isEnumDeclaration(node))
+            return "object";
+        // 3. Parenthesized / Await / Yield
+        if (ts.isParenthesizedExpression(node)) {
+            return this.inferNodeReturnType(node.expression, scope, newVisited, typeEnv);
+        }
+        if (ts.isAwaitExpression(node)) {
+            return this.inferNodeReturnType(node.expression, scope, newVisited, typeEnv);
+        }
+        if (ts.isYieldExpression(node)) {
+            return node.expression
+                ? this.inferNodeReturnType(node.expression, scope, newVisited, typeEnv)
+                : "any";
+        }
+        // 4. Unary
+        if (ts.isTypeOfExpression(node))
+            return "string";
+        if (ts.isVoidExpression(node))
+            return "any";
+        if (ts.isDeleteExpression(node))
+            return "boolean";
+        if (ts.isPrefixUnaryExpression(node)) {
+            switch (node.operator) {
+                case ts.SyntaxKind.ExclamationToken:
+                    return "boolean";
+                case ts.SyntaxKind.PlusToken:
+                case ts.SyntaxKind.MinusToken:
+                case ts.SyntaxKind.TildeToken:
+                case ts.SyntaxKind.PlusPlusToken:
+                case ts.SyntaxKind.MinusMinusToken:
+                    return "number";
+            }
+        }
+        if (ts.isPostfixUnaryExpression(node))
+            return "number";
+        // 5. Binary
+        if (ts.isBinaryExpression(node)) {
+            const op = node.operatorToken.kind;
+            // Assignment
+            if (op >= ts.SyntaxKind.FirstAssignment &&
+                op <= ts.SyntaxKind.LastAssignment) {
+                return this.inferNodeReturnType(node.right, scope, newVisited, typeEnv);
+            }
+            // Comparison
+            if ([
+                ts.SyntaxKind.EqualsEqualsToken,
+                ts.SyntaxKind.ExclamationEqualsToken,
+                ts.SyntaxKind.EqualsEqualsEqualsToken,
+                ts.SyntaxKind.ExclamationEqualsEqualsToken,
+                ts.SyntaxKind.LessThanToken,
+                ts.SyntaxKind.LessThanEqualsToken,
+                ts.SyntaxKind.GreaterThanToken,
+                ts.SyntaxKind.GreaterThanEqualsToken,
+                ts.SyntaxKind.InKeyword,
+                ts.SyntaxKind.InstanceOfKeyword,
+            ].includes(op)) {
+                return "boolean";
+            }
+            // Arithmetic
+            if ([
+                ts.SyntaxKind.MinusToken,
+                ts.SyntaxKind.AsteriskToken,
+                ts.SyntaxKind.SlashToken,
+                ts.SyntaxKind.PercentToken,
+                ts.SyntaxKind.AsteriskAsteriskToken,
+                ts.SyntaxKind.LessThanLessThanToken,
+                ts.SyntaxKind.GreaterThanGreaterThanToken,
+                ts.SyntaxKind.GreaterThanGreaterThanGreaterThanToken,
+                ts.SyntaxKind.AmpersandToken,
+                ts.SyntaxKind.BarToken,
+                ts.SyntaxKind.CaretToken,
+            ].includes(op)) {
+                return "number";
+            }
+            // Plus
+            if (op === ts.SyntaxKind.PlusToken) {
+                const left = this.inferNodeReturnType(node.left, scope, newVisited, typeEnv);
+                const right = this.inferNodeReturnType(node.right, scope, newVisited, typeEnv);
+                if (left === "string" || right === "string")
+                    return "string";
+                if (left === "number" && right === "number")
+                    return "number";
+                return "any";
+            }
+            // Logical
+            if ([
+                ts.SyntaxKind.AmpersandAmpersandToken,
+                ts.SyntaxKind.BarBarToken,
+                ts.SyntaxKind.QuestionQuestionToken,
+            ].includes(op)) {
+                const left = this.inferNodeReturnType(node.left, scope, newVisited, typeEnv);
+                const right = this.inferNodeReturnType(node.right, scope, newVisited, typeEnv);
+                return left === right ? left : "any";
+            }
+        }
+        // 6. Conditional
+        if (ts.isConditionalExpression(node)) {
+            const trueType = this.inferNodeReturnType(node.whenTrue, scope, newVisited, typeEnv);
+            const falseType = this.inferNodeReturnType(node.whenFalse, scope, newVisited, typeEnv);
+            return trueType === falseType ? trueType : "any";
+        }
+        // 7. Template
+        if (ts.isTemplateExpression(node))
+            return "string";
+        // 8. New
+        if (ts.isNewExpression(node))
+            return "object";
+        // 9. Call
+        if (ts.isCallExpression(node)) {
+            const callee = node.expression;
+            if (ts.isIdentifier(callee)) {
+                const name = callee.text;
+                if (name === "Number" || name === "parseInt" ||
+                    name === "parseFloat")
+                    return "number";
+                if (name === "String")
+                    return "string";
+                if (name === "Boolean" || name === "isNaN" ||
+                    name === "isFinite")
+                    return "boolean";
+                const typeInfo = this.scopeManager.lookupFromScope(name, scope);
+                if (typeInfo && typeInfo.declaration) {
+                    const decl = typeInfo.declaration;
+                    const env = new Map();
+                    if (ts.isFunctionLike(decl)) {
+                        decl.parameters.forEach((p, i) => {
+                            if (node.arguments[i]) {
+                                env.set(p, node.arguments[i]);
+                            }
+                        });
+                    }
+                    return this.inferFunctionReturnType(decl, scope, newVisited, env);
+                }
+            }
+            else if (ts.isPropertyAccessExpression(callee)) {
+                const obj = callee.expression.getText();
+                const prop = callee.name.text;
+                const full = `${obj}.${prop}`;
+                if (full.startsWith("Math.") && prop !== "PI" && prop !== "E") {
+                    return "number";
+                }
+                if (full === "Array.isArray")
+                    return "boolean";
+                if (full === "Object.keys" || full === "Object.values" ||
+                    full === "Object.entries")
+                    return "object";
+            }
+            return "any";
+        }
+        // 10. Identifier
+        if (ts.isIdentifier(node)) {
+            if (node.text === "NaN" || node.text === "Infinity") {
+                return "number";
+            }
+            if (node.text === "undefined")
+                return "any";
+            const typeInfo = this.scopeManager.lookupFromScope(node.text, scope);
+            if (typeInfo) {
+                // Check assignments
+                if (typeInfo.assignments && typeInfo.assignments.length > 0) {
+                    const types = typeInfo.assignments.map((expr) => this.inferNodeReturnType(expr, scope, newVisited, typeEnv));
+                    const uniqueTypes = new Set(types.filter((t) => t !== "any"));
+                    if (uniqueTypes.size === 1) {
+                        return uniqueTypes.values().next().value;
+                    }
+                    if (uniqueTypes.size > 1)
+                        return "any";
+                }
+                // Check declaration for TS type
+                if (typeInfo.declaration &&
+                    ts.isVariableDeclaration(typeInfo.declaration) &&
+                    typeInfo.declaration.type) {
+                    const kind = typeInfo.declaration.type.kind;
+                    if (kind === ts.SyntaxKind.StringKeyword)
+                        return "string";
+                    if (kind === ts.SyntaxKind.NumberKeyword)
+                        return "number";
+                    if (kind === ts.SyntaxKind.BooleanKeyword)
+                        return "boolean";
+                    if (kind === ts.SyntaxKind.ObjectKeyword)
+                        return "object";
+                }
+                // Builtin check
+                if (typeInfo.isBuiltin) {
+                    if (["console", "Math", "process", "global", "globalThis"]
+                        .includes(node.text))
+                        return "object";
+                }
+                // Fallback to what we know
+                if (typeInfo.type === "string" || typeInfo.type === "number" ||
+                    typeInfo.type === "boolean" || typeInfo.type === "object" ||
+                    typeInfo.type === "function") {
+                    return typeInfo.type;
+                }
+                if (typeInfo.type === "array")
+                    return "object";
+            }
+        }
+        // 11. Property Access
+        if (ts.isPropertyAccessExpression(node)) {
+            if (node.name.text === "length") {
+                const objType = this.inferNodeReturnType(node.expression, scope, newVisited, typeEnv);
+                if (objType === "string" || objType === "object") {
+                    return "number";
+                }
+            }
+            // Hard to know other properties without full type system
+            return "any";
+        }
+        // 12. Function / Method Return type inference
+        if (ts.isFunctionDeclaration(node) || ts.isFunctionExpression(node) ||
+            ts.isArrowFunction(node) || ts.isMethodDeclaration(node) ||
+            ts.isConstructorDeclaration(node)) {
+            return this.inferFunctionReturnType(node, scope, newVisited, typeEnv);
+        }
+        return "any";
+    }
+    defineParameter(nameNode, p, type) {
+        if (ts.isIdentifier(nameNode)) {
+            this.scopeManager.define(nameNode.text, {
+                type: type || getParameterType(p),
+                isParameter: true,
+                declaration: p,
+            });
+        }
+        else {
+            nameNode.elements.forEach((element) => {
+                if (ts.isBindingElement(element)) {
+                    this.defineParameter(element.name, p, type);
+                }
+            });
+        }
+    }
     analyze(ast) {
         this.nodeToScope.set(ast, this.scopeManager.currentScope);
         const crossScopeModificationVisitor = (node) => {
@@ -237,6 +736,9 @@ export class TypeAnalyzer {
                             type: "function",
                             isClosure: false,
                             captures: new Map(),
+                            assignments: !ts.isBlock(node.body)
+                                ? [node.body]
+                                : [],
                         };
                         this.functionTypeInfo.set(node, funcType);
                         this.scopeManager.enterScope(node);
@@ -246,11 +748,7 @@ export class TypeAnalyzer {
                             if (p.getText() == "this") { // Catch invalid parameters
                                 throw new CompilerError("Cannot use 'this' as a parameter name.", p, "SyntaxError");
                             }
-                            this.scopeManager.define(p.name.getText(), {
-                                type: getParameterType(p),
-                                isParameter: true,
-                                declaration: p,
-                            });
+                            this.defineParameter(p.name, p);
                         });
                         this.functionStack.push(node);
                     }
@@ -270,7 +768,8 @@ export class TypeAnalyzer {
                             isClosure: false,
                             captures: new Map(),
                             declaration: node,
-                            needsHeapAllocation: true, // Added: Functions are always heap-allocated
+                            needsHeapAllocation: true,
+                            assignments: [],
                         };
                         this.functionTypeInfo.set(node, funcType);
                         this.scopeManager.enterScope(node);
@@ -284,11 +783,7 @@ export class TypeAnalyzer {
                             if (p.getText() == "this") { // Catch invalid parameters
                                 throw new CompilerError("Cannot use 'this' as a parameter name.", p, "SyntaxError");
                             }
-                            this.scopeManager.define(p.name.getText(), {
-                                type: getParameterType(p),
-                                isParameter: true,
-                                declaration: p,
-                            });
+                            this.defineParameter(p.name, p);
                         });
                         this.functionStack.push(node);
                     }
@@ -312,6 +807,7 @@ export class TypeAnalyzer {
                                 captures: new Map(),
                                 declaration: node,
                                 needsHeapAllocation: true,
+                                assignments: [],
                             };
                             this.scopeManager.define(funcName, funcType);
                             this.functionTypeInfo.set(node, funcType);
@@ -323,11 +819,7 @@ export class TypeAnalyzer {
                             if (p.getText() == "this") { // Catch invalid parameters
                                 throw new CompilerError("Cannot use 'this' as a parameter name.", p, "SyntaxError");
                             }
-                            this.scopeManager.define(p.name.getText(), {
-                                type: getParameterType(p),
-                                isParameter: true,
-                                declaration: p,
-                            });
+                            this.defineParameter(p.name, p);
                         });
                         this.functionStack.push(node);
                     }
@@ -348,6 +840,7 @@ export class TypeAnalyzer {
                             type: "function", // Classes are functions
                             declaration: classNode,
                             needsHeapAllocation: true,
+                            assignments: [],
                         };
                         this.scopeManager.define(name, typeInfo);
                     }
@@ -388,17 +881,13 @@ export class TypeAnalyzer {
                             captures: new Map(),
                             declaration: node,
                             needsHeapAllocation: true,
-                        };
-                        // Methods don't need to be defined in scope by name generally,
+                            assignments: [],
+                        }; // Methods don't need to be defined in scope by name generally,
                         // but we need to track them for captures.
                         this.functionTypeInfo.set(node, funcType);
                         this.scopeManager.enterScope(node);
                         this.nodeToScope.set(node, this.scopeManager.currentScope);
-                        node.parameters.forEach((p) => this.scopeManager.define(p.name.getText(), {
-                            type: "auto",
-                            isParameter: true,
-                            declaration: p,
-                        }));
+                        node.parameters.forEach((p) => this.defineParameter(p.name, p, "auto"));
                         this.functionStack.push(node);
                     }
                 },
@@ -418,15 +907,12 @@ export class TypeAnalyzer {
                             captures: new Map(),
                             declaration: node,
                             needsHeapAllocation: true,
+                            assignments: [],
                         };
                         this.functionTypeInfo.set(node, funcType);
                         this.scopeManager.enterScope(node);
                         this.nodeToScope.set(node, this.scopeManager.currentScope);
-                        node.parameters.forEach((p) => this.scopeManager.define(p.name.getText(), {
-                            type: "auto",
-                            isParameter: true,
-                            declaration: p,
-                        }));
+                        node.parameters.forEach((p) => this.defineParameter(p.name, p, "auto"));
                         this.functionStack.push(node); // Constructor acts like a function
                     }
                 },
@@ -481,6 +967,9 @@ export class TypeAnalyzer {
                                     declaration: node,
                                     isConst,
                                     needsHeapAllocation: needsHeap,
+                                    assignments: node.initializer
+                                        ? [node.initializer]
+                                        : [],
                                 };
                                 if (isBlockScoped) {
                                     this.scopeManager.define(name, typeInfo);
@@ -547,6 +1036,52 @@ export class TypeAnalyzer {
                                 ts.SyntaxKind.LastAssignment;
                         if (isAssignment) {
                             crossScopeModificationVisitor(node.left);
+                            if (ts.isIdentifier(node.left)) {
+                                const name = node.left.text;
+                                const typeInfo = this.scopeManager.lookup(name);
+                                if (typeInfo) {
+                                    if (!typeInfo.assignments) {
+                                        typeInfo.assignments = [];
+                                    }
+                                    typeInfo.assignments.push(node.right);
+                                }
+                            }
+                        }
+                    }
+                },
+            },
+            CallExpression: {
+                enter: (node) => {
+                    if (ts.isCallExpression(node)) {
+                        const callee = node.expression;
+                        if (ts.isIdentifier(callee)) {
+                            const name = callee.text;
+                            const typeInfo = this.scopeManager.lookup(name);
+                            if (typeInfo && typeInfo.declaration &&
+                                ts.isFunctionLike(typeInfo.declaration)) {
+                                const decl = typeInfo
+                                    .declaration;
+                                if (!this.callSites.has(decl)) {
+                                    this.callSites.set(decl, []);
+                                }
+                                this.callSites.get(decl).push(node);
+                            }
+                        }
+                    }
+                },
+            },
+            ReturnStatement: {
+                enter: (node) => {
+                    if (ts.isReturnStatement(node) && node.expression) {
+                        const currentFuncNode = this.functionStack[this.functionStack.length - 1];
+                        if (currentFuncNode) {
+                            const info = this.functionTypeInfo.get(currentFuncNode);
+                            if (info) {
+                                if (!info.assignments) {
+                                    info.assignments = [];
+                                }
+                                info.assignments.push(node.expression);
+                            }
                         }
                     }
                 },