plum

#![allow(non_snake_case)]
use plum_core::ast::*;
use plum_core::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();
    assert!(!tree.root_node().has_error(), "parse error:\n{}", tree.root_node().to_sexp());
    let ap = AstParser::new(src);
    ap.parseSource(tree.root_node())
}

fn onlyFn(source: &Source) -> &Fn {
    source.items.iter().find_map(|i| match i { Item::Fn(f) => Some(f), _ => None }).expect("expected a Fn item")
}

fn onlyTrait(source: &Source) -> &Trait {
    source.items.iter().find_map(|i| match i { Item::Trait(t) => Some(t), _ => None }).expect("expected a Trait item")
}

fn onlyEnum(source: &Source) -> &Enum {
    source.items.iter().find_map(|i| match i { Item::Enum(e) => Some(e), _ => None }).expect("expected an Enum item")
}

#[test]
fn nestedClassMethodsBecomeTopLevelFnItemsWithTypeParamSet() {
    let src = "\
type Cat =
  name: Str

  fun getName(self) -> Str =
    self.name

  fun setName(self, n: Str) =
    todo
";
    let source = parse(src);
    assert_eq!(source.items.len(), 3, "class + 2 nested methods, in order");
    let Item::Class(class) = &source.items[0] else { panic!("expected a Class item first") };
    assert_eq!(class.name, "Cat");
    let Item::Fn(get_name) = &source.items[1] else { panic!("expected getName immediately after the class") };
    assert_eq!(get_name.name, "getName");
    assert_eq!(get_name.type_param, Some("Cat".to_string()));
    let Item::Fn(set_name) = &source.items[2] else { panic!("expected setName immediately after getName") };
    assert_eq!(set_name.name, "setName");
    assert_eq!(set_name.type_param, Some("Cat".to_string()));
}

#[test]
fn nestedEnumMethodBecomesTopLevelFnItemWithTypeParamSet() {
    let src = "\
enum Step =
  | ReadMin
  | ReadMax

  fun toNumber(self) -> Int =
    match self
      ReadMin => 0
      ReadMax => 1
";
    let source = parse(src);
    assert_eq!(source.items.len(), 2, "enum + 1 nested method, in order");
    let Item::Enum(e) = &source.items[0] else { panic!("expected an Enum item first") };
    assert_eq!(e.name, "Step");
    let Item::Fn(to_number) = &source.items[1] else { panic!("expected toNumber immediately after the enum") };
    assert_eq!(to_number.name, "toNumber");
    assert_eq!(to_number.type_param, Some("Step".to_string()));
}

#[test]
fn enumDiscriminantValuesParseIntoParamsAndVariantValues() {
    let src = "\
enum Step(n: Int) =
  | ReadMin(0)
  | ReadMax(1)
";
    let source = parse(src);
    let e = onlyEnum(&source);
    assert_eq!(e.params, vec![EnumParam { name: "n".to_string(), ty: Type { name: "Int".to_string(), generics: vec![] } }]);
    assert_eq!(e.variants[0].fields, Vec::<String>::new());
    assert_eq!(e.variants[0].values, vec![Expr::Int(0)]);
    assert_eq!(e.variants[1].fields, Vec::<String>::new());
    assert_eq!(e.variants[1].values, vec![Expr::Int(1)]);
}

#[test]
fn classAndEnumWithNoNestedMethodsProduceNoExtraFnItems() {
    let src = "\
type Dog =
  name: Str

enum Bool =
  | True
  | False
";
    let source = parse(src);
    assert_eq!(source.items.len(), 2, "no nested methods means no extra Item::Fn entries");
    assert!(matches!(source.items[0], Item::Class(_)));
    assert!(matches!(source.items[1], Item::Enum(_)));
}

#[test]
fn closureLiteralParsesWithParamsAndBody() {
    let src = "\
fun useClosure() -> Bool =
  cb = |v|
    True
  cb(5)
";
    let source = parse(src);
    let f = onlyFn(&source);
    let FnBody::Block(block) = &f.body else { panic!("expected a block body") };
    let Stmt::Assign(assign) = &block.stmts[0] else { panic!("expected an assign statement") };
    let Expr::Closure(closure) = &assign.values[0] else { panic!("expected a closure expression, got {:?}", assign.values[0]) };
    assert_eq!(closure.params, vec!["v".to_string()]);
    assert_eq!(closure.body.stmts.len(), 1);
}

#[test]
fn closureLiteralParsesWithNoParams() {
    let src = "\
fun useClosure() -> Bool =
  cb = ||
    True
  cb()
";
    let source = parse(src);
    let f = onlyFn(&source);
    let FnBody::Block(block) = &f.body else { panic!("expected a block body") };
    let Stmt::Assign(assign) = &block.stmts[0] else { panic!("expected an assign statement") };
    let Expr::Closure(closure) = &assign.values[0] else { panic!("expected a closure expression") };
    assert!(closure.params.is_empty());
}

#[test]
fn fnWithoutReceiverButWithReturnTypeHasNoTypeParam() {
    let src = "fun main() -> Int =\n  0\n";
    let source = parse(src);
    let f = onlyFn(&source);
    assert_eq!(f.type_param, None, "return type must not be mistaken for a method receiver");
    assert_eq!(f.returns.as_ref().map(|t| t.name.clone()), Some("Int".to_string()));
}

#[test]
fn nestedMethodWithReturnTypeHasCorrectTypeParam() {
    let src = "type Cat =\n  name: Str\n\n  fun toStr() -> Str =\n    \"x\"\n";
    let source = parse(src);
    let f = onlyFn(&source);
    assert_eq!(f.type_param, Some("Cat".to_string()));
    assert_eq!(f.returns.as_ref().map(|t| t.name.clone()), Some("Str".to_string()));
}

#[test]
fn fnValueTypeParamParsesWithPositionalTypesAndReturn() {
    let src = "fun each(cb: fn(Int) -> Bool) -> Bool =\n  True\n";
    let source = parse(src);
    let f = onlyFn(&source);
    let ParamType::Fn(param_types, ret) = &f.params[0].ty else { panic!("expected ParamType::Fn, got {:?}", f.params[0].ty) };
    assert_eq!(param_types.len(), 1);
    assert_eq!(param_types[0].name, "Int");
    assert_eq!(ret.as_ref().map(|t| t.name.clone()), Some("Bool".to_string()));
}

#[test]
fn traitMethodReturnTypeIsNotMangledToArrowToken() {
    let src = "trait Show =\n  show() -> Str\n";
    let source = parse(src);
    let t = onlyTrait(&source);
    let method = &t.methods[0];
    assert_eq!(method.name, "show");
    assert_eq!(
        method.returns.as_ref().map(|r| r.name.clone()),
        Some("Str".to_string()),
        "trait method return type must be the real type, not the '->' token"
    );
}

#[test]
fn fnValueTypeParamParsesWithNoReturn() {
    let src = "fun each(cb: fn(Int)) -> Bool =\n  True\n";
    let source = parse(src);
    let f = onlyFn(&source);
    let ParamType::Fn(param_types, ret) = &f.params[0].ty else { panic!("expected ParamType::Fn") };
    assert_eq!(param_types.len(), 1);
    assert!(ret.is_none());
}