stone-lang 0.1.0

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

@@ -0,0 +1,483 @@
+/**
+ * Control Flow Visitors - mixin for TypeChecker
+ *
+ * Visitors for loops, branches, and control flow constructs.
+ */
+
+import {
+  FunctionType,
+  RecordType,
+  ArrayType,
+  TypeVariable,
+  NUM,
+  UNIT,
+  freshTypeVar,
+  typeEquals,
+  unify,
+  formatType,
+  matchType,
+  canPromote,
+  fromTypeAnnotation,
+} from "../../types/types.js";
+
+/**
+ * Control flow visitor methods to be mixed into TypeChecker.prototype
+ */
+export const controlFlowVisitors = {
+  /**
+   * Visit a loop expression
+   * Validates that state variable transitions preserve types
+   */
+  visitLoopExpression(node, env) {
+    const loopEnv = env.createChild();
+    const stateTypes = new Map();
+
+    // Helper: detect if an expression references a given identifier name
+    // Used to flag self-referential loop initializers (x = x) which evaluate to 'none'
+    const exprContainsIdentifier = (expr, name, parent = null, parentKey = null) => {
+      if (!expr || typeof expr !== "object") return false;
+      if (expr.type === "Identifier" && expr.name === name) {
+        // Ignore property positions like fwd.M (MemberExpression.property) where the identifier is a field name
+        const isMemberProp = (parent?.type === "MemberExpression" || parent?.type === "MemberAccess") &&
+          parent.property === expr &&
+          parent.computed !== true;
+        // Ignore identifier nodes that serve as object/record keys (common key property names)
+        const isObjectKey = typeof parentKey === "string" && /name|key|field/i.test(parentKey);
+        if (!isMemberProp && !isObjectKey) {
+          return true;
+        }
+      }
+      for (const key of Object.keys(expr)) {
+        const val = expr[key];
+        if (Array.isArray(val)) {
+          if (val.some(child => exprContainsIdentifier(child, name, expr, key))) return true;
+        } else if (val && typeof val === "object") {
+          if (exprContainsIdentifier(val, name, expr, key)) return true;
+        }
+      }
+      return false;
+    };
+
+    // Initialize state variables and their primed versions
+    for (const variable of node.variables) {
+      let varType;
+      // Disallow shadowing outer user-defined bindings (Stone has no shadowing)
+      // But allow loop state variables to have the same name as built-in functions
+      // (e.g., a loop state variable named 'sum' shouldn't conflict with the built-in sum function)
+      const existingBinding = env.lookup(variable.name);
+      if (existingBinding && !(existingBinding instanceof FunctionType)) {
+        this.addError(
+          `Loop variable '${variable.name}' shadows an outer binding. Stone disallows shadowing; use a distinct name.`,
+          variable.location || node.location
+        );
+      }
+      // Prevent self-referential loop initializers: h = h yields 'none' at runtime
+      const outerBinding = env.lookup(variable.name); // allow referencing an outer binding with same name
+      if (!outerBinding && variable.initialValue && exprContainsIdentifier(variable.initialValue, variable.name)) {
+        this.addError(
+          `Loop variable '${variable.name}' is initialized with itself. ` +
+          `Loop state initializers cannot reference their own name; compute the value in a separate binding (e.g., 'h_init') and use that instead.`,
+          variable.initialValue.location || node.location
+        );
+      }
+      if (variable.typeAnnotation) {
+        varType = fromTypeAnnotation(variable.typeAnnotation, env);
+        // If there's also an initial value, verify it matches the annotation
+        if (variable.initialValue) {
+          const initType = this.visit(variable.initialValue, env);
+          const unifyResult = unify(varType, initType);
+          if (!unifyResult.success) {
+            this.addError(
+              `Loop variable '${variable.name}' declared as ${varType}, but initialized with ${initType}`,
+              node.location
+            );
+          }
+        }
+      } else if (variable.initialValue) {
+        varType = this.visit(variable.initialValue, env);
+      } else {
+        varType = freshTypeVar();
+      }
+      stateTypes.set(variable.name, varType);
+      loopEnv.define(variable.name, varType);
+      // Define primed version with same type (x' must have same type as x)
+      loopEnv.define(variable.name + "'", varType);
+    }
+
+    // Process drivers - track for back-propagation of element types
+    const forDrivers = []; // [{iterableType, elemTypeVar, pattern}]
+    if (node.drivers) {
+      for (const driver of node.drivers) {
+        if (driver.type === "for") {
+          // For driver introduces iteration variable
+          const iterableType = this.visit(driver.iterable, loopEnv);
+          let elemType;
+          if (iterableType instanceof ArrayType) {
+            elemType = iterableType.rank === 1
+              ? iterableType.elementType
+              : new ArrayType(iterableType.elementType, iterableType.rank - 1);
+          } else {
+            // Unknown iterable type - use fresh TypeVariable and track for back-propagation
+            elemType = freshTypeVar();
+            if (iterableType instanceof TypeVariable) {
+              forDrivers.push({ iterableType, elemType, pattern: driver.pattern });
+            }
+          }
+          loopEnv.define(driver.pattern, elemType);
+        }
+        // While driver doesn't introduce variables
+      }
+    }
+
+    // Type check body with state type validation
+    // Track transition types for each state variable to ensure consistency
+    const transitionTypes = new Map(); // baseName -> [actualType, ...]
+    const assignedPrimes = new Set();  // Track which x' have been assigned
+
+    // Create a custom environment that tracks prime variable references
+    // We'll use this to detect references to x' before assignment
+    const primeRefEnv = loopEnv.createChild();
+
+    // Helper to check if an expression references unassigned primed variables
+    const checkPrimeReferences = (exprNode, location) => {
+      if (!exprNode) return;
+
+      if (exprNode.type === "Identifier" && exprNode.name.endsWith("'")) {
+        const baseName = exprNode.name.slice(0, -1);
+        if (stateTypes.has(baseName) && !assignedPrimes.has(baseName)) {
+          this.addError(
+            `Cannot reference '${exprNode.name}' before it is assigned in the loop body. ` +
+            `State transition '${exprNode.name}' must be assigned before it can be read.`,
+            exprNode.location || location
+          );
+        }
+      }
+
+      // Recursively check child nodes
+      if (exprNode.left) checkPrimeReferences(exprNode.left, location);
+      if (exprNode.right) checkPrimeReferences(exprNode.right, location);
+      if (exprNode.operand) checkPrimeReferences(exprNode.operand, location);
+      if (exprNode.args) exprNode.args.forEach(arg => checkPrimeReferences(arg, location));
+      if (exprNode.value) checkPrimeReferences(exprNode.value, location);
+      if (exprNode.object) checkPrimeReferences(exprNode.object, location);
+      if (exprNode.property && exprNode.computed) checkPrimeReferences(exprNode.property, location);
+      if (exprNode.condition) checkPrimeReferences(exprNode.condition, location);
+      if (exprNode.elements) exprNode.elements.forEach(el => checkPrimeReferences(el, location));
+      if (exprNode.bindings) exprNode.bindings.forEach(b => checkPrimeReferences(b.value, location));
+      if (exprNode.trailingExpr) checkPrimeReferences(exprNode.trailingExpr, location);
+      if (exprNode.paths) {
+        exprNode.paths.forEach(path => {
+          if (path.condition) checkPrimeReferences(path.condition, location);
+          if (path.body) path.body.forEach(stmt => checkPrimeReferences(stmt, location));
+        });
+      }
+    };
+
+    for (const stmt of node.body) {
+      // Check for state variable transitions (assignments to primed variables)
+      if (stmt.type === "Assignment" && stmt.target?.endsWith?.("'")) {
+        const baseName = stmt.target.slice(0, -1);
+        if (stateTypes.has(baseName)) {
+          // Check for multiple assignments to the same x'
+          if (assignedPrimes.has(baseName)) {
+            this.addError(
+              `State variable '${baseName}'' is assigned multiple times in the loop body. ` +
+              `Each state variable can only have one transition per iteration.`,
+              stmt.location
+            );
+          }
+
+          // Check for references to unassigned primed variables in the RHS
+          checkPrimeReferences(stmt.value, stmt.location);
+
+          const expectedType = stateTypes.get(baseName);
+          const actualType = this.visit(stmt.value, loopEnv);
+
+          // Mark this prime as assigned AFTER checking the RHS
+          assignedPrimes.add(baseName);
+
+          // Track all transition types for this variable
+          if (!transitionTypes.has(baseName)) {
+            transitionTypes.set(baseName, []);
+          }
+          transitionTypes.get(baseName).push({ type: actualType, location: stmt.location });
+
+          // Unify the actual type with the expected type
+          // This ensures type inference works for empty arrays and catches mismatches
+          const unifyResult = unify(expectedType, actualType);
+          if (!unifyResult.success) {
+            const expectedStr = formatType(this.toStructuredType(expectedType));
+            const actualStr = formatType(this.toStructuredType(actualType));
+            this.addError(
+              `Loop state type mismatch: '${baseName}' is ${expectedStr}, ` +
+              `transition '${baseName}'' is ${actualStr}`,
+              stmt.location
+            );
+          }
+          continue;
+        } else {
+          // Check if this is an outer loop's state variable
+          // If the primed variable is accessible but not in current loop's state,
+          // it must be from an enclosing loop
+          const primedVarType = loopEnv.lookup(stmt.target);
+          if (primedVarType !== null) {
+            this.addError(
+              `Cannot update outer loop state variable '${baseName}' from inner loop. ` +
+              `State transitions (using '${baseName}\\'') only work for the current loop's state variables. ` +
+              `To accumulate values across nested loops, collect results in the inner loop and merge in the outer loop.`,
+              stmt.location
+            );
+          }
+        }
+      }
+
+      // For non-prime assignments, also check for prime references
+      if (stmt.type === "Assignment" && stmt.value) {
+        checkPrimeReferences(stmt.value, stmt.location);
+      }
+
+      this.visit(stmt, loopEnv);
+    }
+
+    // Verify all transitions for each variable produce consistent types
+    for (const [baseName, transitions] of transitionTypes) {
+      if (transitions.length > 1) {
+        const firstType = transitions[0].type;
+        for (let i = 1; i < transitions.length; i++) {
+          const otherType = transitions[i].type;
+          // After unification, check if they resolved to the same type
+          const resolved1 = firstType instanceof TypeVariable ? firstType.resolve() : firstType;
+          const resolved2 = otherType instanceof TypeVariable ? otherType.resolve() : otherType;
+          if (!typeEquals(resolved1, resolved2)) {
+            // Types don't match - report error
+            const type1Str = formatType(this.toStructuredType(resolved1));
+            const type2Str = formatType(this.toStructuredType(resolved2));
+            this.addError(
+              `Inconsistent transition types for '${baseName}': ` +
+              `first assignment produces ${type1Str}, but another produces ${type2Str}`,
+              transitions[i].location
+            );
+          }
+        }
+      }
+    }
+
+    // Back-propagate element types to iterable TypeVariables
+    // If we learned the element type from usage, constrain the iterable
+    for (const { iterableType, elemType } of forDrivers) {
+      const resolvedElem = elemType instanceof TypeVariable ? elemType.resolve() : elemType;
+      if (resolvedElem !== elemType && !(resolvedElem instanceof TypeVariable)) {
+        // Element type was resolved - unify iterable with appropriate container type
+        // Default to 1D array since we don't know the rank
+        unify(iterableType, new ArrayType(resolvedElem, 1));
+      }
+    }
+
+    // Result is a record of state variables
+    const resultFields = new Map();
+    for (const [name, type] of stateTypes) {
+      resultFields.set(name, type);
+    }
+
+    return this.setType(node, new RecordType(resultFields, false));
+  },
+
+  /**
+   * Detect if a condition is a discriminant check pattern.
+   * Patterns: x.field == "literal" or x.field != "literal"
+   * Returns { varName, fieldName, literal, isNegated } or null.
+   */
+  extractDiscriminantCheck(condition) {
+    if (!condition) return null;
+
+    // Check for binary comparison (== or !=)
+    if (condition.type !== "BinaryOp") return null;
+    if (condition.operator !== "==" && condition.operator !== "!=") return null;
+
+    const left = condition.left;
+    const right = condition.right;
+
+    // Left side should be member access (x.field)
+    if (left?.type !== "MemberAccess" || left.computed) return null;
+    if (left.object?.type !== "Identifier") return null;
+
+    // Right side should be a string literal
+    if (right?.type !== "Literal" || typeof right.value !== "string") return null;
+
+    return {
+      varName: left.object.name,
+      fieldName: left.property?.name,
+      literal: right.value,
+      isNegated: condition.operator === "!="
+    };
+  },
+
+  /**
+   * Check if a branch body is a simple return of a variable.
+   * Returns the variable name if so, null otherwise.
+   */
+  getSimpleReturnVar(body) {
+    if (!body || body.length !== 1) return null;
+    const stmt = body[0];
+
+    // Check for return statement
+    if (stmt.type === "ReturnStatement" && stmt.value?.type === "Identifier") {
+      return stmt.value.name;
+    }
+
+    // Check for direct identifier (implicit return in single-statement branch)
+    if (stmt.type === "Identifier") {
+      return stmt.name;
+    }
+
+    return null;
+  },
+
+  /**
+   * Visit a branch expression (if/elif/else)
+   * Validates that all branches return compatible types
+   */
+  visitBranchExpression(node, env) {
+    const branchResults = [];
+    const branchLabels = [];
+    let hasElseBranch = false;
+
+    // Detect discriminant guard pattern in first branch
+    // Pattern: if (x.field != "literal") { return x }
+    // This is a "type guard passthrough" - the branch returns the input unchanged
+    // when it's already the correct type, so we shouldn't unify its type with other branches
+    let guardPassthroughBranchIndex = -1;
+
+    if (node.paths.length >= 2 && node.paths[0].condition) {
+      const discriminant = this.extractDiscriminantCheck(node.paths[0].condition);
+      if (discriminant && discriminant.isNegated) {
+        // Check if branch body just returns the discriminated variable
+        const returnVar = this.getSimpleReturnVar(node.paths[0].body);
+        if (returnVar === discriminant.varName) {
+          guardPassthroughBranchIndex = 0;
+        }
+      }
+    }
+
+    for (let i = 0; i < node.paths.length; i++) {
+      const path = node.paths[i];
+      if (path.condition) {
+        this.visit(path.condition, env);
+      } else {
+        // Path without condition is the else branch
+        hasElseBranch = true;
+      }
+
+      const branchEnv = env.createChild();
+      let lastType = UNIT;
+
+      for (const stmt of path.body) {
+        // Always visit the statement to ensure proper type tracking
+        lastType = this.visit(stmt, branchEnv);
+      }
+
+      branchResults.push({
+        path,
+        type: lastType,
+        isIncomplete: this.isIncompleteType(lastType),
+        index: i
+      });
+
+      // Label for error messages
+      if (i === 0) {
+        branchLabels.push("'if'");
+      } else if (path.condition) {
+        branchLabels.push(`'elif' #${i}`);
+      } else {
+        branchLabels.push("'else'");
+      }
+    }
+
+    // Check for value-producing conditionals without else branch
+    if (node._isValueProducing && !hasElseBranch) {
+      this.addError(
+        `Conditional expression used as a value must have an 'else' branch. ` +
+        `Add 'else { ... }' to provide a value for all cases.`,
+        node.location
+      );
+    }
+
+    // Separate complete and incomplete types
+    // Exclude guard passthrough branch from unification - it returns the input unchanged
+    // so its type shouldn't constrain other branches' types
+    const complete = branchResults.filter(b => !b.isIncomplete && b.index !== guardPassthroughBranchIndex);
+    const incomplete = branchResults.filter(b => b.isIncomplete && b.index !== guardPassthroughBranchIndex);
+
+    // If there's only one branch (no else), this is a guard clause
+    // The if-expression itself returns UNIT; any return statements inside
+    // are handled at the function level
+    if (branchResults.length === 1) {
+      return this.setType(node, UNIT);
+    }
+
+    // If no branches at all, return UNIT
+    if (branchResults.length === 0) {
+      return this.setType(node, UNIT);
+    }
+
+    // If all branches are incomplete (type variables), unify them and return
+    // This handles polymorphic functions where parameter types are unknown
+    if (complete.length === 0 && incomplete.length > 0) {
+      let resultType = incomplete[0].type;
+      for (let i = 1; i < incomplete.length; i++) {
+        const result = unify(resultType, incomplete[i].type);
+        if (!result.success) {
+          const branchIdx = branchResults.indexOf(incomplete[i]);
+          const firstIdx = branchResults.indexOf(incomplete[0]);
+          this.addError(
+            `Branch type mismatch: ${branchLabels[firstIdx]} and ${branchLabels[branchIdx]} have incompatible types`,
+            node.location
+          );
+        }
+        if (resultType instanceof TypeVariable) {
+          resultType = resultType.resolve();
+        }
+      }
+      return this.setType(node, resultType);
+    }
+
+    // Find common type from complete branches
+    let resultType = complete[0].type;
+    for (let i = 1; i < complete.length; i++) {
+      const result = unify(resultType, complete[i].type);
+      if (!result.success) {
+        // Find branch index for this complete branch
+        const branchIdx = branchResults.indexOf(complete[i]);
+        const firstIdx = branchResults.indexOf(complete[0]);
+        const firstStr = formatType(this.toStructuredType(resultType));
+        const otherStr = formatType(this.toStructuredType(complete[i].type));
+        this.addError(
+          `Branch type mismatch: ${branchLabels[firstIdx]} returns ${firstStr}, ` +
+          `${branchLabels[branchIdx]} returns ${otherStr}`,
+          node.location
+        );
+      } else {
+        // Resolve if it's a TypeVariable
+        if (resultType instanceof TypeVariable) {
+          resultType = resultType.resolve();
+        }
+        // When promotion happened, use the wider type (e.g., complex over num)
+        if (canPromote(resultType, complete[i].type).can) {
+          resultType = complete[i].type;
+        }
+      }
+    }
+
+    // Back-propagate to incomplete branches
+    for (const { path, type } of incomplete) {
+      unify(type, resultType);
+      // Update the last statement's inferredType in the path
+      if (path.body.length > 0) {
+        const lastStmt = path.body[path.body.length - 1];
+        lastStmt.inferredType = resultType;
+      }
+    }
+
+    return this.setType(node, resultType);
+  },
+};