stone-lang 0.1.0

package/frontend/type-checker/OverloadedFunctionType.js

@@ -0,0 +1,43 @@
+/**
+ * OverloadedFunctionType - represents overloaded functions with multiple signatures
+ *
+ * Used for primitives that can be called with different argument types,
+ * and also for user-defined function overloads.
+ */
+
+/**
+ * Represents an overloaded function with multiple signatures.
+ * Used for primitives that can be called with different argument types,
+ * and now also for user-defined function overloads.
+ */
+export class OverloadedFunctionType {
+  constructor(name, overloads = [], userOverloads = []) {
+    this.name = name;
+    this.overloads = overloads; // Array of primitive functions with .stone signatures
+    this.userOverloads = userOverloads; // Array of { funcType, node, overloadId } for user-defined overloads
+  }
+
+  /**
+   * Add a user-defined function overload
+   * @param {FunctionType} funcType - The function's type signature
+   * @param {Object} node - The AST node of the function definition
+   */
+  addUserOverload(funcType, node) {
+    // Use already-set overloadId from node, or from funcType's stored value
+    const overloadId = node.overloadId || funcType._overloadId || `${this.name}_${this.userOverloads.length}`;
+    if (!node.overloadId) {
+      node.overloadId = overloadId; // Store on AST for executor registration
+    }
+    this.userOverloads.push({ funcType, node, overloadId });
+  }
+
+  toString() {
+    const primSigs = this.overloads
+      .filter(fn => fn.stone)
+      .map(fn => `(${fn.stone.in.join(", ")}) -> ${fn.stone.out}`);
+    const userSigs = this.userOverloads
+      .map(({ funcType }) => funcType.toString());
+    const allSigs = [...primSigs, ...userSigs].join(" | ");
+    return `overloaded[${this.name}]: ${allSigs}`;
+  }
+}

package/frontend/type-checker/TypeEnvironment.js

@@ -0,0 +1,165 @@
+/**
+ * Type Environment - tracks type bindings in scope
+ *
+ * Manages variable bindings, type aliases, declared types, extensions, and immutability
+ * with parent-child hierarchy for nested scopes.
+ */
+
+/**
+ * Extension info - tracks metadata for imported extension methods
+ */
+export class ExtensionInfo {
+  constructor(name, selfType, functionType, functionRef = null) {
+    this.name = name;           // Method name (e.g., "T", "shift")
+    this.selfType = selfType;   // TypeAnnotation for self parameter
+    this.functionType = functionType; // Full function type
+    this.functionRef = functionRef;   // Reference to compiled function (set later)
+  }
+}
+
+/**
+ * Type environment - tracks type bindings in scope
+ */
+export class TypeEnvironment {
+  constructor(parent = null) {
+    this.bindings = new Map();
+    this.typeAliases = new Map();
+    this.declaredTypes = new Map();  // Track type annotations for deferred resolution
+    this.immutableBindings = new Set();  // Track immutable bindings (cannot be rebound)
+    this.extensions = new Map();  // Track imported extension methods: name -> ExtensionInfo
+    this.parent = parent;
+  }
+
+  /**
+   * Define a new binding
+   */
+  define(name, type) {
+    this.bindings.set(name, type);
+  }
+
+  /**
+   * Look up a binding
+   */
+  lookup(name) {
+    if (this.bindings.has(name)) {
+      return this.bindings.get(name);
+    }
+    if (this.parent) {
+      return this.parent.lookup(name);
+    }
+    return null;
+  }
+
+  /**
+   * Define a type alias
+   */
+  defineTypeAlias(name, type) {
+    this.typeAliases.set(name, type);
+  }
+
+  /**
+   * Look up a type alias
+   */
+  lookupTypeAlias(name) {
+    if (this.typeAliases.has(name)) {
+      return this.typeAliases.get(name);
+    }
+    if (this.parent) {
+      return this.parent.lookupTypeAlias(name);
+    }
+    return null;
+  }
+
+  /**
+   * Define a declared type (from type annotation)
+   * Used for deferred type resolution of empty arrays
+   */
+  defineDeclared(name, type) {
+    this.declaredTypes.set(name, type);
+  }
+
+  /**
+   * Look up a declared type
+   * For primed variables (e.g., arr'), also checks the base name (arr)
+   */
+  lookupDeclared(name) {
+    if (this.declaredTypes.has(name)) {
+      return this.declaredTypes.get(name);
+    }
+    // For primed variables, check base name (arr' → arr)
+    if (name.endsWith("'")) {
+      const baseName = name.slice(0, -1);
+      if (this.declaredTypes.has(baseName)) {
+        return this.declaredTypes.get(baseName);
+      }
+    }
+    if (this.parent) {
+      return this.parent.lookupDeclared(name);
+    }
+    return null;
+  }
+
+  /**
+   * Check if a binding exists
+   */
+  has(name) {
+    return this.bindings.has(name) || (this.parent && this.parent.has(name));
+  }
+
+  /**
+   * Check if a binding exists in the current scope only (not parent)
+   */
+  hasInCurrentScope(name) {
+    return this.bindings.has(name);
+  }
+
+  /**
+   * Mark a binding as immutable (cannot be rebound)
+   */
+  markImmutable(name) {
+    this.immutableBindings.add(name);
+  }
+
+  /**
+   * Check if a binding is immutable in the scope chain
+   */
+  isImmutable(name) {
+    if (this.immutableBindings.has(name)) return true;
+    if (this.parent) return this.parent.isImmutable(name);
+    return false;
+  }
+
+  /**
+   * Create a child environment
+   */
+  createChild() {
+    return new TypeEnvironment(this);
+  }
+
+  /**
+   * Define an extension method
+   */
+  defineExtension(name, extensionInfo) {
+    this.extensions.set(name, extensionInfo);
+  }
+
+  /**
+   * Look up an extension method by name
+   */
+  lookupExtension(name) {
+    if (this.extensions.has(name)) {
+      return this.extensions.get(name);
+    }
+    if (this.parent) {
+      return this.parent.lookupExtension(name);
+    }
+    return null;
+  }
+
+  /**
+   * Check if an extension exists in the current scope only (not parent)
+   */
+  hasExtensionInCurrentScope(name) {
+    return this.extensions.has(name);
+  }
+}

