[Header]
open Libsm5

let check_cycle (p : Sm5.command) : bool =
  let open Sm5 in
  let rec dfs (cmd : Sm5.command) visited : bool =
    if List.exists (fun c -> c == cmd) visited then true
    else
      match cmd with
      | ( MALLOC | BOX _ | UNBOX _ | BIND _ | UNBIND | GET | PUT | CALL | ADD
        | SUB | MUL | DIV | EQ | LESS | NOT | POP | STORE | LOAD
        | PUSH (Id _ | Val _) )
        :: c ->
          dfs c (cmd :: visited)
      | PUSH (Fn (_, c1)) :: c2 ->
          dfs c1 (cmd :: visited) || dfs c2 (cmd :: visited)
      | JTR (c1, c2) :: c3 ->
          dfs c1 (cmd :: visited)
          || dfs c2 (cmd :: visited)
          || dfs c3 (cmd :: visited)
      | [] -> false
  in
  dfs p []

let test_run pgm_str =
  (* Translate the program *)
  let pgm = Parser.program Lexer.start (Lexing.from_string pgm_str) in
  let sm5_pgm = Translator.trans pgm in

  if check_cycle sm5_pgm then
    (* If a cycle is detected, mark with message *)
    print_endline "__CYCLE__";
  (* Run the translated program *)
  Sm5.run ~debug:false ~malloc:Sm5.default_malloc sm5_pgm

[Test]
(* 1. read, write, arithmetic *)
test_run "let x :=0 in (read x; write (1 + 6 * x - 8 / 4))"

[Input]
4

[Output]
23


[Test]
(* 2. branches *)
test_run
"let x := 0 in
read x;
if (5 < x) then (
  if (false = false) then (
    if not(true = 3) then (write (21)) else (write (22))
  )
  else (write (23))
)
else (write (24))"

[Input]
6

[Output]
21


[Test]
(* 3. Result of sequence expression *)
test_run
"let x := 0 in
read x;
write(11; write(12); x)"

[Input]
13

[Output]
12
13


[Test]
(* 4. Result of WRITE expression *)
test_run
"let x := 0 in
read x;
write((write((write(x)) + 10)) - 20)"

[Input]
20

[Output]
20
30
10

[Test]
(* 5. While + LetV *)
test_run
"let x := 0 in
read x;
let i := 0 in
let sum := 0 in
while (i < x) do (
   sum := sum + i;
   i := i + 1
);
write(sum)"

[Input]
11

[Output]
55


[Test]
(* 6. Factorial with while loop *)
test_run
"
let x := 1 in
let y := 0 in
let z := 0 in
read y;
while (z < y) do (
  z := z + 1;
	x := x * z
);
write x
"
[Input]
7

[Output]
5040


[Test]
(* 7. nested loop *)
test_run
"let x := 0 in
read x;
let sum := 0 in
let i := 0 in
let j := 0 in
while (i < x) do (
   while ( j < i ) do (
     j := i + j
   );
   sum := sum + i;
   i := i + 1
);
write(i);
write(j);
write(sum)
"

[Input]
10

[Output]
10
15
45


[Test]
(* 8.  LetF + CallV *)
test_run
"let proc f (x) = x * 4 in
let x := 0 in
read x;
(write (f (x) + 10))"

[Input]
3

[Output]
22

[Test]
(* 9. Recursive call *)
test_run
"let proc f (x) =
  if (x < 1) then 1 else (x + f (x-1)) in
let x:= 0 in
read x;
(write (f (x)))"

[Input]
11

[Output]
67


[Test]
(* 10. LetF + CallV + CallR *)
test_run
"let proc f (x) = (x := 3) in
let proc g (x) = (x := 3) in
let x := 1 in
let y := 2 in
f <x>; g (y); (write (x+y))"

[Input]
1
2

[Output]
5


[Test]
(* 11. Read and write *)
test_run "let x := 0 in while (0 < (read x)) do (write x)"

[Input]
3
2
5
0

[Output]
3
2
5

[Test]
(* 12. Simple bonus *)
test_run "let x := false in if x then (write 1) else (write 0)"

[Output]
0

[Test]
(* 13. Assign, revisited *)
test_run
"let x := false in
let y := false in
let z := false in
let w := 0 in
read w;
if ( (x := 1) + (y := 2) + (z := 3) = 6 ) then (write (w *10)) else (write 0)"

[Input]
10

[Output]
100


[Test]
(* 14. Scope *)
test_run
"let x := 3 in
write(
  (let x := 1 in (x := 5); ((x := x + 1) * (let x := 2 in (x := x + 1)))) +
  (let y := 2 in 2 * x)
)"

[Output]
24


[Test]
(* 15. Assignment in recursive function *)
test_run
"let sum := 0 in
let k := 0 in
let proc f (x) = (
   if (x < k) then (
     f(x+1);
     x := 2 * x;
     sum := sum + x
   ) else 0
) in
read k;
f(0);
(write sum)"

[Input]
10

[Output]
90


[Test]
(* 16. Arguments in CallF *)
test_run
"let y := 2 in
let z := 0 in
let proc f (x) = x * y + z in
read z;
write (f (3))"

[Input]
2

[Output]
8


[Test]
(* 17. Arguments in CallR *)
test_run
"let x := 1 in
let y := 2 in
let z := 3 in
let proc f (x) = (x := x + (y := y + 1) + (z := z + 1)) in
f <x>; (write (x+y+z))"

[Output]
15

[Test]
(* 18. Scope case II *)
test_run
"let a := 0 in
let proc f (x) = x + 1 in
let x := 2 in
read a;
write (a * (x + x))"

[Input]
10

[Output]
40


[Test]
(* Scope case III *)
test_run
"let x := 1 in
let y := 2 in
let z := 3 in
let proc f (x) = (x := x + y) in
let proc g (x) = (x := x + z) in
let y := 4 in
let z := 5 in
write (f ( g <y> ))"

[Output]
9

[Test]
(* 20. Complex scope *)
test_run
"let x := 1 in
let y := 5 in
let proc f (z) =
  if z then (x := x + 1) else (y := y - 2)
in
let x := 10 in
let y := 50 in
let proc g (z) =
  if z then (x := x + 5) else (y := y - 10)
in
write (
  y + x + f (false) + f (true) + y + x + g (false) + g (true)
)"

[Output]
170