yukigo-haskell-parser 0.1.0

package/src/typechecker/checker.ts

@@ -0,0 +1,395 @@
+import {
+  AST,
+  Function,
+  UnguardedBody,
+  GuardedBody,
+  Visitor,
+  Return,
+} from "yukigo-ast";
+import { typeMappings } from "../utils/types.js";
+import { InferenceEngine, PatternVisitor } from "./inference.js";
+import { CoreHM } from "./core.js";
+import { inspect } from "util";
+import { DeclarationCollectorVisitor } from "./DeclarationCollector.js";
+
+export interface TypeVar {
+  type: "TypeVar";
+  id: number;
+  name?: string;
+  constraints: string[];
+}
+
+export interface TypeConstructor {
+  type: "TypeConstructor";
+  name: string;
+  args: Type[];
+}
+export interface FunctionType {
+  type: "TypeConstructor";
+  name: "->";
+  args: [Type, Type];
+}
+export interface ListType {
+  type: "TypeConstructor";
+  name: "List";
+  args: [Type];
+}
+export interface TupleType {
+  type: "TypeConstructor";
+  name: "Tuple";
+  args: Type[];
+}
+
+export interface TypeScheme {
+  type: "TypeScheme";
+  quantifiers: number[];
+  body: Type;
+  constraints: Map<number, string[]>;
+}
+
+export type Type = TypeVar | TypeConstructor;
+
+export type Environment = Map<string, TypeScheme>;
+
+export type Substitution = Map<number, Type>;
+
+export type Result<T> =
+  | { success: true; value: T }
+  | { success: false; error: string };
+
+export const booleanType: TypeConstructor = {
+  type: "TypeConstructor",
+  name: "YuBoolean",
+  args: [],
+};
+export const numberType: TypeConstructor = {
+  type: "TypeConstructor",
+  name: "YuNumber",
+  args: [],
+};
+export const stringType: TypeConstructor = {
+  type: "TypeConstructor",
+  name: "YuString",
+  args: [],
+};
+
+export class FunctionRegistrarVisitor implements Visitor<void> {
+  constructor(
+    private env: Environment,
+    private signatureMap: Map<string, TypeScheme>,
+    private coreHM: CoreHM
+  ) {}
+  visitFunction(node: Function): void {
+    const functionName = node.identifier.value;
+    let funcScheme = this.signatureMap.get(functionName);
+    if (!funcScheme) {
+      const funcTypeVar = this.coreHM.freshVar();
+      funcScheme = {
+        type: "TypeScheme",
+        quantifiers: [],
+        body: funcTypeVar,
+        constraints: new Map(),
+      };
+      this.env.set(functionName, funcScheme);
+    }
+    for (const equation of node.equations) {
+      if (isUnguardedBody(equation.body)) {
+        const { statements } = equation.body.sequence;
+        statements
+          .filter((stmt) => stmt instanceof Function)
+          .forEach((func) => {
+            this.env.set(func.identifier.value, funcScheme);
+          });
+      }
+    }
+  }
+}
+export class FunctionCheckerVisitor implements Visitor<void> {
+  constructor(
+    private environments: Environment[],
+    private signatureMap: Map<string, TypeScheme>,
+    private coreHM: CoreHM,
+    private errors: string[]
+  ) {}
+  visitFunction(node: Function): void {
+    const functionName = node.identifier.value;
+    let funcScheme = this.signatureMap.get(functionName);
+    // Handle function without signature
+    if (!funcScheme) {
+      const firstEq = node.equations[0];
+      this.environments.unshift(new Map());
+      const inferenceEngine = new InferenceEngine(
+        this.signatureMap,
+        this.coreHM,
+        this.environments
+      );
+      const paramTypes = firstEq.patterns.map(() => this.coreHM.freshVar());
+
+      paramTypes.forEach((type, i) => {
+        try {
+          new PatternVisitor(
+            this.coreHM,
+            this.signatureMap,
+            type,
+            this.environments,
+            inferenceEngine
+          ).visit(firstEq.patterns[i]);
+        } catch (error) {
+          this.errors.push(`Type error in '${functionName}': ${error.message}`);
+        }
+      });
+
+      let inferredBodyType: Type;
+      if (isUnguardedBody(firstEq.body)) {
+        const equationStatements = firstEq.body.sequence.statements;
+        equationStatements.forEach((stmt) =>
+          stmt.accept(
+            new FunctionCheckerVisitor(
+              this.environments,
+              this.signatureMap,
+              this.coreHM,
+              this.errors
+            )
+          )
+        );
+        const returnResult = equationStatements
+          .find((stmt) => stmt instanceof Return)
+          .accept(inferenceEngine);
+        if (returnResult.success === false) {
+          this.errors.push(
+            `Type error in '${functionName}': ${returnResult.error}`
+          );
+          return;
+        }
+        inferredBodyType = returnResult.value;
+      } else {
+        // Handle guarded body inference if necessary
+        inferredBodyType = this.coreHM.freshVar(); // Placeholder
+      }
+
+      const fullFuncType = paramTypes.reduceRight(
+        (acc, param) => functionType(param, acc),
+        inferredBodyType
+      );
+
+      // Generalize the inferred type to create a polymorphic type scheme
+      this.environments.shift();
+      funcScheme = this.coreHM.generalize(this.environments[0], fullFuncType);
+      this.signatureMap.set(functionName, funcScheme);
+      this.environments[0].set(functionName, funcScheme);
+    }
+
+    const expectedArity = getArity(this.coreHM.instantiate(funcScheme));
+    for (const [index, equation] of node.equations.entries()) {
+      if (equation.patterns.length !== expectedArity) {
+        this.errors.push(`Type error in '${functionName}': Arity mismatch in equation number ${index}`);
+        continue;
+      }
+      try {
+        const funcType = this.coreHM.instantiate(funcScheme);
+
+        const returnType = getReturnType(funcType);
+        const patternTypes = getArgumentTypes(funcType);
+
+        this.environments.unshift(new Map());
+        const inferenceEngine = new InferenceEngine(
+          this.signatureMap,
+          this.coreHM,
+          this.environments
+        );
+        patternTypes.forEach((argType, i) => {
+          try {
+            equation.patterns[i].accept(
+              new PatternVisitor(
+                this.coreHM,
+                this.signatureMap,
+                argType,
+                this.environments,
+                inferenceEngine
+              )
+            );
+          } catch (error) {
+            this.errors.push(
+              `Type error in '${functionName}': ${error.message}`
+            );
+          }
+        });
+        if (isUnguardedBody(equation.body)) {
+          const equationStatements = equation.body.sequence.statements;
+          equationStatements.forEach((stmt) =>
+            stmt.accept(
+              new FunctionCheckerVisitor(
+                this.environments,
+                this.signatureMap,
+                this.coreHM,
+                this.errors
+              )
+            )
+          );
+          const returnResult = equationStatements
+            .find((stmt) => stmt instanceof Return)
+            .accept(inferenceEngine);
+          if (returnResult.success === false) {
+            this.errors.push(
+              `Type error in '${functionName}': ${returnResult.error}`
+            );
+            return;
+          }
+
+          const sub = this.coreHM.unify(returnResult.value, returnType);
+          if (sub.success === false) throw Error(sub.error);
+        } else {
+          // Handles GuardedBody case
+          for (const guard of equation.body) {
+            // checks if condition expression in guard is a resolves to YuBoolean
+
+            const condition = guard.condition.accept(inferenceEngine);
+            if (condition.success === false) throw Error(condition.error);
+
+            const conditionSub = this.coreHM.unify(
+              condition.value,
+              booleanType
+            );
+            if (conditionSub.success === false) throw Error(conditionSub.error);
+
+            const bodyResult = guard.body.accept(inferenceEngine);
+            if (bodyResult.success === false) throw Error(bodyResult.error);
+
+            const sub = this.coreHM.unify(bodyResult.value, returnType);
+            if (sub.success === false) throw Error(sub.error);
+          }
+        }
+        this.environments.shift();
+      } catch (error: any) {
+        this.errors.push(`Type error in '${functionName}': ${error.message}`);
+      }
+    }
+  }
+}
+
+export class TypeChecker {
+  private signatureMap: Map<string, TypeScheme>;
+  private coreHM: CoreHM;
+  private errors: string[];
+
+  constructor() {
+    this.signatureMap = new Map<string, TypeScheme>();
+    this.coreHM = new CoreHM();
+    this.errors = [];
+  }
+  check(ast: AST): string[] {
+    const typeAliasMap = new Map<string, Type>();
+    const recordMap = new Map<string, Type>();
+
+    // Phase 1: Collect declarations
+    const collector = new DeclarationCollectorVisitor(
+      this.errors,
+      typeAliasMap,
+      recordMap,
+      this.signatureMap,
+      this.coreHM
+    );
+    for (const node of ast) {
+      try {
+        node.accept(collector);
+      } catch (error) {
+        this.errors.push(error);
+      }
+    }
+
+    if (this.errors.length > 0) return this.errors;
+
+    // Phase 2: Infer and check functions
+    this.inferencePass(ast);
+
+    return this.errors;
+  }
+  inferencePass(ast: AST): void {
+    // Step 1: Register all functions in env
+    const globalEnv = new Map<string, TypeScheme>(this.signatureMap);
+    const visitor1 = new FunctionRegistrarVisitor(
+      globalEnv,
+      this.signatureMap,
+      this.coreHM
+    );
+    for (const node of ast) {
+      node.accept(visitor1);
+    }
+    // Step 2: Infer and check each function
+    const visitor2 = new FunctionCheckerVisitor(
+      [globalEnv],
+      this.signatureMap,
+      this.coreHM,
+      this.errors
+    );
+    for (const node of ast) {
+      node.accept(visitor2);
+    }
+  }
+}
+
+export function showType(t: Type): string {
+  if (t.type === "TypeVar") return t.name ?? `t${t.id}`;
+
+  if (isFunctionType(t)) {
+    const a = showType(t.args[0]);
+    const b = showType(t.args[1]);
+    const aDisp = isFunctionType(t.args[0]) ? `(${a})` : a;
+    return `${aDisp} -> ${b}`;
+  }
+  if (isListType(t)) return `[${showType(t.args[0])}]`;
+  if (isTupleType(t)) return `(${t.args.map(showType.bind(this)).join(", ")})`;
+
+  return t.args.length
+    ? `${t.name} ${t.args.map(showType.bind(this)).join(" ")}`
+    : t.name;
+}
+
+export function getReturnType(type: Type): Type {
+  let t: Type = type;
+  while (isFunctionType(t)) t = t.args[1];
+  return t;
+}
+export function getArgumentTypes(type: Type): Type[] {
+  const args: Type[] = [];
+  let t: Type = type;
+  while (isFunctionType(t)) {
+    args.push(t.args[0]);
+    t = t.args[1];
+  }
+  return args;
+}
+export function getArity(type: Type): number {
+  return getArgumentTypes(type).length;
+}
+
+export function isUnguardedBody(
+  body: UnguardedBody | GuardedBody[]
+): body is UnguardedBody {
+  return !Array.isArray(body);
+}
+
+export function isFunctionType(t: Type): t is FunctionType {
+  return t.type === "TypeConstructor" && t.name === "->" && t.args.length === 2;
+}
+export function isListType(t: Type): t is ListType {
+  return t.name === "List";
+}
+export function isTupleType(t: Type): t is TupleType {
+  return t.name === "Tuple";
+}
+
+export function functionType(params: Type, returnType: Type): FunctionType {
+  return {
+    type: "TypeConstructor",
+    name: "->",
+    args: [params, returnType],
+  };
+}
+export function listType(elementsType: Type): ListType {
+  return {
+    type: "TypeConstructor",
+    name: "List",
+    args: [elementsType],
+  };
+}

