open TypeInference

let tests_typeof =
  let x = Identifier.fresh () in
  let y = Identifier.fresh () in
  let z = Identifier.fresh () in
  [ (* IntConst *)
    "0", Term.IntConst 0, Some Type.Int
  ; (* int -> int -> int = <fun> *)
    ( "fun x -> fun y -> x + y"
    , Term.(Fun (x, Fun (y, Binop (Var x, Plus, Var y))))
    , Some Type.(Arrow (Int, Arrow (Int, Int))) )
  ; (* Not typed variable *)
    "x", Term.(Var "x"), None
  ; (* Binary operation *)
    "1 + 2", Term.(Binop (IntConst 1, Plus, IntConst 2)), Some Type.Int
  ; (* Pair *)
    "(1, 2)", Term.(Pair (IntConst 1, IntConst 2)), Some Type.(Product (Int, Int))
  ; (* Projection with first *)
    "fst (1, 2)", Term.(Proj (First, Pair (IntConst 1, IntConst 2))), Some Type.Int
  ; (* Projection with second *)
    "snd (1, 2)", Term.(Proj (Second, Pair (IntConst 1, IntConst 2))), Some Type.Int
  ; (* Apply (int) into (fun : int -> int) *)
    ( "(fun x -> x + 1) 5"
    , Term.(App (Fun (x, Binop (Var x, Plus, IntConst 1)), IntConst 5))
    , Some Type.Int )
  ; (* Apply product (int * int) into a not compatible function (fun : int -> int) *)
    ( "(fun x -> x + 1) (1, 2)"
    , Term.(App (Fun (x, Binop (Var x, Plus, IntConst 1)), Pair (IntConst 1, IntConst 2)))
    , None )
  ; (* int -> int -> int = <fun> *)
    ( "fun x -> fun y -> x * y"
    , Term.(Fun (x, Fun (y, Binop (Var x, Times, Var y))))
    , Some Type.(Arrow (Int, Arrow (Int, Int))) )
  ; (* int -> int -> int = <fun> *)
    ( "fun x -> fun y -> x - y"
    , Term.(Fun (x, Fun (y, Binop (Var x, Minus, Var y))))
    , Some Type.(Arrow (Int, Arrow (Int, Int))) )
  ; (* int -> int -> int = <fun> *)
    ( "fun x -> fun y -> y / x"
    , Term.(Fun (x, Fun (y, Binop (Var y, Div, Var x))))
    , Some Type.(Arrow (Int, Arrow (Int, Int))) )
  ; (* Use of a non declared variable *)
    "fun x -> fun y -> y / z", Term.(Fun (x, Fun (y, Binop (Var y, Div, Var z)))), None
    (* Type mismatch between argument and function *)
  ; ( "(fun x -> x + 1) true"
    , Term.(App (Fun (x, Binop (Var x, Plus, IntConst 1)), Var x))
    , None )
  ; (* Nested functions *)
    ( "(fun x -> fun y -> x + y) 3 4"
    , Term.(
        App (App (Fun (x, Fun (y, Binop (Var x, Plus, Var y))), IntConst 3), IntConst 4))
    , Some Type.Int )
  ; (* Bad Nested functions *)
    ( "(fun x -> fun y -> x + y) 3 (2, 4)"
    , Term.(
        App
          ( App (Fun (x, Fun (y, Binop (Var x, Plus, Var y))), IntConst 3)
          , Pair (IntConst 2, IntConst 4) ))
    , None )
  ; (* Function with nested pair *)
    ( "(fun x -> fst x) (1, 2)"
    , Term.(App (Fun (x, Proj (First, Var x)), Pair (IntConst 1, IntConst 2)))
    , Some Type.Int )
  ; (* Bad Application on a function application with nested pair and projection *)
    ( "(fun x -> fst (snd x)) ((1, 2), 3)"
    , Term.(
        App
          ( Fun (x, Proj (First, Proj (Second, Var x)))
          , Pair (Pair (IntConst 1, IntConst 2), IntConst 3) ))
    , None )
  ; (* Function application with nested pair and projection *)
    ( "(fun x -> fst (snd x)) ((1, 2), (3, 2))"
    , Term.(
        App
          ( Fun (x, Proj (First, Proj (Second, Var x)))
          , Pair (Pair (IntConst 1, IntConst 2), Pair (IntConst 3, IntConst 2)) ))
    , Some Type.Int )
    (* ; (* x -> y -> (x -> y -> z) -> z *)
       ( "fun x -> fun y -> fun z -> z x y"
       , Term.(Fun (x, Fun (y, Fun (z, App (Var z, App (Var x, Var y))))))
       , Some
       Type.(
       Arrow (Var x, Arrow (Var y, Arrow (Arrow (Var x, Arrow (Var y, Var z)), Var z))))
       ) *)
  ]
;;

let typeModule = (module Type : Alcotest.TESTABLE with type t = Type.t)

let check_typeof term_text term expected_type =
  let open Alcotest in
  test_case term_text `Quick (fun () ->
    check (option typeModule) "Same type" expected_type (Inference.typeof term))
;;

let () =
  let open Alcotest in
  run
    "Inference"
    [ ( "typeof"
      , List.map
          (fun (term_text, term, expected_type) ->
            check_typeof term_text term expected_type)
          tests_typeof )
    ]
;;