fix(deducer): avoid retrieving declarations for non-infrastructure call nodes (#329)

jianzs · web-flow · commit 5e7702efc970 · 2024-08-09T11:43:05.000+08:00
Previously, the deducer attempted to retrieve all call node declarations and match them against custom infrastructure functions. This method was flawed as functions with multiple declarations caused the deducer to fail. To address this, we now first verify that a call node pertains to a custom infrastructure function by comparing function names. Only then do we fetch the call node's declaration, effectively bypassing the collection of extraneous call node declarations.
diff --git a/.changeset/wild-planets-report.md b/.changeset/wild-planets-report.md
@@ -0,0 +1,7 @@
+---
+"@plutolang/pyright-deducer": patch
+---
+
+fix(deducer): avoid retrieving declarations for non-infrastructure call nodes
+
+Previously, the deducer attempted to retrieve all call node declarations and match them against custom infrastructure functions. This method was flawed as functions with multiple declarations caused the deducer to fail. To address this, we now first verify that a call node pertains to a custom infrastructure function by comparing function names. Only then do we fetch the call node's declaration, effectively bypassing the collection of extraneous call node declarations.
diff --git a/components/deducers/python-pyright/src/custom-infra-fn/custom-infra-fn-types.ts b/components/deducers/python-pyright/src/custom-infra-fn/custom-infra-fn-types.ts
@@ -1,9 +1,9 @@
 import { Scope } from "pyright-internal/dist/analyzer/scope";
-import { ParseNode } from "pyright-internal/dist/parser/parseNodes";
+import { FunctionNode } from "pyright-internal/dist/parser/parseNodes";
 import { SourceFile } from "pyright-internal/dist/analyzer/sourceFile";
 
 export interface CustomInfraFn {
-  readonly topNode: ParseNode;
+  readonly topNode: FunctionNode;
   readonly hierarchy: Scope[];
   readonly sourceFile: SourceFile;
 }
diff --git a/components/deducers/python-pyright/src/index.ts b/components/deducers/python-pyright/src/index.ts
@@ -210,12 +210,22 @@ export default class PyrightDeducer extends core.Deducer {
         }
 
         const fnName = statement.leftExpression;
+
+        // First, check if the function name is same as one custom infrastructure function name.
+        const infraFn = customInfraFns.find((fn) => fn.topNode.name.value === fnName.value);
+        if (!infraFn) {
+          continue;
+        }
+
+        // Find the declaration of the called function.
         const decls = this.typeEvaluator!.getDeclarationsForNameNode(fnName);
-        assert(decls?.length === 1, "The function name must be declared only once.");
+        assert(
+          decls?.length === 1,
+          `The function name \`${fnName.value}\` must be declared only once.`
+        );
 
-        const customInfraFn = customInfraFns.find((fn) => fn.topNode.id === decls[0].node.id);
-        if (customInfraFn) {
-          // This call node is a call to a custom infrastructure function.
+        // Second, check if the call node is calling the found custom infrastructure function.
+        if (infraFn.topNode.id === decls[0].node.id) {
           customInfraFnCallNodes.push({
             callNode: statement,
             functionNode: decls[0].node as FunctionNode,
@@ -365,6 +375,10 @@ function findCandidateInfraFns(
     const topNode = ScopeUtils.findTopNodeInScope(node, scopeHierarchy[0]);
     assert(topNode, `No top node found in scope.`);
     if (!customInfraFns.has(topNode.id)) {
+      assert(
+        topNode.nodeType === ParseNodeType.Function,
+        `The custom infra fn must be a function.`
+      );
       customInfraFns.set(topNode.id, { topNode, sourceFile, hierarchy: scopeHierarchy });
     }
   }
