stone-lang 0.1.0

package/frontend/type-checker/typeChecker.js

@@ -0,0 +1,452 @@
+/**
+ * Stone Type Checker
+ *
+ * Performs compile-time type checking on Stone AST.
+ *
+ * Design principles:
+ * - Types are ALWAYS optional (no errors for missing annotations)
+ * - Compile-time errors when annotated types don't match inferred types
+ * - Infers types from usage context
+ * - Stores inferred types on AST nodes for IDE tooling
+ */
+
+import {
+  FunctionType,
+  RecordType,
+  ArrayType,
+  TypeVariable,
+  INT,
+  FLOAT,
+  BOOL,
+  STRING,
+  UNIT,
+  NUM,
+  freshTypeVar,
+  flattenNestedArrayType,
+} from "../types/types.js";
+import { OVERLOADS } from "../../backends/js-vm/primitives/overloads.js";
+
+// Import modular components
+import { TypeEnvironment } from "./TypeEnvironment.js";
+import { OverloadedFunctionType } from "./OverloadedFunctionType.js";
+import { typeCompatibilityMethods } from "./typeCompatibility.js";
+import { bidirectionalInferenceMethods } from "./bidirectionalInference.js";
+import { overloadResolutionMethods } from "./overloadResolution.js";
+import { moduleAnalysisMethods } from "./moduleAnalysis.js";
+import { registerVisitors } from "./visitors/index.js";
+
+// Re-export for external consumers
+export { OverloadedFunctionType };
+
+/**
+ * Stone Type Checker
+ *
+ * Performs compile-time type checking on Stone AST using Hindley-Milner style
+ * type inference with structural typing.
+ *
+ * Design principles:
+ * - Types are always optional (no errors for missing annotations)
+ * - Compile-time errors when annotated types don't match inferred types
+ * - Infers types from usage context
+ * - Stores inferred types on AST nodes (node.inferredType) for tooling and optimization
+ *
+ * Type system features:
+ * - Primitive types: num, bool, string, unit, complex
+ * - Compound types: arrays (with rank), records, functions
+ * - Type variables for polymorphism
+ * - Overloaded functions with signature matching
+ * - Module import type analysis
+ *
+ * @example
+ * const checker = new TypeChecker({ moduleResolver, parser });
+ * const result = checker.check(ast);
+ * if (!result.success) {
+ *   console.error(result.errors);
+ * }
+ */
+export class TypeChecker {
+  /**
+   * Create a new TypeChecker instance
+   *
+   * @param {Object} [options={}] - Configuration options
+   * @param {ModuleResolver} [options.moduleResolver] - Module resolver for import type analysis
+   * @param {Function} [options.parser] - Parser function for analyzing module source
+   */
+  constructor(options = {}) {
+    /**
+     * Collected type errors
+     * @type {Array<{message: string, location: Object|null}>}
+     */
+    this.errors = [];
+
+    /**
+     * Collected warnings
+     * @type {Array<{message: string, location: Object|null}>}
+     */
+    this.warnings = [];
+
+    /**
+     * Map from AST node to inferred type
+     * @type {Map<Object, Object>}
+     */
+    this.typeMap = new Map();
+
+    /**
+     * Configuration options
+     * @type {Object}
+     */
+    this.options = options;
+
+    /**
+     * Overload registry for static resolution.
+     * Creates a shallow copy to avoid modifying the global OVERLOADS object
+     * when adding type-checking-only overloads.
+     * @type {Object}
+     */
+    this.overloads = {};
+    for (const name of Object.keys(OVERLOADS)) {
+      this.overloads[name] = [...OVERLOADS[name]];
+    }
+
+    /**
+     * Module resolver for import type analysis (optional)
+     * @type {ModuleResolver|null}
+     */
+    this.moduleResolver = options.moduleResolver || null;
+
+    /**
+     * Parser function for analyzing module source (optional)
+     * @type {Function|null}
+     */
+    this.parser = options.parser || null;
+
+    /**
+     * Cache of module export types (modulePath -> Map<exportName, Type>)
+     * @type {Map<string, Map<string, Object>>}
+     */
+    this.moduleTypeCache = new Map();
+
+    /**
+     * Map of TypeVariable id to TypeVariable object for bidirectional inference
+     * @type {Map<number, TypeVariable>}
+     */
+    this.typeVarMap = new Map();
+
+    /**
+     * Stack of deferred constraint collections for nested function definitions.
+     * Each entry is an array of constraints for the currently-being-analyzed function.
+     * @type {Array<Array>}
+     */
+    this.constraintStack = [];
+  }
+
+  /**
+   * Record a deferred constraint for the currently-being-analyzed function.
+   * Called when scoreOverload encounters a variable rank that must match a concrete rank.
+   *
+   * @param {Object} constraint - The deferred constraint
+   */
+  recordDeferredConstraint(constraint) {
+    if (this.constraintStack.length > 0) {
+      this.constraintStack[this.constraintStack.length - 1].push(constraint);
+    }
+  }
+
+  /**
+   * Start tracking constraints for a new function definition.
+   */
+  pushConstraintContext() {
+    this.constraintStack.push([]);
+  }
+
+  /**
+   * Finish tracking constraints and return collected constraints.
+   * @returns {Array} The deferred constraints collected for this function
+   */
+  popConstraintContext() {
+    return this.constraintStack.pop() || [];
+  }
+
+  /**
+   * Register complex module arithmetic overloads for type checking.
+   * Called when any import from 'complex' module is seen.
+   * This allows operators like + - * / to work with complex numbers.
+   */
+  _registerComplexOverloads() {
+    // Don't register twice
+    if (this._complexOverloadsRegistered) return;
+    this._complexOverloadsRegistered = true;
+
+    // Create mock "primitive" functions with complex signatures for type checking
+    // These signatures match what the Stone complex module provides
+    // Register under both function names (for import) and operator names (for +, -, etc.)
+    const complexOps = [
+      { name: 'add', opName: '__op_add', sig: { in: ['complex', 'complex'], out: 'complex' } },
+      { name: 'sub', opName: '__op_sub', sig: { in: ['complex', 'complex'], out: 'complex' } },
+      { name: 'mul', opName: '__op_mul', sig: { in: ['complex', 'complex'], out: 'complex' } },
+      { name: 'div', opName: '__op_div', sig: { in: ['complex', 'complex'], out: 'complex' } },
+      { name: 'neg', opName: '__op_neg', sig: { in: ['complex'], out: 'complex' } },
+      { name: 'pow', opName: '__op_pow', sig: { in: ['complex', 'complex'], out: 'complex' } },
+    ];
+
+    for (const { name, opName, sig } of complexOps) {
+      // Create a mock function with the .stone signature
+      const mockFn = { stone: sig };
+
+      // Add to overloads under function name (for import from complex module)
+      if (!this.overloads[name]) {
+        this.overloads[name] = [];
+      }
+      this.overloads[name].push(mockFn);
+
+      // Also add under operator name (for +, -, *, /, ^ operators)
+      if (!this.overloads[opName]) {
+        this.overloads[opName] = [];
+      }
+      this.overloads[opName].push({ stone: sig });
+    }
+  }
+
+  /**
+   * Check deferred constraints from a function definition against actual argument types.
+   * Called at function call sites to verify that inferred rank constraints are satisfied.
+   *
+   * @param {string} funcName - The function name (for error messages)
+   * @param {FunctionType} funcType - The function type with constraints
+   * @param {Array} argTypes - The actual argument types at the call site
+   * @param {Object} node - The call site node (for error location)
+   */
+  checkDeferredConstraints(funcName, funcType, argTypes, node) {
+    if (!funcType.constraints || funcType.constraints.length === 0) {
+      return;
+    }
+
+    for (const constraint of funcType.constraints) {
+      if (constraint.type === 'rank') {
+        const paramIndex = constraint.paramIndex;
+        if (paramIndex >= argTypes.length) {
+          continue; // Arity mismatch handled elsewhere
+        }
+
+        let argType = argTypes[paramIndex];
+        let actualRank = null;
+
+        // Flatten nested array types before extracting rank
+        // e.g., array<array<num, 1>, 1> should have effective rank 2
+        if (argType instanceof ArrayType || (argType && argType.kind === 'array')) {
+          argType = flattenNestedArrayType(argType);
+        }
+
+        // Extract the rank from the argument type
+        if (argType instanceof ArrayType) {
+          actualRank = argType.rank;
+        } else if (argType && argType.kind === 'array') {
+          actualRank = argType.rank;
+        }
+
+        if (actualRank === null) {
+          // Argument isn't an array - type mismatch error handled elsewhere
+          continue;
+        }
+
+        // Check if rank is concrete and matches
+        if (typeof actualRank === 'number') {
+          if (actualRank !== constraint.requiredRank) {
+            this.addError(
+              `Argument ${paramIndex + 1} to '${funcName}' requires array of rank ${constraint.requiredRank}, but got rank ${actualRank}`,
+              node.location
+            );
+          }
+        } else if (typeof actualRank === 'object' && actualRank.var) {
+          // Argument also has variable rank - propagate the constraint
+          // This means we're calling a constrained function from another function with variable rank
+          this.recordDeferredConstraint({
+            type: 'rank',
+            paramIndex: paramIndex,
+            rankVar: actualRank.var,
+            requiredRank: constraint.requiredRank
+          });
+        }
+        // If rank is neither number nor variable object, it's a type we don't handle
+      }
+    }
+  }
+
+  /**
+   * Type check an AST and return type information
+   *
+   * This is the main entry point for type checking. It:
+   * 1. Creates a fresh type environment
+   * 2. Registers built-in types and functions
+   * 3. Visits all AST nodes, inferring and checking types
+   * 4. Returns collected errors, warnings, and type map
+   *
+   * @param {Object} ast - The parsed AST (Program node)
+   * @returns {Object} Type checking result
+   * @returns {Array} returns.errors - Array of type errors
+   * @returns {Array} returns.warnings - Array of type warnings
+   * @returns {Map} returns.types - Map from AST nodes to their inferred types
+   * @returns {boolean} returns.success - True if no errors were found
+   *
+   * @example
+   * const result = checker.check(ast);
+   * if (result.success) {
+   *   // Access inferred types via result.types or node.inferredType
+   *   const funcType = result.types.get(functionNode);
+   * }
+   */
+  check(ast) {
+    this.errors = [];
+    this.warnings = [];
+    this.typeMap = new Map();
+
+    const env = new TypeEnvironment();
+    this.registerBuiltins(env);
+
+    try {
+      this.visitProgram(ast, env);
+    } catch (error) {
+      this.errors.push({
+        message: error.message,
+        location: null,
+      });
+    }
+
+    // Collect extension names from the type environment
+    // This includes both locally defined extensions and imported ones
+    const extensions = new Set();
+    this._collectExtensions(env, extensions);
+
+    return {
+      errors: this.errors,
+      warnings: this.warnings,
+      types: this.typeMap,
+      success: this.errors.length === 0,
+      extensions, // Set of extension function names for the compiler
+    };
+  }
+
+  /**
+   * Recursively collect extension names from environment chain
+   * @private
+   */
+  _collectExtensions(env, extensions) {
+    if (!env) return;
+    for (const name of env.extensions.keys()) {
+      extensions.add(name);
+    }
+    // Note: extensions are typically only defined in the current env,
+    // not inherited, but we recurse for completeness
+    if (env.parent) {
+      this._collectExtensions(env.parent, extensions);
+    }
+  }
+
+  /**
+   * Register built-in functions and types.
+   * Note: Functions that have proper overload definitions in OVERLOADS are NOT registered here,
+   * as they will be handled by overload resolution which provides better type checking.
+   */
+  registerBuiltins(env) {
+    // Only register functions that are NOT in OVERLOADS or have no overloaded variants.
+    // Functions like abs, sqrt, pow, sin, cos, min, max etc. are in OVERLOADS and
+    // will be resolved through the overload system for proper type handling.
+
+    // Array functions (len has overloads, sum is simple)
+    env.define("len", new FunctionType([freshTypeVar()], INT));
+    env.define("sum", new FunctionType([new ArrayType(NUM, 1)], NUM));
+
+    // Type conversion (these are not overloaded)
+    env.define("int", new FunctionType([freshTypeVar()], INT));
+    env.define("float", new FunctionType([freshTypeVar()], FLOAT));
+    env.define("str", new FunctionType([freshTypeVar()], STRING));
+    env.define("bool", new FunctionType([freshTypeVar()], BOOL));
+
+    // I/O
+    env.define("print", new FunctionType([freshTypeVar()], UNIT));
+
+    // Terminal constructors
+    // Note: graph2d/graph3d are now in the 'plots' module (import graph2d from plots)
+    const terminalType = new RecordType(new Map(), true);
+    env.define("console_terminal", new FunctionType([terminalType], terminalType));
+  }
+
+  /**
+   * Check if a name is defined in the scope chain (not just via overload resolution)
+   */
+  isDefinedInScope(env, name) {
+    let current = env;
+    while (current) {
+      if (current.bindings.has(name)) {
+        return true;
+      }
+      current = current.parent;
+    }
+    return false;
+  }
+
+  /**
+   * Add an error with stage information
+   */
+  addError(message, location) {
+    this.errors.push({
+      stage: 'type_check',
+      type: 'type',
+      message,
+      location
+    });
+  }
+
+  /**
+   * Add a warning with stage information
+   */
+  addWarning(message, location) {
+    this.warnings.push({
+      stage: 'type_check',
+      type: 'warning',
+      message,
+      location
+    });
+  }
+
+  /**
+   * Set the inferred type for an AST node
+   */
+  setType(node, type) {
+    this.typeMap.set(node, type);
+    node.inferredType = type;
+    return type;
+  }
+
+  /**
+   * Get the inferred type for an AST node
+   */
+  getType(node) {
+    return this.typeMap.get(node);
+  }
+}
+
+// Mix in all method groups from extracted modules
+Object.assign(TypeChecker.prototype,
+  typeCompatibilityMethods,
+  bidirectionalInferenceMethods,
+  overloadResolutionMethods,
+  moduleAnalysisMethods
+);
+
+// Register all visitor methods
+registerVisitors(TypeChecker);
+
+// ============================================================================
+// EXPORTS
+// ============================================================================
+
+/**
+ * Type check an AST
+ */
+export function typeCheck(ast, options = {}) {
+  const checker = new TypeChecker(options);
+  return checker.check(ast);
+}
+
+export { TypeEnvironment };