module Local

  use int.Int

  let test1 () =
    let x = 1 in
    assert { x >= 0 }; (* ok *)
    assert { forall y. y = 1 -> x = y }; (* ok *)
    assert { x = 0 } (* failure *)

  predicate p (_: unit) = false

  let test2 () =
    let ref x = () in
    assert { not p x } (* ok *)

  let test3 () =
    let ref x = 1 in
    assert { x >= 0 }; (* ok *)
    assert { x > 10 } (* failure *)

  let test4 () =
    assert { 4 >= 3 || 1 < 3 } (* ok *)

  let test5 () =
    let x = True in
    assert { x = True }
end

module Global

  use int.Int

  val ref x : int

  let test1 () =
    assert { x = 0 }; (* ok *)
    x <- x + 1;
    assert { x = 1 } (* ok *)

  let ref y = 42

  let test2 () =
    assert { y = 42 } (* ok *)
end

module Functions

  use int.Int

  let function f1_aux (x: int): int = x+100

  let ref x = 0

  let test1 () =
    let function f1_aux' (x: int) (_: int) : int = x+200 in
    assert { 101 = f1_aux 1 }; (* ok *)
    assert { 200 = f1_aux' 0 1 }; (* ok *)
    assert { f1_aux 0 = f1_aux' 0 1 } (* failure *)

 let test2 () =
    let function f1_aux (x: int) (_: int) : int = x+200 in
    assert { f1_aux 0 0 = 200 } (* ok *)

  let test3_aux (y: int) : int
    requires { y > 0 }
    ensures { result = 1 }
  = y

  let test3a () =
    test3_aux 0 (* failure in precondition *)

  let test3b () =
    let _ = test3_aux 1 in (* ok *)
    test3_aux 2 (* failure in postcondition *)

  exception E int

  let test4_aux () : unit
    raises { E j -> j = x }
  = x <- 1;
    raise (E 42)

  let test4 () =
    try test4_aux () (* failure in exceptional postcondition *)
    with E _ -> () end

  let test5_aux (d: int)
    ensures { x = old x + 1 }
  = x <- x + d

  let test5 () =
    test5_aux 1; (* ok *)
    test5_aux 2 (* failure in postcondition *)

  let test6_aux (ref x: 'a) (y: 'a)
    requires { x <> y }
    ensures { result = y }
  = x <- y;
    x

  let test6 () =
    let ref r = 1 in (* pres, posts of ref ok *)
    let _ = test6_aux r 42 in (* ok *)
    let _ = test6_aux r 42 in (* failure in precondition *)
    ()
end

module Loops
  use int.Int

  val ref x : int

  let test1 () requires { x = 0 } =
    for i = 1 to 10 do
      invariant { 0 < i <= 10 } (* ok *)
    done

  let test2 () requires { x = 0 } =
    for i = 0 to 5 do
      invariant { x = 0 || x = 1 } (* failure when i = 1, x = 2 *)
      x <- x + 1
    done

  let test3 ()  =
    while x < 10 do
      variant { 10 - x }
      invariant { 5 >= x } (* failure when x = 6 *)
      x <- x + 1
    done;
    assert { x = 100 }
end

module Aliasing
  use int.Int

  type r = { mutable value: int }

  let test1 () = (* aliasing *)
    let x = {value= 0} in
    let y = x in
    y.value <- 100;
    assert { (x.value, y.value) = (100, 100) };
    x.value, y.value

  let test2_aux (b: bool) = (* aliasing *)
    let x = {value= 0} in
    let y = x in
    (if b then
      y.value <- 100
    else
      x.value <- 111);
    assert { y.value >= 100 };
    x.value, y.value

  let test2 () =
    let x, y = test2_aux True in
    assert { (x, y) = (100, 100) };
    let x, y = test2_aux False in
    assert { (x, y) = (111, 111) }
end

module Labels
  use int.Int

  type r = { mutable value: int }

  let test1 () =
    let x = {value= 0} in
    label A in
    x.value <- x.value + 1;
    label B in
    x.value <- x.value + 1;
    assert { x.value = x.value at B + 1 = x.value at A + 2 }
end


(* TODO *)
module Chars
  use int.Int

  type char = < range 0 255 >

  val (<=) (c1 c2: char) : bool
     ensures { result <-> char'int c1 <= char'int c2 }

  val (-) (c1 c2: char) : char
     ensures { char'int result = char'int c1 - char'int c2 }

  let uppercase (c:char) =
     if (97:char) <= c <= (122:char)  then
        c - 32 else c

  let test1 () : char =
    uppercase 97
end

(* TODO *)
module Strings
  use string.OCaml
  let test1 () =
    let s = "hello" in
    concat s " world"
end

module RecordMutGhost
  use int.Int

  type t = {mutable f: int -> int; ghost n: int}

  let test () =
    let x = { f= (fun (x:int) -> x); n= 2 } in
    assert { x.n = 2 }; (* ok *)
    assert { x.f 3 = 3 }; (* ok *)
    x.f <- (let g = x.f in (fun (z: int) -> if z = 3 then 42 else g z));
    assert { x.f 3 = 42 }; (* ok *)
    1 (* + x.n *) + x.f 3
end

module PolyContext

  let poly (x: 'a) : 'a = x

  (* Call polymorphic function from context with statically concrete types *)
  let test1 () =
    let x = 42 in
    poly x

  (* Call polymorphic function from context with statically polymorphic types *)
  let poly' (x: 'a) : 'a =
    poly x (* ← *)

  let test2 () =
    let x = 42 in
    poly' x
end

module PolyRefContracts
  use int.Int
  use ref.Ref

  let test1 () = (* check contracts in ref *)
    let r = ref 0 in
    r := !r + 1

  let test2_aux (ref argref: 'a) (argx argy: 'a)
    requires { argref <> argx }
    ensures { result = argx }
    ensures { result = argy }
  = argref <- argx;
    argref

  let ref myref = 0

  let test2a () =
    test2_aux myref 1 1 (* OK *)

  let test2b () =
    test2_aux myref 0 1 (* PRE fails *)

  let test2c () =
    test2_aux myref 1 2 (* POST fails *)
end

module RecordPoly
  use int.Int
  type t1 'a = { mutable x: 'a }

  let test1 () =
    let r = {x= 1} in
    assert { r.x = 1 }; (* ok *)
    r.x <- 100;
    assert { r.x = 100 } (* ok *)

  type t2 'a = { mutable f: int -> 'a }
  let test2 () =
    let r = { f= fun (_: int) -> True } in
    assert { r.f 0 = True }; (* ok *)
    assert { r.f 1 = True }; (* ok *)
    r.f <- (let f = r.f in (fun (j: int) -> if j = 0 then False else f j));
    assert { r.f 0 = False }; (* ok *)
    assert { r.f 1 = True }; (* ok *)

  type t3 'a = { mutable elts: int -> 'a; length: int }

  let update (r: t3 'a) (i: int) (x: 'a) : unit
    requires { 0 < i < r.length }
    ensures { forall j. 0 < j < r.length -> j <> i -> r.elts j = old r.elts j }
    ensures { r.elts i = x }
  = r.elts <- let f = r.elts in (fun (j: int) -> if j = i then x else f j)

  use bool.Bool

  let test3 () =
    let r = { length= 5; elts= fun (_:int) -> True } in
    assert { r.elts 3 = True }; (* ok *)
    update r 3 False; (* post condition cannot be evaluated *)
    assert { r.elts 3 = False }; (* ok *)
    assert { r.elts 2 = True }; (* ok *)
    assert { r.elts 1 = False } (* failure *)