package/frontend/type-checker/bidirectionalInference.js

@@ -0,0 +1,149 @@
+/**
+ * Bidirectional Inference Methods - mixin for TypeChecker
+ *
+ * Methods for collecting return statements, unifying return types,
+ * and back-propagating types to incomplete expressions.
+ */
+
+import {
+  TypeVariable,
+  UNIT,
+  unify,
+  formatType,
+} from "../types/types.js";
+
+/**
+ * Bidirectional inference methods to be mixed into TypeChecker.prototype
+ */
+export const bidirectionalInferenceMethods = {
+  /**
+   * Recursively collect all return statements from a function body
+   * Uses types already computed during the first pass (stored on node.inferredType)
+   * Returns array of { node, type, isIncomplete }
+   */
+  collectReturnStatements(body) {
+    const returns = [];
+
+    const collect = (statements) => {
+      for (const stmt of statements) {
+        if (stmt.type === "ReturnStatement") {
+          // Use the type already computed during the first pass
+          // Do NOT re-visit, as that may give incorrect results
+          const valueType = stmt.inferredType || UNIT;
+          const isIncomplete = this.isIncompleteType(valueType);
+          returns.push({ node: stmt, type: valueType, isIncomplete });
+        } else if (stmt.type === "BranchExpression") {
+          // Recurse into branches
+          for (const path of stmt.paths) {
+            collect(path.body);
+          }
+        } else if (stmt.type === "LoopExpression") {
+          collect(stmt.body);
+        } else if (stmt.type === "BlockExpression") {
+          collect(stmt.statements || []);
+        } else if (stmt.type === "BindingStructure") {
+          // Handle block expressions like { if (...) { return x } else { return y } }
+          for (const binding of stmt.bindings || []) {
+            if (binding.value) {
+              collect([binding.value]);
+            }
+          }
+          if (stmt.trailingExpr) {
+            collect([stmt.trailingExpr]);
+          }
+        }
+        // Continue visiting other statements (they might contain nested returns)
+      }
+    };
+
+    collect(body);
+    return returns;
+  },
+
+  /**
+   * Unify all return types, finding common type or reporting errors
+   * Returns { unifiedType, errors }
+   */
+  unifyReturnTypes(returnInfos) {
+    if (returnInfos.length === 0) {
+      return { unifiedType: UNIT, errors: [] };
+    }
+
+    const errors = [];
+
+    // Separate complete and incomplete types
+    // Exclude guard passthrough returns - these return the input unchanged
+    // when it's already the correct type, so they shouldn't constrain the return type
+    const complete = returnInfos.filter(r => !r.isIncomplete && !r.isGuardPassthrough);
+    const incomplete = returnInfos.filter(r => r.isIncomplete && !r.isGuardPassthrough);
+
+    // If no returns at all
+    if (complete.length === 0 && incomplete.length === 0) {
+      return { unifiedType: UNIT, errors };
+    }
+
+    // If all returns are incomplete, try to unify them
+    // This allows polymorphic functions like reverse(arr) that return types dependent on parameters
+    if (complete.length === 0) {
+      // Try to unify all incomplete types - if they conflict, error
+      let unifiedIncomplete = incomplete[0].type;
+      for (let i = 1; i < incomplete.length; i++) {
+        const result = unify(unifiedIncomplete, incomplete[i].type);
+        if (!result.success) {
+          errors.push({
+            message: `Return type mismatch: ${formatType(this.toStructuredType(unifiedIncomplete))} vs ${formatType(this.toStructuredType(incomplete[i].type))}`,
+            location: incomplete[i].node.location
+          });
+        } else if (result.unifiedType) {
+          // Use the unified type from the result if provided (e.g., LiteralUnionType)
+          unifiedIncomplete = result.unifiedType;
+        } else if (unifiedIncomplete instanceof TypeVariable) {
+          unifiedIncomplete = unifiedIncomplete.resolve();
+        }
+      }
+      return { unifiedType: unifiedIncomplete, errors };
+    }
+
+    // Unify all complete types
+    let unifiedType = complete[0].type;
+    for (let i = 1; i < complete.length; i++) {
+      const result = unify(unifiedType, complete[i].type);
+      if (!result.success) {
+        errors.push({
+          message: `Return type mismatch: ${formatType(unifiedType)} vs ${formatType(complete[i].type)}`,
+          location: complete[i].node.location
+        });
+      } else {
+        // Use the unified type from the result if provided (e.g., LiteralUnionType from merging literals)
+        if (result.unifiedType) {
+          unifiedType = result.unifiedType;
+        } else if (unifiedType instanceof TypeVariable) {
+          // Resolve if it's a TypeVariable
+          unifiedType = unifiedType.resolve();
+        }
+      }
+    }
+
+    return { unifiedType, errors };
+  },
+
+  /**
+   * Back-propagate unified type to incomplete return expressions
+   */
+  resolveIncompleteTypes(returnInfos, unifiedType) {
+    for (const { node, type, isIncomplete, isGuardPassthrough } of returnInfos) {
+      // Skip guard passthrough returns - they return input unchanged
+      // and shouldn't be constrained to the converted output type
+      if (isIncomplete && !isGuardPassthrough) {
+        // Unify the incomplete type with the resolved type
+        unify(type, unifiedType);
+
+        // Update the AST node's inferred type
+        node.inferredType = unifiedType;
+        if (node.value) {
+          node.value.inferredType = unifiedType;
+        }
+      }
+    }
+  },
+};

