plum

#![allow(non_snake_case)]

use plum_checker::types::*;

use plum_checker::{checkSource, plumTypeFromAst, unify};

use plum_core::ast::Type as AstType;

use plum_core::{ast::*, AstParser};

fn parse(src: &str) -> Source {

    let mut parser = tree_sitter::Parser::new();

    parser.set_language(&tree_sitter_plum::LANGUAGE.into()).unwrap();

    let tree = parser.parse(src, None).unwrap();

    let ap = AstParser::new(src);

    ap.parseSource(tree.root_node())

}

#[test]

fn astTypeIntMapsToTint() {

    let ast_ty = AstType { name: "Int".to_string(), generics: vec![] };

    assert_eq!(plumTypeFromAst(&ast_ty), PlumType::TInt);

}

#[test]

fn astTypeUnknownMapsToNamed() {

    let ast_ty = AstType { name: "MyClass".to_string(), generics: vec![] };

    assert_eq!(plumTypeFromAst(&ast_ty), PlumType::TNamed("MyClass".to_string()));

}

#[test]

fn unifySameTypesOk() {

    assert!(unify(&PlumType::TInt, &PlumType::TInt).is_ok());

    assert!(unify(&PlumType::TFloat, &PlumType::TFloat).is_ok());

}

#[test]

fn unifyDifferentTypesErr() {

    assert!(unify(&PlumType::TInt, &PlumType::TFloat).is_err());

}

#[test]

fn freshVarsAreUnique() {

    let mut state = InferState::new();

    let a = state.freshVar();

    let b = state.freshVar();

    assert_ne!(a, b);

    assert_eq!(a, "a0");

    assert_eq!(b, "a1");

}

#[test]

fn monoScheme() {

    let scheme = TypeScheme::mono(PlumType::TInt);

    assert!(scheme.vars.is_empty());

    assert_eq!(*scheme.body, PlumType::TInt);

}

#[test]

fn validAddFnPasses() {

    let src = "fun add(a: Int, b: Int) -> Int =\n  a + b\n";

    let source = parse(src);

    assert!(checkSource(&source).is_ok(), "expected Ok");

}

#[test]

fn wrongReturnTypeIsError() {

    let src = "fun bad() -> Int =\n  True\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

    let errs = result.unwrap_err();

    assert!(errs[0].message.contains("return type mismatch"), "got: {}", errs[0].message);

}

#[test]

fn undeclaredVarIsError() {

    let src = "fun bad() -> Int =\n  x\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn typeMismatchInBinaryOpIsError() {

    let src = "fun bad() -> Int =\n  1 + 2.0\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn boolLiteralTrueFalseAreBool() {

    let src = "fun isTrue() -> Bool =\n  True\n";

    let source = parse(src);

    assert!(checkSource(&source).is_ok(), "expected Ok, got {:?}", checkSource(&source).err());

}

#[test]