end

module PolyFunc
  use int.Int

  let mk v = fun (_: int) -> v

  let test1 () =
      mk 0

  type t 'a = {mutable f: int -> 'a}

  let function map_update (f: int -> 'a) (i: int) (x: 'a) =
    fun (j: int) -> if j = i then x else f j

  let ([]<-) (x: t 'a) (i: int) (v: 'a) : unit =
    x.f <- map_update x.f i v

  let test2 () =
    let x = {f= fun _ -> 0} in
    (* x.f <- map_update x.f 0 111; *)
    x[0] <- 111;
    x.f 0, x.f 1
end

module ArrayExec
  use int.Int

  type array 'a = {length: int; mutable elts: int -> 'a}

  let empty () : array 'a = {length= 0; elts= fun _ -> any 'a}

  let make (n: int) (v: 'a) : array 'a
    ensures { result.length = n }
    ensures { forall i. 0 <= i < n -> result.elts i = v }
  = {length= n; elts= fun _ -> v}

  let function map_update (f: int -> 'a) (i: int) (x: 'a) =
    fun (j: int) -> if j = i then x else f j

  let ([]) (a: array 'a) (i: int) : 'a = a.elts i

  function ([]') (a: array 'a) (i: int) : 'a = a.elts i

  let ([]<-) (a: array 'a) (i: int) (v: 'a) : unit
    = a.elts <- map_update a.elts i v

  let test1 () =
    let a = {length= 5; elts= fun (_: int) -> 0} in
    a.elts 0, a.elts 1

  let test2 () =
    let a = {length= 5; elts= fun (_: int) -> 0} in
    a.elts <- (let f = a.elts in (fun (i: int) -> if i = 0 then 11 else f i));
    a.elts 0, a.elts 1

  let test3 () =
    let a = make 10 0 in
    a[0] <- 10;
    a[0], a[1]

  let test4 () =
    let a = make 10 0 in
    assert { length a = 10 };
    let x = a[0] in
    assert { x = 0 }; (* ok *)
    a[1] <- 11;
    let x = a[0] in
    assert { x = 0 }; (* ok *)
    let x = a[1] in
    assert { x = 11 }; (* ok *)
    let x = a[2] in
    assert { x = 11 }; (* failure *)
    ()

  clone export array.Array with
    type array = array,
    val make = make,
    val empty = empty,
    val ([]) = ([]),
    val ([]<-) = ([]<-)
end

(* TODO dispatch array.Array/rac.ArrayExec by module substitution *)
module Arrays
  use array.Array

  meta compute_max_steps 10_000

  let test0 () =
    let a = make 10 0 in
    a[0] <- 10;
    a[0], a[1]

  let test1 () =
    let a = make 10 0 in
    assert { length a = 10 };
    assert { a[0] = 0 }; (* ok *)
    a[1] <- 11;
    assert { a[0] = 0 }; (* ok *)
    assert { a[1] = 11 }; (* ok *)
    assert { a[2] = 11 }; (* failure *)
    ()

  use int.Int

  let test2 () =
    let a = make 1 11 in
    assert { a.length = 42 } (* failure *)

  let test3 () =
    let a = (empty () : array int) in
    assert { a.length = 42 } (* failure *)

  let test4 () =
    let a : array int = empty () in
    assert { a.length > 0 } (* failure, cannot evaluate TODO *)
end

module Ghost
  use int.Int
  type t = {i1: int; ghost i2: int}

  let test1 () : (int, ghost int) =
    let x = {i1= 123; i2= 456} in
    assert { x.i1 = 123 /\ x.i2 = 456 };
    (x.i1, x.i2)

  let ghost f1 (x y : int) : int =
    if x = y then 123 else 456

  let ghost test2 () =
    let ghost x = f1 0 1 in
    x + 1
end

module Predicates
  use int.Int

  predicate p1 (n: int) =
    n = 123

  let test1 () =
    assert { p1 123 };
    let ghost x1 = p1 123 in
    assert { x1 };
    let ghost x2 = not p1 0 in
    assert { x2 };
    let predicate p2 (x: int) = x = 456 in
    assert { p2 456 }

  let ghost test2 () =
    let b1 = pure { true } in
    assert { b1 };
    let x = if b1 then 0 else 1 in
    assert { x = 0 };
    let b2 = p1 123 in
    assert { b2 };
    (b1, x, b2)
end

module Variants
  use int.Int
  let loop_var_ok1 () =
    let ref i = 1 in
    while i >= 0 do
      variant { i }
      i <- i - 1
    done
  let loop_var_ok2 () =
    let ref i = 1 in
    let ref j = 1 in
    while j >= 0 do
      variant { i, j }
      if i > 0 then
        i <- i - 1
      else
        j <- j - 1
    done
  let loop_var_fail1 () =
    let ref i = 2 in
    while i >= 0 do
      variant { i }
      if i > 1 then
        i <- i - 1
    done
  let loop_var_fail2 () =
    let ref i = 1 in
    let ref j = 2 in
    while j >= 0 do
      variant { i, j }
      if i > 0 then
        i <- i - 1
      else if j > 1 then
        j <- j - 1
    done

  use relations.WellFounded
  predicate p (i i_old: int) =
    i_old <= 2 /\ i > i_old
  let loop_custom_variant () =
    let ref i = 1 in
    while i >= 0 do
      variant { 2 - i with p }
      i <- i - 1
    done;
    assert { i = -1 };
    i <- 1;
    while i >= -1 do
      variant { 2 - i with p }
      i <- i - 1
    done
end

module TestUInt64
  use mach.int.UInt64
  let test () = 0 - 1
end

module Premises
  val predicate p (x: int)
  val predicate q (x: int)

  axiom a: forall x. p x -> q x

  let test1 () : unit
    requires { p 42 }
    ensures { q 42 }
  = ()
end

module FunctionVariant
  use int.Int

  let rec test1 (i: int) variant { i } (* ok *)
  = if i > 0 then test1 (i-1) else 0

  let rec test2 (i: int) variant { i }
  = if i > 0 then test2 (i-1) else test2 i (* failure *)

  let rec test3 (i: int) variant { i } (* ok *)
  = if i = 0 then 0 else test3' (i-1) (i-1)
  with test3' (i j: int) variant { i, j }
  = if j = 0 then test3 i else test3' i (j - 1)

  let rec test4 (i: int) variant { i }
  = if i = 0 then 0 else test4' (i-1) (i-1)
  with test4' (i j: int) variant { i, j }
  = test4' i (j - 1) (* failure *)

  let rec test5 (i: int) variant { i }
  = if i = 0 then 0 else test5' i i (* failure *)
  with test5' (i j: int) variant { i, j }
  = if j = 0 then test5 i else test5' i (j - 1)
end

module ArrayList
  use int.Int
  use array.Array
  use list.List
  use array.ToList

  let main () =
    let f = make 4 0 in
    to_list f 0 4
end