package/frontend/type-checker/index.js

@@ -0,0 +1,10 @@
+/**
+ * Frontend Type Checker - Barrel exports
+ */
+export { TypeChecker, OverloadedFunctionType } from './typeChecker.js';
+export { TypeEnvironment, ExtensionInfo } from './TypeEnvironment.js';
+export { typeCompatibilityMethods } from './typeCompatibility.js';
+export { bidirectionalInferenceMethods } from './bidirectionalInference.js';
+export { overloadResolutionMethods } from './overloadResolution.js';
+export { moduleAnalysisMethods } from './moduleAnalysis.js';
+export { BINARY_OP_NAMES, UNARY_OP_NAMES } from './operatorMappings.js';

package/frontend/type-checker/moduleAnalysis.js

@@ -0,0 +1,248 @@
+/**
+ * Module Analysis Methods - mixin for TypeChecker
+ *
+ * Methods for analyzing module imports and exports to determine types
+ * for static type checking.
+ */
+
+import {
+  FunctionType,
+  RecordType,
+  freshTypeVar,
+} from "../types/types.js";
+import { TypeEnvironment } from "./TypeEnvironment.js";
+import { OverloadedFunctionType } from "./OverloadedFunctionType.js";
+
+/**
+ * Module analysis methods to be mixed into TypeChecker.prototype
+ */
+export const moduleAnalysisMethods = {
+  /**
+   * Get export types for a module.
+   * Returns a Map<exportName, Type> for all exports in the module.
+   *
+   * @param {string} modulePath - The module path (e.g., "primitives", "math", "array")
+   * @returns {Map<string, Object>|null} Map of export names to their types
+   */
+  getModuleExportTypes(modulePath) {
+    // Check cache first
+    if (this.moduleTypeCache.has(modulePath)) {
+      return this.moduleTypeCache.get(modulePath);
+    }
+
+    // Handle 'primitives' module specially - these are JS-wrapped overloaded functions
+    if (modulePath === "primitives") {
+      const exports = this.analyzePrimitivesModule();
+      this.moduleTypeCache.set(modulePath, exports);
+      return exports;
+    }
+
+    // Handle 'plots' module - runtime terminal constructors
+    if (modulePath === "plots") {
+      const exports = this.analyzePlotsModule();
+      this.moduleTypeCache.set(modulePath, exports);
+      return exports;
+    }
+
+    // For .stn modules, we need the moduleResolver and parser
+    if (!this.moduleResolver || !this.parser) {
+      return null; // Can't analyze without these - fall back to type variables
+    }
+
+    // Check if it's a .stn builtin module
+    if (this.moduleResolver.isStnBuiltinModule(modulePath)) {
+      const exports = this.analyzeStnModule(modulePath);
+      if (exports) {
+        this.moduleTypeCache.set(modulePath, exports);
+      }
+      return exports;
+    }
+
+    // For file-based modules, we'd need async loading which isn't supported
+    // in the synchronous type checker. Return null to fall back to type variables.
+    return null;
+  },
+
+  /**
+   * Analyze the primitives module and extract types for all exports.
+   * Primitives are overloaded functions from the OVERLOADS registry.
+   */
+  analyzePrimitivesModule() {
+    const exports = new Map();
+
+    // Each overload set in OVERLOADS is an exportable primitive function
+    for (const [name, overloadSet] of Object.entries(this.overloads)) {
+      if (overloadSet && overloadSet.length > 0) {
+        // Create an OverloadedFunctionType for functions with multiple signatures
+        exports.set(name, new OverloadedFunctionType(name, overloadSet));
+      }
+    }
+
+    return exports;
+  },
+
+  /**
+   * Analyze the plots module and extract types for terminal constructors.
+   * Exports graph2d and graph3d functions.
+   *
+   * Terminal constructors take a config record and return a terminal handle.
+   * The terminal handle has dynamic methods (add, remove, etc.) that can't be
+   * statically typed, so we use type variables to allow any operations.
+   */
+  analyzePlotsModule() {
+    const exports = new Map();
+
+    // Use type variables for config and terminal handle since they're dynamically typed
+    // Config can be any record, terminal handle has runtime-defined methods
+    const configType = freshTypeVar();
+    const terminalHandleType = freshTypeVar();
+
+    // graph2d(config) -> TerminalHandle
+    exports.set("graph2d", new FunctionType([configType], terminalHandleType));
+    // graph3d(config) -> TerminalHandle
+    exports.set("graph3d", new FunctionType([freshTypeVar()], freshTypeVar()));
+
+    return exports;
+  },
+
+  /**
+   * Analyze a .stn module and extract export types.
+   * Recursively type-checks the module to determine export types.
+   *
+   * @param {string} modulePath - The module path
+   * @returns {Map<string, Object>|null} Map of export names to types
+   */
+  analyzeStnModule(modulePath) {
+    if (!this.moduleResolver || !this.parser) {
+      return null;
+    }
+
+    try {
+      const stnContent = this.moduleResolver.getStnBuiltinContent(modulePath);
+      if (!stnContent) {
+        return null;
+      }
+
+      // Parse the module
+      const ast = this.parser(stnContent, { filename: `<builtin:${modulePath}>` });
+      if (!ast || !ast.exports) {
+        return null;
+      }
+
+      // Create a fresh environment for the module
+      const moduleEnv = new TypeEnvironment();
+      this.registerBuiltins(moduleEnv);
+
+      // First pass: process imports to get types from other modules
+      for (const stmt of ast.statements) {
+        if (stmt.type === "ImportStatement") {
+          this.analyzeModuleImport(stmt, moduleEnv);
+        } else if (stmt.type === "NamespaceImportStatement") {
+          this.visit(stmt, moduleEnv);
+        }
+      }
+
+      // Second pass: process all statements to build environment
+      // This includes ExportStatements which are in both ast.statements and ast.exports
+      for (const stmt of ast.statements) {
+        if (stmt.type !== "ImportStatement" && stmt.type !== "NamespaceImportStatement") {
+          this.visit(stmt, moduleEnv);
+        }
+      }
+
+      // Extract export types and extension info
+      const exports = new Map();
+      const extensionExports = new Set(); // Track which exports are extensions
+
+      for (const exportStmt of ast.exports) {
+        if (exportStmt.type === "ExportStatement") {
+          // Direct export: export fn foo() or export x = 5
+          const decl = exportStmt.declaration;
+          if (decl.type === "FunctionDefinition") {
+            // For overloaded functions, the env already has the OverloadedFunctionType
+            // We only need to add once per name
+            if (!exports.has(decl.name)) {
+              const funcType = moduleEnv.lookup(decl.name);
+              if (funcType) {
+                exports.set(decl.name, funcType);
+              }
+              // Track if this is an extension function
+              if (decl.isExtension) {
+                extensionExports.add(decl.name);
+              }
+            }
+          } else if (decl.type === "Assignment") {
+            const varType = moduleEnv.lookup(decl.target);
+            if (varType) {
+              exports.set(decl.target, varType);
+            }
+          }
+        } else if (exportStmt.type === "SelectiveReExportStatement") {
+          // Re-export: export add, sub from primitives
+          const sourceExports = this.getModuleExportTypes(exportStmt.modulePath);
+          if (sourceExports) {
+            for (const spec of exportStmt.specifiers) {
+              const sourceType = sourceExports.get(spec.name);
+              if (sourceType) {
+                const exportName = spec.alias || spec.name;
+                exports.set(exportName, sourceType);
+                // Check if the source export is an extension
+                if (sourceExports._extensions && sourceExports._extensions.has(spec.name)) {
+                  extensionExports.add(exportName);
+                }
+              }
+            }
+          }
+        }
+        // NamespaceReExportStatement handled differently (exports namespace object)
+      }
+
+      // Attach extension info to the exports map
+      exports._extensions = extensionExports;
+      return exports;
+    } catch (error) {
+      // Module analysis failed - log but don't throw
+      this.addWarning(`Failed to analyze module '${modulePath}': ${error.message}`, null);
+      return null;
+    }
+  },
+
+  /**
+   * Process an import statement during module analysis.
+   * Registers imported names with their types from the source module.
+   * Reports errors for imports that don't exist in the module.
+   */
+  analyzeModuleImport(importStmt, env) {
+    // Register complex module overloads when any import from complex module is seen
+    // This allows operators like + - * / to work with complex numbers
+    if (importStmt.modulePath === 'complex') {
+      this._registerComplexOverloads();
+    }
+
+    const sourceExports = this.getModuleExportTypes(importStmt.modulePath);
+
+    for (const spec of importStmt.specifiers) {
+      // Support both spec.name (old format) and spec.importedName (new format)
+      const importedName = spec.importedName || spec.name;
+      const localName = spec.localName || importedName;
+
+      if (sourceExports) {
+        const sourceType = sourceExports.get(importedName);
+        if (sourceType) {
+          env.define(localName, sourceType);
+          continue;
+        }
+        // Module was analyzed but import name doesn't exist - report error
+        this.addError(
+          `'${importedName}' is not exported from module '${importStmt.modulePath}'`,
+          spec.location || importStmt.location
+        );
+        // Still define with type variable to allow continued type checking
+        env.define(localName, freshTypeVar());
+      } else {
+        // Module couldn't be analyzed - fall back to type variable silently
+        env.define(localName, freshTypeVar());
+      }
+    }
+  },
+};

package/frontend/type-checker/operatorMappings.js

@@ -0,0 +1,35 @@
+/**
+ * Operator Mappings - maps operators to overload function names
+ */
+
+/**
+ * Map binary operators to overload function names
+ */
+export const BINARY_OP_NAMES = {
+  "+": "__op_add",
+  "-": "__op_sub",
+  "*": "__op_mul",
+  "/": "__op_div",
+  "%": "__op_mod",
+  "^": "__op_pow",
+  ".+": "__op_elem_add",
+  ".-": "__op_elem_sub",
+  ".*": "__op_elem_mul",
+  "./": "__op_elem_div",
+  ".^": "__op_elem_pow",
+  "@": "__op_matmul",
+  "<": "__op_lt",
+  ">": "__op_gt",
+  "<=": "__op_lte",
+  ">=": "__op_gte",
+  "==": "__op_eq",
+  "!=": "__op_neq",
+};
+
+/**
+ * Map unary operators to overload function names
+ */
+export const UNARY_OP_NAMES = {
+  "-": "__op_neg",
+  "!": "__op_not",
+};