stone-lang 0.1.0

package/frontend/type-checker/overloadResolution.js

@@ -0,0 +1,605 @@
+/**
+ * Overload Resolution Methods - mixin for TypeChecker
+ *
+ * Methods for static overload resolution of both primitive (from OVERLOADS)
+ * and user-defined function overloads.
+ */
+
+import {
+  FunctionType,
+  ArrayType,
+  TypeVariable,
+  unify,
+  parseTypeString,
+  matchType,
+  substituteType,
+  canPromote,
+  formatType,
+  normalizeType,
+} from "../types/types.js";
+import { OverloadedFunctionType } from "./OverloadedFunctionType.js";
+import { getOverloads } from "../../backends/js-vm/primitives/overloads.js";
+import { CompiledFunction } from "../../backends/js-vm/bytecode.js";
+
+/**
+ * Overload resolution methods to be mixed into TypeChecker.prototype
+ */
+export const overloadResolutionMethods = {
+  /**
+   * Score how well argument types match a function signature
+   * Supports bidirectional inference: TypeVariables are unified with expected types.
+   *
+   * @param {Array<Object>} argTypes - Structured argument types
+   * @param {Array<string>} paramTypeStrings - Parameter type strings from signature
+   * @returns {{score: number, bindings: Object, promotions: boolean[], unifications: Array}|null}
+   */
+  scoreOverload(argTypes, paramTypeStrings) {
+    // Arity must match
+    if (argTypes.length !== paramTypeStrings.length) {
+      return null;
+    }
+
+    let totalScore = 0;
+    let bindings = {};
+    const promotions = new Array(argTypes.length).fill(false);
+    const unifications = []; // Pending TypeVariable unifications: [{typeVarId, targetType}]
+    const deferredConstraints = []; // Deferred rank constraints: [{paramIndex, type, rankVar, requiredRank}]
+
+    const polymorphicPositions = []; // Track which positions have unresolved TypeVariables
+
+    for (let i = 0; i < argTypes.length; i++) {
+      const argType = argTypes[i];
+      const paramPattern = parseTypeString(paramTypeStrings[i]);
+
+      // Bidirectional inference: if argType is an unresolved TypeVariable,
+      // DON'T unify immediately - keep it polymorphic and track the position
+      if (argType.kind === "typevar" && !argType.resolved) {
+        // Track this position for polymorphic result type inference
+        polymorphicPositions.push({ index: i, typeVarId: argType.id, argType });
+        // TypeVariable matches any type with score 0.5 (prefer concrete types)
+        totalScore += 0.5;
+        continue;
+      }
+
+      // Try exact match
+      const matchResult = matchType(paramPattern, argType, { ...bindings });
+      if (matchResult.success) {
+        bindings = matchResult.bindings || bindings;
+        // Collect any TypeVariable unifications from function type matching
+        if (matchResult.typeVarUnifications) {
+          unifications.push(...matchResult.typeVarUnifications);
+        }
+        // Score 0 for exact match, no promotion needed
+        continue;
+      }
+
+      // Handle deferred constraints (variable rank vs concrete rank)
+      if (matchResult.deferred) {
+        bindings = matchResult.bindings || bindings;
+        // Collect deferred constraints for this parameter
+        if (matchResult.constraints) {
+          for (const constraint of matchResult.constraints) {
+            deferredConstraints.push({
+              paramIndex: i,
+              ...constraint
+            });
+          }
+        }
+        // Collect any TypeVariable unifications from function type matching
+        if (matchResult.typeVarUnifications) {
+          unifications.push(...matchResult.typeVarUnifications);
+        }
+        // Small penalty for deferred match (prefer exact matches)
+        totalScore += 0.1;
+        continue;
+      }
+
+      // Try promotion (num -> complex or user-defined)
+      const promoResult = canPromote(argType, paramPattern);
+      if (promoResult.can) {
+        totalScore += 1;
+        promotions[i] = promoResult.converter; // Store converter name instead of boolean
+        continue;
+      }
+
+      // Handle primitive num -> complex promotion (for structured types in primitive signatures)
+      if (
+        argType.kind === "primitive" &&
+        argType.name === "num" &&
+        paramPattern.kind === "primitive" &&
+        paramPattern.name === "complex"
+      ) {
+        totalScore += 1;
+        promotions[i] = "to_complex"; // Store converter name
+        continue;
+      }
+
+      // Handle array<num> -> array<complex> promotion (for structured types in primitive signatures)
+      if (
+        argType.kind === "array" &&
+        argType.elem?.kind === "primitive" &&
+        argType.elem.name === "num" &&
+        paramPattern.kind === "array" &&
+        paramPattern.elem?.kind === "primitive" &&
+        paramPattern.elem.name === "complex"
+      ) {
+        // Match ranks
+        if (typeof paramPattern.rank === "object" && paramPattern.rank.var) {
+          const varName = paramPattern.rank.var;
+          if (varName in bindings) {
+            if (bindings[varName] !== argType.rank) return null;
+          } else {
+            bindings[varName] = argType.rank;
+          }
+        } else if (paramPattern.rank !== argType.rank) {
+          return null;
+        }
+        totalScore += 1;
+        promotions[i] = "to_complex"; // Store converter name
+        continue;
+      }
+
+      // No match possible
+      return null;
+    }
+
+    return { score: totalScore, bindings, promotions, unifications, deferredConstraints, polymorphicPositions };
+  },
+
+  /**
+   * Resolve overload statically using inferred types
+   * Stores result on AST node for executor to use
+   *
+   * Supports bidirectional inference: when argument types are TypeVariables,
+   * they get unified with the expected parameter types from the best matching overload.
+   *
+   * @param {string} name - Function name (e.g., "add", "sin")
+   * @param {Array<Object>} argTypes - Class-based types of arguments
+   * @param {Object} node - AST node to annotate with resolved overload
+   * @param {Object} errorLocation - Optional location to use for error reporting
+   * @returns {Object|null} Output type (structured) or null if no match
+   */
+  resolveOverloadStatic(name, argTypes, node, errorLocation = null) {
+    const candidates = this.overloads[name] || [];
+    if (candidates.length === 0) {
+      return null; // Not a known overloaded function
+    }
+
+    // Convert class-based types to structured types
+    // This also populates this.typeVarMap for any TypeVariables encountered
+    // Then normalize to flatten nested arrays: array<array<num,1>,1> -> array<num,2>
+    const structuredArgTypes = argTypes.map(t => normalizeType(this.toStructuredType(t)));
+
+    let bestMatch = null;
+    let bestScore = Infinity;
+    let bestPromotions = [];
+    let bestBindings = {};
+    let bestUnifications = [];
+    let bestDeferredConstraints = [];
+    let bestPolymorphicPositions = [];
+
+    for (const fn of candidates) {
+      const sig = fn.stone;
+      if (!sig) continue;
+
+      const result = this.scoreOverload(structuredArgTypes, sig.in);
+      if (result && result.score < bestScore) {
+        bestScore = result.score;
+        bestMatch = fn;
+        bestPromotions = result.promotions;
+        bestBindings = result.bindings;
+        bestUnifications = result.unifications || [];
+        bestDeferredConstraints = result.deferredConstraints || [];
+        bestPolymorphicPositions = result.polymorphicPositions || [];
+      }
+    }
+
+    if (!bestMatch) {
+      // Only report error if there were no TypeVariables involved
+      // (TypeVariables might match at runtime with different types)
+      const hasTypeVars = structuredArgTypes.some(t => this.hasUnresolvedTypeVars(t));
+      if (!hasTypeVars) {
+        const argTypeStrs = structuredArgTypes.map(formatType).join(", ");
+        const available = candidates
+          .filter(f => f.stone)
+          .map(f => `  (${f.stone.in.join(", ")}) -> ${f.stone.out}`)
+          .join("\n");
+        this.addError(
+          `No matching overload for ${name}(${argTypeStrs})\nAvailable:\n${available}`,
+          errorLocation || node?.location
+        );
+      }
+      return null;
+    }
+
+    // Apply TypeVariable unifications from bidirectional inference
+    // This constrains TypeVariables to their expected types
+    for (const { typeVarId, targetType } of bestUnifications) {
+      const typeVar = this.typeVarMap.get(typeVarId);
+      if (typeVar && !typeVar.resolved) {
+        // Convert structured type to class-based type and unify
+        const classBasedTarget = this.fromStructuredType(targetType);
+        unify(typeVar, classBasedTarget);
+      }
+    }
+
+    // Record deferred constraints for the current function context
+    // These will be attached to the function type being defined
+    for (const constraint of bestDeferredConstraints) {
+      this.recordDeferredConstraint(constraint);
+    }
+
+    // Unify rank TypeVariables with their constrained values
+    // This ensures subsequent code in the function body sees the constrained rank
+    for (const constraint of bestDeferredConstraints) {
+      if (constraint.type === 'rank' && constraint.paramIndex !== undefined) {
+        const argType = argTypes[constraint.paramIndex];
+        if (argType instanceof ArrayType && argType.rank instanceof TypeVariable) {
+          // Directly bind the TypeVariable to the required rank
+          const rankVar = argType.rank.resolve();
+          if (rankVar instanceof TypeVariable && !rankVar.resolved) {
+            rankVar.resolved = true;
+            rankVar.resolvedTo = constraint.requiredRank;
+          }
+        }
+      }
+    }
+
+    // Check if we should keep the result type polymorphic
+    // This happens when:
+    // 1. There are TypeVariable arguments (polymorphicPositions)
+    // 2. All matching overloads have consistent pattern: input[i] == output
+    let keptPolymorphic = false;
+    if (bestPolymorphicPositions.length > 0) {
+      const polyPattern = this.inferPolymorphicResultPattern(name, candidates, bestPolymorphicPositions, structuredArgTypes);
+      if (polyPattern && polyPattern.preservesInputType) {
+        // Return the TypeVariable from the source position as the result type
+        const sourceArgType = bestPolymorphicPositions.find(
+          p => p.index === polyPattern.sourcePosition
+        );
+        if (sourceArgType) {
+          // Store on AST node for executor (will resolve at runtime)
+          if (node) {
+            node.resolvedOverload = {
+              fn: bestMatch,
+              promotions: bestPromotions,
+              polymorphic: true,
+            };
+          }
+          keptPolymorphic = true;
+          // Return the original TypeVariable - keeps function polymorphic
+          return sourceArgType.argType;
+        }
+      }
+
+      // Pattern check failed - apply unifications for TypeVariables
+      // This is the fallback: unify TypeVars with the best match's parameter types
+      if (!keptPolymorphic) {
+        const sig = bestMatch.stone;
+        for (const { index, typeVarId } of bestPolymorphicPositions) {
+          const targetType = parseTypeString(sig.in[index]);
+          const typeVar = this.typeVarMap.get(typeVarId);
+          if (typeVar && !typeVar.resolved) {
+            const classBasedTarget = this.fromStructuredType(targetType);
+            unify(typeVar, classBasedTarget);
+          }
+        }
+      }
+    }
+
+    // Compute output type with bindings substituted
+    const outputType = substituteType(
+      parseTypeString(bestMatch.stone.out),
+      bestBindings
+    );
+
+    // Store on AST node for executor
+    if (node) {
+      // If bestMatch is a mock (not a real function/CompiledFunction), find the real function from dynamic overloads
+      // This happens for complex operations where the type checker has mock signatures
+      // but the actual functions are registered dynamically when the complex module loads
+      let fnToUse = bestMatch;
+      const isRealFunction = typeof bestMatch === 'function' || bestMatch instanceof CompiledFunction;
+      if (!isRealFunction && bestMatch.stone) {
+        // Look for a real function with matching signature in dynamic overloads
+        const dynamicOverloads = getOverloads(name);
+        const sigKey = JSON.stringify(bestMatch.stone);
+        for (const fn of dynamicOverloads) {
+          const fnIsReal = typeof fn === 'function' || fn instanceof CompiledFunction;
+          if (fnIsReal && fn.stone && JSON.stringify(fn.stone) === sigKey) {
+            fnToUse = fn;
+            break;
+          }
+        }
+      }
+
+      node.resolvedOverload = {
+        fn: fnToUse,
+        promotions: bestPromotions,
+      };
+    }
+
+    return outputType;
+  },
+
+  /**
+   * Infer if all MATCHING overloads have a consistent pattern where input type == output type.
+   * This allows keeping TypeVariables polymorphic instead of unifying them prematurely.
+   *
+   * Only considers overloads that could actually match the argument types (ignoring TypeVar positions).
+   *
+   * @param {string} name - Function name
+   * @param {Array} candidates - Overload candidates
+   * @param {Array} polymorphicPositions - Positions with TypeVariable arguments
+   * @param {Array} structuredArgTypes - The actual argument types (with TypeVars)
+   * @returns {{preservesInputType: boolean, sourcePosition: number}|null}
+   */
+  inferPolymorphicResultPattern(name, candidates, polymorphicPositions, structuredArgTypes) {
+    if (polymorphicPositions.length === 0) return null;
+
+    const polyIndices = new Set(polymorphicPositions.map(p => p.index));
+
+    // Filter candidates to only those that match the non-polymorphic argument types
+    const matchingCandidates = candidates.filter(fn => {
+      const sig = fn.stone;
+      if (!sig || sig.in.length !== structuredArgTypes.length) return false;
+
+      // Check that non-polymorphic positions match
+      for (let i = 0; i < structuredArgTypes.length; i++) {
+        if (polyIndices.has(i)) continue; // Skip TypeVariable positions
+
+        const argType = structuredArgTypes[i];
+        const paramPattern = parseTypeString(sig.in[i]);
+
+        // Simple compatibility check - does the arg type fit the param pattern?
+        const matchResult = matchType(paramPattern, argType, {});
+        if (!matchResult.success && !matchResult.deferred) {
+          // Also check for promotion
+          const promoResult = canPromote(argType, paramPattern);
+          if (!promoResult.can) {
+            return false;
+          }
+        }
+      }
+      return true;
+    });
+
+    if (matchingCandidates.length === 0) return null;
+
+    // Check each polymorphic position to see if it determines the output type
+    for (const { index } of polymorphicPositions) {
+      let allMatch = true;
+
+      for (const fn of matchingCandidates) {
+        const sig = fn.stone;
+        if (!sig || sig.in.length <= index) {
+          allMatch = false;
+          break;
+        }
+
+        // Check if input type at this position matches output type
+        const inputType = sig.in[index];
+        const outputType = sig.out;
+
+        if (inputType !== outputType) {
+          allMatch = false;
+          break;
+        }
+      }
+
+      if (allMatch) {
+        return { preservesInputType: true, sourcePosition: index };
+      }
+    }
+
+    return null;
+  },
+
+  /**
+   * Resolve user-defined function overloads at compile time.
+   * Selects the best matching overload based on argument types and stores
+   * the resolved overloadId on the AST node for the executor.
+   *
+   * @param {string} funcName - Function name
+   * @param {Array} argTypes - Array of argument types (class-based)
+   * @param {OverloadedFunctionType} overloadedType - The overloaded function type
+   * @param {Object} callNode - AST node of the function call
+   * @returns {FunctionType|null} The resolved function type, or null if no match
+   */
+  resolveUserOverloadStatic(funcName, argTypes, overloadedType, callNode) {
+    // Phase 1: Identify concrete vs ambiguous arguments
+    const isAmbiguousArg = (argType) => {
+      if (argType instanceof TypeVariable) return true;
+      if (argType instanceof FunctionType) {
+        return argType.paramTypes.some(p => p instanceof TypeVariable && !p.resolved) ||
+               (argType.returnType instanceof TypeVariable && !argType.returnType.resolved);
+      }
+      if (argType instanceof ArrayType) {
+        // Check if element type contains unresolved TypeVariables
+        return this.containsUnresolvedTypeVar(argType.elementType);
+      }
+      return false;
+    };
+
+    // Phase 1: Filter overloads using only concrete arguments
+    let phase1Candidates = [];
+    for (const { funcType, node, overloadId } of overloadedType.userOverloads) {
+      if (funcType.paramTypes.length !== argTypes.length) continue;
+
+      let valid = true;
+      for (let i = 0; i < argTypes.length; i++) {
+        const paramType = funcType.paramTypes[i];
+        const argType = argTypes[i];
+
+        if (isAmbiguousArg(argType)) continue;
+
+        if (this.typesEqual(paramType, argType)) {
+          // Exact match
+        } else if (this.canPromote(argType, paramType).can) {
+          // Promotion possible
+        } else if (this.typesCompatible(paramType, argType)) {
+          // Compatible
+        } else {
+          valid = false;
+          break;
+        }
+      }
+
+      if (valid) {
+        phase1Candidates.push({ funcType, node, overloadId });
+      }
+    }
+
+    // Phase 1.5: If multiple candidates, use function type structure to filter
+    if (phase1Candidates.length > 1) {
+      const funcArgIndices = [];
+      for (let i = 0; i < argTypes.length; i++) {
+        if (argTypes[i] instanceof FunctionType) funcArgIndices.push(i);
+      }
+
+      if (funcArgIndices.length > 0) {
+        const phase15Candidates = [];
+
+        for (const candidate of phase1Candidates) {
+          let compatible = true;
+
+          for (const funcIdx of funcArgIndices) {
+            const expectedFuncType = candidate.funcType.paramTypes[funcIdx];
+            const actualFuncType = argTypes[funcIdx];
+
+            if (expectedFuncType instanceof FunctionType && actualFuncType instanceof FunctionType) {
+              if (expectedFuncType.paramTypes.length !== actualFuncType.paramTypes.length) {
+                compatible = false;
+                break;
+              }
+
+              for (let p = 0; p < actualFuncType.paramTypes.length; p++) {
+                let actualParamType = actualFuncType.paramTypes[p];
+                if (actualParamType instanceof TypeVariable) actualParamType = actualParamType.resolve();
+                if (!(actualParamType instanceof TypeVariable)) {
+                  if (!this.typesCompatibleOrUnifiable(actualParamType, expectedFuncType.paramTypes[p])) {
+                    compatible = false;
+                    break;
+                  }
+                }
+              }
+
+              let actualReturnType = actualFuncType.returnType;
+              if (actualReturnType instanceof TypeVariable) actualReturnType = actualReturnType.resolve();
+              if (!(actualReturnType instanceof TypeVariable)) {
+                if (!this.typesCompatibleOrUnifiable(actualReturnType, expectedFuncType.returnType)) {
+                  compatible = false;
+                }
+              }
+            }
+          }
+
+          if (compatible) {
+            phase15Candidates.push(candidate);
+          }
+        }
+
+        if (phase15Candidates.length === 1) {
+          phase1Candidates = phase15Candidates;
+        } else if (phase15Candidates.length > 0 && phase15Candidates.length < phase1Candidates.length) {
+          phase1Candidates = phase15Candidates;
+        }
+      }
+    }
+
+    // If phase 1 narrowed to exactly one overload, use bidirectional inference
+    if (phase1Candidates.length === 1) {
+      const best = phase1Candidates[0];
+
+      for (let i = 0; i < argTypes.length; i++) {
+        const argType = argTypes[i];
+        const paramType = best.funcType.paramTypes[i];
+
+        if (isAmbiguousArg(argType)) {
+          unify(argType, paramType);
+        }
+      }
+
+      callNode.resolvedUserOverload = {
+        funcType: best.funcType,
+        overloadId: best.overloadId,
+      };
+      return best.funcType;
+    }
+
+
+    const candidates = [];
+
+    for (const { funcType, node, overloadId } of overloadedType.userOverloads) {
+      // Skip if arity doesn't match
+      if (funcType.paramTypes.length !== argTypes.length) continue;
+
+      // Score based on type compatibility
+      let score = 0;
+      let valid = true;
+
+      for (let i = 0; i < argTypes.length; i++) {
+        const paramType = funcType.paramTypes[i];
+        const argType = argTypes[i];
+
+        // Try to check type compatibility
+        if (this.typesEqual(paramType, argType)) {
+          // Exact match
+          score += 0;
+        } else if (this.containsUnresolvedTypeVar(paramType) || this.containsUnresolvedTypeVar(argType)) {
+          // Type contains unresolved type variable - can unify, slight penalty
+          score += 0.5;
+        } else if (this.canPromote(argType, paramType).can) {
+          // Promotion possible (num -> complex or user-defined)
+          score += 1;
+        } else if (this.typesCompatible(paramType, argType)) {
+          // Compatible but not exact (e.g., array rank matching)
+          score += 0.25;
+        } else if (this.canUnifyFunctionTypes(paramType, argType)) {
+          // Function types that can unify (e.g., (num, array) -> array with (T, T) -> array)
+          score += 0.5;
+        } else {
+          valid = false;
+          break;
+        }
+      }
+
+      if (valid) {
+        candidates.push({ funcType, node, score, overloadId });
+      }
+    }
+
+    if (candidates.length === 0) {
+      const argTypeStrs = argTypes.map(t => formatType(t)).join(", ");
+      const available = overloadedType.userOverloads
+        .map(({ funcType }) => `  ${funcType.toString()}`)
+        .join("\n");
+      this.addError(
+        `No matching overload for '${funcName}(${argTypeStrs})'\nAvailable overloads:\n${available}`,
+        callNode?.location
+      );
+      return null;
+    }
+
+    // Select best (lowest score)
+    candidates.sort((a, b) => a.score - b.score);
+
+    if (candidates.length > 1 && candidates[0].score === candidates[1].score) {
+      this.addError(
+        `Ambiguous overload for '${funcName}' - multiple overloads match with equal specificity`,
+        callNode?.location
+      );
+    }
+
+    const best = candidates[0];
+
+    // Store resolved overload on AST node for executor
+    // The overloadId allows executor to look up the specific definition
+    callNode.resolvedUserOverload = {
+      funcType: best.funcType,
+      overloadId: best.overloadId,
+    };
+
+    return best.funcType;
+  },
+};