package/src/typechecker/core.ts

@@ -0,0 +1,162 @@
+import { typeClasses } from "../utils/types.js";
+import {
+  Environment,
+  Result,
+  showType,
+  Substitution,
+  Type,
+  TypeConstructor,
+  TypeScheme,
+  TypeVar,
+} from "./checker.js";
+
+export class CoreHM {
+  private nextVarId = 0;
+
+  public freshVar(constraints: string[] = []): TypeVar {
+    return { type: "TypeVar", id: this.nextVarId++, constraints };
+  }
+
+  public freeTypeVars(t: Type): Map<number, string[]> {
+    const freeVars = new Map<number, string[]>();
+    const collect = (type: Type) => {
+      if (type.type === "TypeVar") {
+        freeVars.set(type.id, type.constraints);
+      } else if (type.type === "TypeConstructor") {
+        type.args.forEach(collect);
+      }
+    };
+    collect(t);
+    return freeVars;
+  }
+  public generalize(env: Environment, t: Type): TypeScheme {
+    const envFreeVars = new Set<number>();
+    for (const scheme of env.values()) {
+      const schemeFreeVars = this.freeTypeVars(scheme.body);
+      for (const id of schemeFreeVars.keys()) {
+        if (!scheme.quantifiers.includes(id)) {
+          envFreeVars.add(id);
+        }
+      }
+    }
+
+    const typeFreeVars = this.freeTypeVars(t);
+
+    const quantifiers: number[] = [];
+    const constraints = new Map<number, string[]>();
+
+    for (const [id, consts] of typeFreeVars.entries()) {
+      if (!envFreeVars.has(id)) {
+        quantifiers.push(id);
+        if (consts.length > 0) {
+          constraints.set(id, consts);
+        }
+      }
+    }
+
+    return { type: "TypeScheme", quantifiers, body: t, constraints };
+  }
+
+  public instantiate(scheme: TypeScheme): Type {
+    const substitutions = new Map<number, Type>();
+    scheme.quantifiers.forEach((id) => {
+      const constraints = scheme.constraints.get(id) || [];
+      substitutions.set(id, this.freshVar(constraints));
+    });
+    return this.applySubst(substitutions, scheme.body);
+  }
+
+  public applySubst(subst: Substitution, t: Type): TypeVar | TypeConstructor {
+    if (t.type === "TypeVar") {
+      const replacement = subst.get(t.id);
+      return replacement ? replacement : t;
+    } else if (t.type === "TypeConstructor") {
+      return {
+        type: "TypeConstructor",
+        name: t.name,
+        args: t.args.map((arg) => this.applySubst(subst, arg)),
+      };
+    }
+
+    throw new Error("Unexpected TypeScheme in applySubst");
+  }
+
+  public composeSubst(s1: Substitution, s2: Substitution): Substitution {
+    const result = new Map(s2);
+    for (const [id, type] of s1) {
+      result.set(id, this.applySubst(s2, type));
+    }
+    return result;
+  }
+
+  public unify(t1: Type, t2: Type): Result<Substitution> {
+    if (t1.type === "TypeVar") return this.unifyVar(t1, t2);
+    if (t2.type === "TypeVar") return this.unifyVar(t2, t1);
+    if (t1.type === "TypeConstructor" && t2.type === "TypeConstructor") {
+      if (t1.name !== t2.name || t1.args.length !== t2.args.length) {
+        return {
+          success: false,
+          error: `Cannot unify ${showType(t1)} with ${showType(t2)}`,
+        };
+      }
+      let sub: Substitution = new Map();
+      for (let i = 0; i < t1.args.length; i++) {
+        const arg1 = this.applySubst(sub, t1.args[i]);
+        const arg2 = this.applySubst(sub, t2.args[i]);
+        const argSubRes = this.unify(arg1, arg2);
+        if (!argSubRes.success) return argSubRes;
+        sub = this.composeSubst(sub, argSubRes.value);
+      }
+      return { success: true, value: sub };
+    }
+    return { success: false, error: `Cannot unify non-types` };
+  }
+
+  public unifyVar(v: TypeVar, t: Type): Result<Substitution> {
+    if (t.type === "TypeVar" && t.id === v.id) {
+      return { success: true, value: new Map() };
+    }
+    if (this.occurs(v.id, t)) {
+      return { success: false, error: `Occurs check failed` };
+    }
+
+    for (const constraint of v.constraints) {
+      this.checkConstraint(constraint, t);
+    }
+    if (t.type === "TypeVar") {
+      for (const constraint of t.constraints) {
+        this.checkConstraint(constraint, v);
+      }
+      // Merge constraints
+      const mergedConstraints = [
+        ...new Set([...v.constraints, ...t.constraints]),
+      ];
+      v.constraints = mergedConstraints;
+      t.constraints = mergedConstraints;
+    }
+
+    return { success: true, value: new Map([[v.id, t]]) };
+  }
+
+  public checkConstraint(constraintName: string, t: Type) {
+    if (t.type === "TypeVar") {
+      if (!t.constraints.includes(constraintName)) {
+        t.constraints.push(constraintName);
+      }
+    } else if (t.type === "TypeConstructor") {
+      const instances = typeClasses.get(constraintName);
+      if (!instances || !instances.includes(t.name)) {
+        throw new Error(
+          `Type '${showType(t)}' is not an instance of '${constraintName}'`
+        );
+      }
+    }
+  }
+
+  public occurs(varId: number, t: Type): boolean {
+    if (t.type === "TypeVar") return t.id === varId;
+    if (t.type === "TypeConstructor")
+      return t.args.some((a) => this.occurs(varId, a));
+    return false;
+  }
+}