open Alcotest
open Syntax
open Parse_tree

let parse program = Option.get @@ Parsing.Driver.parse_string program

let check_pretty_printing_consistency program_str =
  let program = parse program_str in
  let pp_meta _ _ = () in
  let pretty_printed_program =
    Format.asprintf "%a" (Program.pp ~pp_meta) program
  in
  let parsed_program = parse pretty_printed_program in
  let pretty_printed_again =
    Format.asprintf "%a" (Program.pp ~pp_meta) parsed_program
  in
  Alcotest.(check string)
    "Pretty printed program is consistent" pretty_printed_program
    pretty_printed_again

let test_variant _ =
  let program_str =
    {|
    type ty =
        Int
        | Bool
        | Arrow of ty * ty

    let ty : ty = Arrow (Int, Arrow (Int, Int)) in
    let is_arrow : ty -> bool = fun ty : ty ->
        match ty with
        | Int -> false
        | Bool -> false
        | Arrow ty1 ty2 -> true
    in
    print (is_arrow ty) ;
    print (is_arrow (Int))
|}
  in
  check_pretty_printing_consistency program_str

let test_fib _ =
  let program_str =
    {|
    let rec fib = fun n ->
        if n < 2 then
            n
        else
            fib (n - 1) + fib (n - 2)
    in
    print_int (fib 47)
|}
  in
  check_pretty_printing_consistency program_str

let test_matmult _ =
  let program_str =
    {|
let create =
    fun size ->
        let x = vector (size * size) 0 in
        for i = 0 to size do
            for j = 0 to size do
                vector_set x ((size * i) + j) (j + i)
            done
        done ;
        x
in

let mult =
    fun x ->
    fun x1 -> fun x2 ->
    fun y ->
    fun y1 -> fun y2 ->
    let r = vector (y2 * x1) 0 in
    for i = 0 to x1 do
        for j = 0 to y2 do
            for k = 0 to y1 do
                vector_set r ((i * y2) + j)
                    ((vector_get r ((i * y2) + j)) +
                        (vector_get x ((i * x2) + k)) *
                        (vector_get y ((k * y2) + j)))
            done
        done
    done ;
    r
in

let size = 300 in

let ar = size in
let ac = size in
let br = size in
let bc = size in

let a = create ar in
let b = create br in

print (vector_get (mult a ar ac b br bc) ((ar * bc) - 1))
|}
  in
  check_pretty_printing_consistency program_str

let tests =
  [
    ("Pretty printing variant", `Quick, test_variant);
    ("Pretty printing fib", `Quick, test_fib);
    ("Pretty printing matmult", `Quick, test_matmult);
  ]

let () = run "test_pretty_printer" [ ("pretty_printing", tests) ]
