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());
}