stone-lang 0.1.0

package/frontend/type-checker/visitors/functions.js

@@ -0,0 +1,604 @@
+/**
+ * Function Visitors - mixin for TypeChecker
+ *
+ * Visitors for function definitions and function calls.
+ */
+
+import {
+  FunctionType,
+  RecordType,
+  ArrayType,
+  TypeVariable,
+  PromotedType,
+  LiteralType,
+  LiteralUnionType,
+  UNIT,
+  freshTypeVar,
+  unify,
+  formatType,
+  fromTypeAnnotation,
+  typeEquals,
+  isAssignableTo,
+} from "../../types/types.js";
+import { OverloadedFunctionType } from "../OverloadedFunctionType.js";
+import { ExtensionInfo } from "../TypeEnvironment.js";
+
+/**
+ * Function visitor methods to be mixed into TypeChecker.prototype
+ */
+export const functionVisitors = {
+  /**
+   * Visit a function definition
+   */
+  visitFunctionDefinition(node, env) {
+    const funcEnv = env.createChild();
+
+    // Start tracking deferred constraints for this function
+    this.pushConstraintContext();
+
+    // Process parameters and create a map from parameter name to index
+    const paramTypes = [];
+    const paramNameToIndex = new Map();
+    for (let i = 0; i < node.parameters.length; i++) {
+      const param = node.parameters[i];
+      let paramType;
+      if (param.typeAnnotation) {
+        paramType = fromTypeAnnotation(param.typeAnnotation, env);
+
+        // Validate promoted parameters: look up converter(s) and set sourceType(s)
+        if (paramType instanceof PromotedType) {
+          // Validate that paramRef matches the parameter name
+          if (paramType.paramRef !== param.name) {
+            this.addError(
+              `Promotion converter references '${paramType.paramRef}' but parameter is '${param.name}'`,
+              param.location
+            );
+          }
+
+          // Validate each converter in the multi-converter list
+          const allSourceTypes = [];
+          for (const converter of paramType.converters) {
+            const converterType = env.lookup(converter.funcName);
+            if (!converterType) {
+              this.addError(
+                `Converter function '${converter.funcName}' not found`,
+                param.location
+              );
+              continue;
+            }
+
+            if (converterType instanceof FunctionType) {
+              // Single function - validate it has exactly one parameter
+              if (converterType.paramTypes.length !== 1) {
+                this.addError(
+                  `Converter '${converter.funcName}' must take exactly one parameter`,
+                  param.location
+                );
+              } else {
+                // Set the sourceType from the converter's input type
+                converter.sourceType = converterType.paramTypes[0];
+
+                // Validate that converter returns the target type
+                if (!typeEquals(converterType.returnType, paramType.targetType)) {
+                  this.addError(
+                    `Converter '${converter.funcName}' returns ${formatType(converterType.returnType)}, expected ${formatType(paramType.targetType)}`,
+                    param.location
+                  );
+                }
+
+                allSourceTypes.push({ type: converter.sourceType, funcName: converter.funcName });
+              }
+            } else if (converterType instanceof OverloadedFunctionType) {
+              // Overloaded function - find matching overloads
+              const validSourceTypes = [];
+
+              // Check user-defined overloads
+              for (const { funcType } of converterType.userOverloads) {
+                if (funcType.paramTypes.length === 1 &&
+                    typeEquals(funcType.returnType, paramType.targetType)) {
+                  validSourceTypes.push(funcType.paramTypes[0]);
+                }
+              }
+
+              // Check primitive overloads
+              for (const primFn of converterType.overloads) {
+                if (primFn.stone && primFn.stone.in.length === 1) {
+                  const outType = fromTypeAnnotation({ type: 'TypeAnnotation', name: primFn.stone.out }, env);
+                  if (typeEquals(outType, paramType.targetType)) {
+                    const inType = fromTypeAnnotation({ type: 'TypeAnnotation', name: primFn.stone.in[0] }, env);
+                    validSourceTypes.push(inType);
+                  }
+                }
+              }
+
+              if (validSourceTypes.length === 0) {
+                this.addError(
+                  `No overload of '${converter.funcName}' takes one parameter and returns ${formatType(paramType.targetType)}`,
+                  param.location
+                );
+              } else {
+                converter.sourceType = validSourceTypes[0];
+                converter.validSourceTypes = validSourceTypes;
+                allSourceTypes.push({ type: validSourceTypes[0], funcName: converter.funcName });
+              }
+            } else {
+              this.addError(
+                `'${converter.funcName}' is not a function`,
+                param.location
+              );
+            }
+          }
+
+          // Check for overlapping source types (error if two converters accept the same type)
+          for (let i = 0; i < allSourceTypes.length; i++) {
+            for (let j = i + 1; j < allSourceTypes.length; j++) {
+              if (allSourceTypes[i].type && allSourceTypes[j].type &&
+                  typeEquals(allSourceTypes[i].type, allSourceTypes[j].type)) {
+                this.addError(
+                  `Overlapping source types: converters '${allSourceTypes[i].funcName}' and '${allSourceTypes[j].funcName}' both accept ${formatType(allSourceTypes[i].type)}`,
+                  param.location
+                );
+              }
+            }
+          }
+
+          // For backwards compatibility, set the legacy single sourceType
+          if (paramType.converters.length > 0 && paramType.converters[0].sourceType) {
+            paramType.sourceType = paramType.converters[0].sourceType;
+          }
+        }
+      } else {
+        paramType = freshTypeVar();
+      }
+      paramTypes.push(paramType);
+
+      // For promoted types, define the parameter with the target type in the function scope
+      let effectiveType = paramType instanceof PromotedType
+        ? paramType.targetType
+        : paramType;
+
+      // If the effective type is a schema (RecordType with requiredFields), convert to instance type
+      // This allows field access in the function body without "required type descriptor" errors
+      if (effectiveType instanceof RecordType && effectiveType.requiredFields) {
+        effectiveType = new RecordType(
+          effectiveType.fields,  // Same fields
+          false,                 // Not open
+          false,                 // Not explicit
+          null                   // No required fields - this is an instance, not a schema
+        );
+      }
+      funcEnv.define(param.name, effectiveType);
+      funcEnv.markImmutable(param.name);  // Function parameters are immutable
+      paramNameToIndex.set(param.name, i);
+    }
+
+    // Infer return type from body
+    let returnType;
+    if (node.isExpression) {
+      // First visit the body to type-check it
+      returnType = this.visit(node.body, funcEnv);
+
+      // For expression bodies that may contain return statements (e.g., block expressions with if/else),
+      // also collect and unify all return statements to ensure consistent types
+      const returnInfos = this.collectReturnStatements([node.body]);
+      if (returnInfos.length > 0) {
+        // Unify all return statement types
+        const { unifiedType, errors } = this.unifyReturnTypes(returnInfos);
+        for (const error of errors) {
+          this.addError(error.message, error.location);
+        }
+        // Use the unified return type instead of the expression type
+        returnType = unifiedType instanceof TypeVariable ? unifiedType.resolve() : unifiedType;
+      }
+    } else {
+      // Block function - look for return statements
+      returnType = this.visitFunctionBody(node.body, funcEnv);
+    }
+
+    // Collect deferred constraints from body analysis
+    const collectedConstraints = this.popConstraintContext();
+
+    // Check against declared return type
+    if (node.returnType) {
+      const declaredReturnType = fromTypeAnnotation(node.returnType, env);
+
+      // For literal types, use strict assignability check (inferred must be subset of declared)
+      // For other types, use unify for compatibility
+      const isLiteralReturn = declaredReturnType instanceof LiteralType ||
+                              declaredReturnType instanceof LiteralUnionType;
+
+      if (isLiteralReturn) {
+        // Strict check: inferred return type must be assignable to declared type
+        if (!isAssignableTo(returnType, declaredReturnType)) {
+          this.addError(
+            `Function '${node.name}' declared to return ${formatType(declaredReturnType)} but inferred ${formatType(returnType)}`,
+            node.location
+          );
+        }
+      } else {
+        // Use unify for regular types
+        const unifyResult = unify(declaredReturnType, returnType);
+        if (!unifyResult.success) {
+          this.addError(
+            `Function '${node.name}' declared to return ${formatType(declaredReturnType)} but inferred ${formatType(returnType)}`,
+            node.location
+          );
+        }
+      }
+      returnType = declaredReturnType;
+    }
+
+    // Resolve any TypeVariables that were constrained during body analysis
+    // This makes function signatures more precise based on actual usage
+    const resolvedParamTypes = paramTypes.map((t, i) => {
+      if (t instanceof TypeVariable) {
+        const resolved = t.resolve();
+        return resolved;
+      }
+      return t;
+    });
+    const resolvedReturnType = returnType instanceof TypeVariable
+      ? returnType.resolve()
+      : returnType;
+
+    // Create function type with any collected deferred constraints
+    const funcType = new FunctionType(resolvedParamTypes, resolvedReturnType, collectedConstraints);
+
+    // Support function overloading: accumulate overloads instead of overwriting
+    const existing = env.lookup(node.name);
+
+    // Note: We allow user functions to shadow primitives. The executor will use
+    // the pre-resolved overloadId from type checking, falling back to primitives
+    // only if no user overload matches. If a user overload was resolved but can't
+    // be found at runtime, the executor throws a clear error.
+
+    // Generate overloadId for this definition using location or a counter as unique identifier
+    let locId;
+    if (node.location && node.location.start) {
+      locId = `${node.location.start.line}_${node.location.start.column}`;
+    } else if (node.location && typeof node.location.line === 'number') {
+      locId = `${node.location.line}_${node.location.column || 0}`;
+    } else {
+      // Fallback to a counter
+      this._overloadCounter = (this._overloadCounter || 0) + 1;
+      locId = `gen_${this._overloadCounter}`;
+    }
+    const overloadId = `${node.name}_${locId}`;
+    node.overloadId = overloadId;
+    // Store overloadId on funcType so we can retrieve it if this becomes overloaded
+    funcType._overloadId = overloadId;
+    funcType._sourceNode = node;
+
+    if (existing instanceof OverloadedFunctionType) {
+      // Add to existing user-defined overload set
+      existing.addUserOverload(funcType, node);
+    } else if (existing instanceof FunctionType) {
+      // Convert to overloaded: previous definition + new one
+      const overloaded = new OverloadedFunctionType(node.name, [], []);
+      // Use the stored overloadId and node from the first definition
+      const firstOverloadId = existing._overloadId || `${node.name}_first`;
+      const firstNode = existing._sourceNode || null;
+      overloaded.userOverloads.push({
+        funcType: existing,
+        node: firstNode,
+        overloadId: firstOverloadId
+      });
+      overloaded.addUserOverload(funcType, node);
+      env.define(node.name, overloaded);
+    } else {
+      // First definition - store as regular FunctionType
+      env.define(node.name, funcType);
+    }
+
+    // If this is an extension function, also register it in the extension map
+    if (node.isExtension && node.selfType) {
+      const extensionInfo = new ExtensionInfo(
+        node.name,
+        node.selfType,
+        funcType
+      );
+      env.defineExtension(node.name, extensionInfo);
+    }
+
+    return this.setType(node, funcType);
+  },
+
+  /**
+   * Visit function body to find return type using two-pass approach
+   * 1. Collect all return statements with their types
+   * 2. Unify all complete types, report mismatches
+   * 3. Back-propagate unified type to incomplete expressions (e.g., empty arrays)
+   */
+  visitFunctionBody(body, env) {
+    // First, visit all statements to type-check the body
+    for (const stmt of body) {
+      this.visit(stmt, env);
+    }
+
+    // Collect all return statements (they've already been visited above)
+    const returnInfos = [];
+    const collectReturns = (statements, isGuardPassthrough = false) => {
+      for (const stmt of statements) {
+        if (stmt.type === "ReturnStatement") {
+          const type = stmt.inferredType || (stmt.value ? stmt.value.inferredType : UNIT) || freshTypeVar();
+          const isIncomplete = this.isIncompleteType(type);
+          returnInfos.push({ node: stmt, type, isIncomplete, isGuardPassthrough });
+        } else if (stmt.type === "BranchExpression") {
+          // Check for guard passthrough pattern
+          // Pattern: if (x.field != "literal") { return x } - early return guard
+          let guardPassthroughInFirstBranch = false;
+          if (stmt.paths.length >= 1 && stmt.paths[0].condition) {
+            const discriminant = this.extractDiscriminantCheck(stmt.paths[0].condition);
+            if (discriminant && discriminant.isNegated) {
+              const returnVar = this.getSimpleReturnVar(stmt.paths[0].body);
+              if (returnVar === discriminant.varName) {
+                guardPassthroughInFirstBranch = true;
+              }
+            }
+          }
+
+          for (let pathIdx = 0; pathIdx < stmt.paths.length; pathIdx++) {
+            const branchPath = stmt.paths[pathIdx];
+            const isGuardPath = guardPassthroughInFirstBranch && pathIdx === 0;
+            collectReturns(branchPath.body, isGuardPath);
+          }
+        } else if (stmt.type === "LoopExpression") {
+          collectReturns(stmt.body);
+        } else if (stmt.type === "BlockExpression") {
+          collectReturns(stmt.statements || []);
+        }
+      }
+    };
+    collectReturns(body);
+
+    // If no return statements, return UNIT
+    if (returnInfos.length === 0) {
+      return UNIT;
+    }
+
+    // Unify return types
+    const { unifiedType, errors } = this.unifyReturnTypes(returnInfos);
+
+    // Report any errors
+    for (const error of errors) {
+      this.addError(error.message, error.location);
+    }
+
+    // Back-propagate unified type to incomplete expressions
+    this.resolveIncompleteTypes(returnInfos, unifiedType);
+
+    return unifiedType instanceof TypeVariable ? unifiedType.resolve() : unifiedType;
+  },
+
+  /**
+   * Visit a function call
+   * Uses static overload resolution for primitive functions
+   */
+  visitFunctionCall(node, env) {
+    if (node.callee) {
+      // Method call - get the method's FunctionType from visiting the callee
+      const calleeType = this.visit(node.callee, env);
+      const argTypes = (node.args || []).map((arg) => this.visit(arg, env));
+
+      // If the callee is a FunctionType, unify argument types and return the return type
+      if (calleeType instanceof FunctionType) {
+        // Special handling for push method on arrays without declared type:
+        // Infer and store the element type from the first push argument
+        if (node.callee.type === "MemberAccess" &&
+            node.callee.property?.name === "push" &&
+            node.callee.object?.type === "Identifier") {
+          let arrayVarName = node.callee.object.name;
+          // For primed variables (arr'), use base name for declared type storage
+          const baseName = arrayVarName.endsWith("'") ? arrayVarName.slice(0, -1) : arrayVarName;
+          const arrayType = env.lookup(arrayVarName);
+
+          // If array has unresolved element type and no declared type, infer from push arg
+          if (arrayType instanceof ArrayType &&
+              arrayType.elementType instanceof TypeVariable &&
+              !arrayType.elementType.resolved &&
+              !env.lookupDeclared(baseName) &&
+              argTypes.length > 0) {
+            // The first argument determines the array's element type
+            const argType = argTypes[0];
+            // Store as the "inferred declared type" for future pushes (use base name)
+            const inferredArrayType = new ArrayType(argType, arrayType.rank);
+            env.defineDeclared(baseName, inferredArrayType);
+            // Also unify the TypeVariable to resolve it
+            unify(arrayType.elementType, argType);
+          }
+        }
+
+        // Check argument types against parameter types
+        for (let i = 0; i < Math.min(argTypes.length, calleeType.paramTypes.length); i++) {
+          const unifyResult = unify(calleeType.paramTypes[i], argTypes[i]);
+          if (!unifyResult.success) {
+            // Only error if both types are fully resolved (no unresolved type variables)
+            const paramHasTypeVar = this.containsUnresolvedTypeVar(calleeType.paramTypes[i]);
+            const argHasTypeVar = this.containsUnresolvedTypeVar(argTypes[i]);
+            if (!paramHasTypeVar && !argHasTypeVar) {
+              this.addError(
+                `Method argument ${i + 1} type mismatch: expected ${formatType(calleeType.paramTypes[i])}, got ${formatType(argTypes[i])}`,
+                node.location
+              );
+            }
+          }
+        }
+        return this.setType(node, calleeType.returnType);
+      }
+
+      // Unknown callee type - return fresh type
+      return this.setType(node, freshTypeVar());
+    }
+
+    const funcName = node.funcName;
+    const argTypes = (node.args || []).map((arg) => this.visit(arg, env));
+
+    // Check environment first for user-defined functions
+    let funcType = env.lookup(funcName);
+
+    // Instantiate fresh type variables for each call (polymorphism)
+    // This ensures that calls like print("str") and print([1,2]) work independently
+    if (funcType instanceof FunctionType) {
+      funcType = funcType.instantiate();
+    }
+
+    // Handle user-defined function overloads
+    // If the function has multiple user-defined overloads, resolve statically
+    if (funcType instanceof OverloadedFunctionType && funcType.userOverloads.length > 0) {
+      const resolved = this.resolveUserOverloadStatic(funcName, argTypes, funcType, node);
+      if (resolved) {
+        return this.setType(node, resolved.returnType);
+      }
+      // If resolution failed, error was already added - return fresh type
+      return this.setType(node, freshTypeVar());
+    }
+
+    // Check if this function is defined in the scope chain as a FunctionType (not just an overload).
+    // User-defined functions shadow builtins.
+    const isUserDefinedFunction = funcType instanceof FunctionType && this.isDefinedInScope(env, funcName);
+
+    // If it's a user-defined FunctionType, use it directly.
+    // This allows user-defined functions to shadow builtins.
+    if (isUserDefinedFunction) {
+      // Track promotions for codegen
+      const promotions = [];
+
+      // Check argument types against function signature
+      if (argTypes.length !== funcType.paramTypes.length) {
+        this.addError(
+          `Function '${funcName}' expects ${funcType.paramTypes.length} arguments, got ${argTypes.length}`,
+          node.location
+        );
+      } else {
+        for (let i = 0; i < argTypes.length; i++) {
+          const unifyResult = unify(funcType.paramTypes[i], argTypes[i]);
+
+          // Track promotion info (converter function name or null)
+          promotions.push(unifyResult.promoted ? unifyResult.converter : null);
+
+          if (!unifyResult.success) {
+            const paramType = funcType.paramTypes[i];
+            const argType = argTypes[i];
+            const paramHasTypeVar = this.containsUnresolvedTypeVar(paramType);
+            const argHasTypeVar = this.containsUnresolvedTypeVar(argType);
+
+            // Check if param type is an explicit constraint (from destructuring)
+            // Explicit constraints should always report errors, while implicit constraints
+            // (from member access inference) may be flow-dependent and shouldn't error
+            const resolvedParam = paramType.resolve ? paramType.resolve() : paramType;
+            const isExplicitConstraint = resolvedParam instanceof RecordType && resolvedParam.explicit;
+
+            // "Missing required field" errors from explicit destructuring constraints
+            // should always be reported - the user explicitly asked for those fields
+            const isMissingRequiredFieldError = unifyResult.error &&
+              unifyResult.error.includes('Missing required field') &&
+              isExplicitConstraint;
+
+            // Other structural errors - only report when param type is fully resolved
+            // to avoid false positives when type variables could match different types
+            const isOtherStructuralError = !paramHasTypeVar && unifyResult.error &&
+              !isMissingRequiredFieldError &&
+              (unifyResult.error.includes('Missing') ||
+               unifyResult.error.includes('field'));
+
+            const isStructuralError = isMissingRequiredFieldError || isOtherStructuralError;
+
+            // Report error if:
+            // 1. Both types are fully resolved (no type variables), OR
+            // 2. Type kinds are incompatible (e.g., function vs primitive) and arg is concrete, OR
+            // 3. It's a structural error (missing required fields) AND param is fully resolved
+            const kindsIncompatible = paramType.kind !== argType.kind &&
+                                      !(paramType instanceof TypeVariable) &&
+                                      !(argType instanceof TypeVariable);
+            const argIsConcrete = !argHasTypeVar;
+
+            if ((!paramHasTypeVar && !argHasTypeVar) || (kindsIncompatible && argIsConcrete) || isStructuralError) {
+              const errorMsg = isStructuralError
+                ? `Argument ${i + 1}: ${unifyResult.error}`
+                : `Argument ${i + 1} type mismatch: expected ${formatType(paramType)}, got ${formatType(argType)}`;
+              this.addError(errorMsg, node.location);
+            }
+          }
+        }
+
+        // Check deferred constraints from function definition
+        this.checkDeferredConstraints(funcName, funcType, argTypes, node);
+      }
+
+      // Store promotions on the call node for codegen
+      // Use resolvedUserOverload to match the pattern in compiler.js
+      if (promotions.some(p => p !== null)) {
+        node.resolvedUserOverload = { promotions };
+      }
+
+      return this.setType(node, funcType.returnType);
+    }
+
+    // Special handling for parse(raw, schema) - returns an instance of the schema type
+    // Works for parse(string, schema) and parse(record, schema)
+    if (funcName === 'parse' && argTypes.length === 2 && argTypes[1] instanceof RecordType) {
+      // parse(raw, schema) returns a record matching the schema type
+      // The second argument is the schema record, but we return an "instance" type
+      // (same fields, but without requiredFields markers so field access is allowed)
+      const schemaType = argTypes[1];
+      const instanceType = new RecordType(
+        schemaType.fields,      // Same fields
+        false,                  // Not open
+        false,                  // Not explicit
+        null                    // No required fields - this is an instance, not a schema
+      );
+      return this.setType(node, instanceType);
+    }
+
+    // Try overload resolution for primitives or OverloadedFunctionType
+    if (this.overloads[funcName] || funcType instanceof OverloadedFunctionType) {
+      // Use function's overloads if available, otherwise use global overloads
+      const overloadsToUse = funcType instanceof OverloadedFunctionType
+        ? funcType.overloads
+        : this.overloads[funcName];
+
+      // Temporarily set overloads for resolution
+      const originalOverloads = this.overloads[funcName];
+      this.overloads[funcName] = overloadsToUse;
+
+      const outputType = this.resolveOverloadStatic(funcName, argTypes, node);
+
+      // Restore original overloads
+      if (originalOverloads) {
+        this.overloads[funcName] = originalOverloads;
+      } else if (funcType instanceof OverloadedFunctionType) {
+        delete this.overloads[funcName];
+      }
+
+      if (outputType) {
+        const resultType = this.fromStructuredType(outputType);
+        return this.setType(node, resultType);
+      }
+
+      // Overload resolution failed - error already added
+      return this.setType(node, freshTypeVar());
+    }
+
+    // If it's a type variable (e.g., function parameter), constrain it to FunctionType
+    // This enables type inference for higher-order functions like ODE solvers
+    if (funcType instanceof TypeVariable) {
+      const returnTypeVar = freshTypeVar();
+      const expectedFuncType = new FunctionType(argTypes, returnTypeVar);
+      // Try to unify - this constrains the TypeVariable to be a function type
+      unify(funcType, expectedFuncType);
+      return this.setType(node, returnTypeVar);
+    }
+
+    // Not found in environment and not an overload
+    if (!funcType) {
+      return this.setType(node, freshTypeVar());
+    }
+
+    // Unknown type - not a function
+    this.addError(
+      `'${funcName}' is not a function`,
+      node.location
+    );
+    return this.setType(node, freshTypeVar());
+  },
+};

package/frontend/type-checker/visitors/index.js

@@ -0,0 +1,38 @@
+/**
+ * Visitor Registration - aggregates all visitor modules
+ *
+ * Registers all visitor methods on the TypeChecker prototype.
+ */
+
+import { literalVisitors } from './literals.js';
+import { operatorVisitors } from './operators.js';
+import { functionVisitors } from './functions.js';
+import { controlFlowVisitors } from './controlFlow.js';
+import { statementVisitors } from './statements.js';
+import { moduleVisitors } from './modules.js';
+
+/**
+ * Register all visitor methods on the TypeChecker prototype.
+ *
+ * @param {Function} TypeChecker - The TypeChecker class constructor
+ */
+export function registerVisitors(TypeChecker) {
+  Object.assign(TypeChecker.prototype,
+    literalVisitors,
+    operatorVisitors,
+    functionVisitors,
+    controlFlowVisitors,
+    statementVisitors,
+    moduleVisitors
+  );
+}
+
+// Re-export individual visitor modules for testing
+export {
+  literalVisitors,
+  operatorVisitors,
+  functionVisitors,
+  controlFlowVisitors,
+  statementVisitors,
+  moduleVisitors,
+};