fn boolLiteralWrongReturnTypeIsError() {

    let src = "fun bad() -> Int =\n  False\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn methodSelfFieldAccessPasses() {

    let src = "type Cat =\n  name: Str\n  age: Int\n\n  fun getName() -> Str =\n    self.name\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn methodSelfUnknownFieldIsError() {

    let src = "type Cat =\n  name: Str\n\n  fun getAge() -> Int =\n    self.age\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn nestedMethodTypeChecks() {

    let src = "type Cat =\n  name: Str\n  age: Int\n\n  fun getName(self) -> Str =\n    self.name\n";

    let result = checkSource(&parse(src));

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn nestedMethodUnknownFieldIsError() {

    let src = "type Cat =\n  name: Str\n\n  fun getAge(self) -> Int =\n    self.age\n";

    assert!(checkSource(&parse(src)).is_err());

}

#[test]

fn selfOutsideMethodIsError() {

    let src = "fun bad() -> Int =\n  self\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn classCallChecksFieldTypes() {

    let src = "type Cat =\n  name: Str\n  age: Int\n\nfun makeCat() -> Cat =\n  Cat(name: \"x\", age: 1)\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn classCallWrongFieldTypeIsError() {

    let src = "type Cat =\n  name: Str\n  age: Int\n\nfun makeCat() -> Cat =\n  Cat(name: \"x\", age: \"y\")\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn classCallUnknownFieldIsError() {

    let src = "type Cat =\n  name: Str\n\nfun makeCat() -> Cat =\n  Cat(name: \"x\", age: 1)\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn methodCallViaAttributeTypeChecksArgs() {

    let src = "type Cat =\n  name: Str\n\n  fun rename(n: Str) -> Str =\n    n\n\nfun use(c: Cat) -> Str =\n  c.rename(\"x\")\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn matchBindsNamePatternToSubjectType() {

    let src = "fun main(a: Int) -> Int =\n  match a\n    x =>\n      x\n    _ =>\n      0\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn matchInlineCaseBodyTypeChecks() {

    // Case bodies can be a single inline expression, not just an indented block.

    let src = "fun main(a: Int) -> Int =\n  match a\n    1 => 10\n    _ => 0\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn matchTrueFalseAreVariantPatternsNotBindingsWithoutEnumDecl() {

    // True/False are built-in Bool variants — they must be recognized as tag

    // comparisons even when the source doesn't redeclare `enum Bool`, so a

    // later `_` wildcard arm remains reachable (each pattern binds/compares,

    // it doesn't just re-bind the subject under the name "True").

    let src = "fun pick(a: Bool) -> Int =\n  match a\n    True =>\n      1\n    False =>\n      0\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn matchIntPatternAgainstStrSubjectIsError() {

    let src = "fun main(a: Str) -> Int =\n  match a\n    1 =>\n      1\n    _ =>\n      0\n";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn bareEnumTagUnifiesWithOwningEnumType() {

    // Regression: a bare non-Bool tag like `None` used to type as `TNamed("None")`

    // (itself, not its enum), so comparing it against an `Option` value would wrongly

    // fail with a type mismatch.

    let src = "\

enum Option =

  | Some[Int]

  | None

fun isNone(o: Option) -> Bool = o == None

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn variantConstructionChecksArgCountAndTypes() {

    let src = "\

enum Option =

  | Some[Int]

  | None

fun makeSome(v: Int) -> Option =

  Some(v)

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn variantConstructionWrongArgTypeIsError() {

    let src = "\

enum Option =

  | Some(Int)

  | None

fun bad() -> Option =

  Some(\"x\")

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn variantConstructionWrongArgCountIsError() {

    let src = "\

enum Shape =

  | Rect[Float, Float]

  | Circle[Float]

fun bad() -> Shape =

  Rect(1.0)

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn enumDiscriminantValuesTypeCheckWithNoErrors() {

    let src = "\

enum Step(n: Int) =

  | ReadMin(0)

  | ReadMax(1)

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn enumDiscriminantWrongValueCountIsError() {

    let src = "\

enum Step(n: Int) =

  | ReadMin(0)

  | ReadMax

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn enumDiscriminantWronglyTypedValueIsError() {

    let src = "\

enum Step(n: Int) =

  | ReadMin(0)

  | ReadMax(\"x\")

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn enumDiscriminantValueOnParamLessEnumIsError() {

    let src = "\

enum Option =

  | Some(5)

  | None

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn fieldAccessOnDiscriminantEnumReceiverTypeChecks() {

    let src = "\

enum Step(n: Int) =

  | ReadMin(0)

  | ReadMax(1)

  fun toNumber(self) -> Int =

    self.n

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn fieldAccessOnOrdinaryEnumReceiverIsUnaffectedByEnumParams() {

    // An ordinary enum (no discriminant `params`) isn't in `ctx.classes` OR

    // `ctx.enum_params`, so it falls through to the same permissive "unmodeled

    // type" escape hatch every other type not in `ctx.classes` gets (codegen,

    // not the checker, is what would catch a genuinely bad field access here) —

    // exactly as it did before discriminant enums existed.

    let src = "\

enum Option =

  | Some[Int]

  | None

  fun bad(self) -> Int =

    self.n

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn constructorPatternBindsFieldsToDeclaredTypes() {

    let src = "\

enum Shape =

  | Rect[Float, Float]

  | Circle[Float]

fun area(s: Shape) -> Float =

  match s

    Rect(w, h) =>

      w * h

    Circle(r) =>

      r * r

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn constructorPatternWrongFieldCountIsError() {

    let src = "\

enum Shape =

  | Rect[Float, Float]

  | Circle[Float]

fun bad(s: Shape) -> Float =

  match s

    Rect(w) =>

      w

    _ =>

      0.0

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn classNameCollidingWithEnumVariantIsAClearError() {

    // `Cat(...)` is ambiguous when `Cat` is both a class and an enum variant:

    // downstream code consults `enum_variants` first, so the class

    // constructor would otherwise be silently shadowed with no diagnostic.

    let src = "\

type Cat =

  name: Str

enum Animal =

  | Cat

  | Dog

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err(), "expected Err");

    let errs = result.unwrap_err();

    assert!(

        errs.iter().any(|e| e.message.contains("is declared as both a class and an enum variant")),

        "got: {:?}",

        errs

    );

}

#[test]

fn genericClassInstantiatedAtTwoConcreteTypesTypeChecks() {

    let src = "\

type Box[T] =

  value: T

fun makeIntBox() -> Box =

  Box(value: 5)

fun makeStrBox() -> Box =

  Box(value: \"x\")

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn genericFunctionCalledWithDifferentConcreteTypesPerSiteTypeChecks() {

    let src = "\

fun wrap(value: T) -> Bool =

  True

fun useInt() -> Bool =

  wrap(5)

fun useStr() -> Bool =

  wrap(\"x\")

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn genericFunctionWithTwoIndependentTypeParamsTypeChecks() {

    let src = "\

fun pair(first: T, second: U) -> Bool =

  True

fun use() -> Bool =

  pair(1, \"x\")

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn genericFunctionWithInternallyInconsistentBodyIsRejectedAfterSpecialization() {

    // A generic function whose body is genuinely inconsistent with its concrete,

    // correct declared return type must still be REJECTED after specialization.

    // The `-> Str` on `useIt` means the buggy (unconditional return-overwrite)

    // behavior — rewriting `wrong$Str`'s `-> Int` to `-> Str` — would make the

    // whole program type-check, masking the real `expected Int, found Str` error.

    let src = "\

fun wrong(x: T) -> Int =

  \"hello\"

fun useIt() -> Str =

  wrong(\"s\")

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(

        result.is_err(),

        "expected the internally-inconsistent generic function to be rejected, got Ok"

    );

}

#[test]

fn genericEnumSingleInstantiationTypeChecks() {

    // A generic Option-shaped enum, constructed at one concrete type (`Some(5)`),

    // matched, must resolve end-to-end via checkSource.

    let src = "\

enum Option =

  | Some[T]

  | None

fun get() -> Int =

  o = Some(5)

  match o

    Some(v) =>

      v

    None =>

      0

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

    // Directly prove resolution happened: the monomorphized output must contain a

    // concrete `Option$Int` enum whose `Some$Int` variant carries an `Int` field

    // (not the generic `a`), and must NOT retain the generic `Option` template.

    // The variant name is ALSO mangled (`Some` -> `Some$Int`), the same suffix as

    // the enum's own name.

    let mono = plum_checker::monomorphize::monomorphizeSource(&source)

        .expect("monomorphize should succeed");

    let opt = mono.items.iter().find_map(|it| match it {

        Item::Enum(e) if e.name == "Option$Int" => Some(e),

        _ => None,

    });

    let opt = opt.expect("expected a specialized `Option$Int` enum in the output");

    let some = opt.variants.iter().find(|v| v.name == "Some$Int")

        .expect("expected `Some$Int` (mangled) variant on `Option$Int`");

    assert_eq!(some.fields, vec!["Int".to_string()], "Some's field should be concrete Int");

    assert!(

        !mono.items.iter().any(|it| matches!(it, Item::Enum(e) if e.name == "Option")),

        "the generic `Option` template must be dropped from the output"

    );

}

#[test]

fn genericEnumMultipleInstantiationsCoexistAndTypeCheck() {

    // The SAME generic enum instantiated at two different concrete types in one

    // program must now type-check correctly for BOTH instantiations — this is the

    // behavior this task adds (previously this was a documented, rejected limitation).

    let src = "\

enum Option =

  | Some[T]

  | None

fun useInt() -> Int =

  o = Some(5)

  match o

    Some(v) =>

      v

    None =>

      0

fun useStr() -> Str =

  o = Some(\"x\")

  match o

    Some(v) =>

      v

    None =>

      \"z\"

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

    // Directly prove both specializations exist independently, with distinct

    // mangled variant names, so neither collides with the other.

    let mono = plum_checker::monomorphize::monomorphizeSource(&source)

        .expect("monomorphize should succeed");

    let has_enum_with_variant = |enum_name: &str, variant_name: &str| {

        mono.items.iter().any(|it| matches!(it, Item::Enum(e) if e.name == enum_name

            && e.variants.iter().any(|v| v.name == variant_name)))

    };

    assert!(has_enum_with_variant("Option$Int", "Some$Int"), "expected Option$Int with Some$Int");

    assert!(has_enum_with_variant("Option$Str", "Some$Str"), "expected Option$Str with Some$Str");

}

#[test]

fn genericMethodOnGenericClassTypeChecks() {

    let src = "\

type Box[T] =

  value: T

  fun getValue() -> T =

    self.value

fun use() -> Int =

  b = Box(value: 5)

  b.getValue()

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

#[ignore = "slow (~100s) — exercises the 10,000-specialization runaway-recursion guard; run explicitly with --ignored when touching monomorphize.rs's guard logic"]

fn unboundedRecursiveGenericInstantiationIsAClearError() {

    // `recurse` re-wraps its argument in a `Box` on every recursive call, so each

    // specialization's own body demands a specialization of `recurse` at a STRICTLY

    // bigger type (`recurse$Int`, then `recurse$Box$Int`, then `recurse$Box$Box$Int`,

    // ...), forever. This genuinely grows the worklist without bound (unlike a

    // generic class field merely NAMING a recursive generic type in its own

    // declaration, which is never itself a call site and so never reaches the

    // worklist at all) and must fail with a clear, bounded error rather than hang.

    let src = "\

type Box[T] =

  value: T

fun recurse(v: T) -> Int =

  b = Box(value: v)

  recurse(b)

fun use() -> Int =

  recurse(5)

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

    let errs = result.unwrap_err();

    assert!(errs[0].message.contains("monomorphize"), "got: {:?}", errs);

}

#[test]

fn ordinaryFunctionWithBareGenericEnumParamTypeChecks() {

    // The shape that broke the pre-existing codegen test: an otherwise-ordinary

    // function taking a bare generic-enum-typed parameter.

    let src = "\

enum Option =

  | Some[T]

  | None

fun unwrapOr(o: Option, default: Int) -> Int =

  match o

    Some(v) =>

      v

    None =>

      default

fun use() -> Int =

  unwrapOr(Some(5), 0)

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

    // Directly prove `unwrapOr` itself got specialized (not left bare/unresolved).

    let mono = plum_checker::monomorphize::monomorphizeSource(&source)

        .expect("monomorphize should succeed");

    let has_specialized_unwrap_or = mono.items.iter().any(|it| matches!(it, Item::Fn(f)

        if f.name.starts_with("unwrapOr$") && f.type_param.is_none()));

    assert!(has_specialized_unwrap_or, "expected a specialized `unwrapOr$...` function in the output");

}

#[test]

fn ordinaryFunctionWithBareGenericClassParamTypeChecks() {

    // The same shape, for a generic CLASS param instead of an enum — untested until

    // now, but the identical root cause: `Box` is dropped from the monomorphized

    // output, so a bare `Box`-typed param would otherwise reference nothing.

    let src = "\

type Box[T] =

  value: T

  fun getBoxValue() -> T =

    self.value

fun sumBox(b: Box) -> Int =

  b.getBoxValue()

fun use() -> Int =

  sumBox(Box(value: 5))

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn closureLiteralInfersAsAFunctionType() {

    let src = "\

fun useClosure() -> Bool =

  cb = |v|

    True

  cb(5)

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn fnValueTypedParamCanBeCalled() {

    let src = "\

fun each(cb: fn(Int) -> Bool) -> Bool =

  cb(5)

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn fieldAssignmentTargetWithMatchingTypePasses() {

    let src = "\

type Cat =

  name: Str

  age: Int

  fun haveBirthday() =

    self.age = self.age + 1

";

    let source = parse(src);

    assert!(checkSource(&source).is_ok(), "expected Ok");

}

#[test]

fn fieldAssignmentTargetWithMismatchedTypeIsError() {

    let src = "\

type Cat =

  name: Str

  age: Int

  fun breakCat() =

    self.age = \"oops\"

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn fieldAssignmentTargetUnknownFieldIsError() {

    let src = "\

type Cat =

  name: Str

  age: Int

  fun breakCat() =

    self.nope = 1

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_err());

}

#[test]

fn closurePassedToFnValueTypedParamTypeChecks() {

    let src = "\

fun each(cb: fn(Int) -> Bool) -> Bool =

  cb(5)

fun use() -> Bool =

  each(|v|

    True)

";

    let source = parse(src);

    let result = checkSource(&source);

    assert!(result.is_ok(), "expected Ok, got {:?}", result.err());

}

#[test]

fn variadicCallWithZeroTrailingArgsPasses() {

    let src = "\

fun sumAll(nums: ...Int) -> Int =

  0

fun useSumAll() -> Int =

  sumAll()

";

    let source = parse(src);

    assert!(checkSource(&source).is_ok(), "expected Ok");

}

#[test]

fn variadicCallWithSeveralTrailingArgsPasses() {

    let src = "\

fun sumAll(nums: ...Int) -> Int =

  0

fun useSumAll() -> Int =

  sumAll(1, 2, 3)

";

    let source = parse(src);

    assert!(checkSource(&source).is_ok(), "expected Ok");

}

#[test]

fn variadicCallWithMismatchedTrailingArgTypeIsError() {

    let src = "\

fun sumAll(nums: ...Int) -> Int =

  0

fun useSumAll() -> Int =

  sumAll(1, \"two\")

";

    let source = parse(src);

    assert!(checkSource(&source).is_err());

}

#[test]

fn variadicCallWithFixedPrefixPasses() {

    let src = "\

fun combine(prefix: Int, rest: ...Int) -> Int =

  prefix

fun useCombine() -> Int =

  combine(1, 2, 3)

";

    let source = parse(src);

    assert!(checkSource(&source).is_ok(), "expected Ok");

}

#[test]

fn twoVariadicParamsIsError() {

    let src = "\

fun bad(a: ...Int, b: ...Int) -> Int =

  0

";

    let source = parse(src);

    assert!(checkSource(&source).is_err());

}

#[test]

fn variadicParamNotLastIsError() {

    let src = "\

fun bad(a: ...Int, b: Int) -> Int =

  0

";

    let source = parse(src);

    assert!(checkSource(&source).is_err());

}

#[test]

fn forLoopOverVariadicBindsElementType() {

    let src = "\

fun sumAll(nums: ...Int) -> Int =

  total = 0

  for v := range nums

    total = total + v

  total

";

    let source = parse(src);

    assert!(checkSource(&source).is_ok(), "expected Ok");

}

#[test]

fn forLoopOverVariadicWithTwoVarsIsError() {

    let src = "\

fun bad(nums: ...Int) -> Int =

  for v, i := range nums

    v

  0

";

    let source = parse(src);

    assert!(checkSource(&source).is_err());

}