Require Import Coq.ZArith.ZArith.
Require Import Coq.Lists.List.
Require Import Crypto.Util.ZRange.
Require Import Crypto.Util.LetIn.
Require Import Crypto.Language.
Require Import Crypto.UnderLets.
Require Import Crypto.AbstractInterpretation.
Require Import Crypto.RewriterAll.
Require Import Crypto.MiscCompilerPasses.
Require Import Crypto.CStringification.
Import ListNotations. Local Open Scope Z_scope.

Import Language.Compilers.
Import UnderLets.Compilers.
Import AbstractInterpretation.Compilers.
Import RewriterAll.Compilers.
Import MiscCompilerPasses.Compilers.
Import CStringification.Compilers.
Local Coercion Z.of_nat : nat >-> Z.
Import Compilers.defaults.

Local Notation "x + y"
  := ((#ident.Z_add @ x @ y)%expr)
     : expr_scope.
Local Notation "x * y"
  := ((#ident.Z_mul @ x @ y)%expr)
     : expr_scope.
Local Notation "x" := (expr.Var x) (only printing, at level 9) : expr_scope.

Example test1 : True.
Proof.
  let v := Reify ((fun x => 2^x) 255)%Z in
  pose v as E.
  vm_compute in E.
  pose (PartialEvaluate E) as E'.
  vm_compute in E'.
  lazymatch (eval cbv delta [E'] in E') with
  | (fun var => expr.Ident (ident.Literal ?v)) => idtac
  end.
  constructor.
Qed.
Module testrewrite.
  Import expr.
  Import ident.

  Eval compute in RewriteRules.RewriteNBE (fun var =>
                          (#ident.fst @ (expr_let x := ##10 in ($x, $x)))%expr).

  Notation "x + y" := (@expr.Ident base.type ident _ _ ident.Z_add @ x @ y)%expr : expr_scope.

  Eval compute in RewriteRules.RewriteNBE (fun var =>
                          ((\ x , expr_let y := ##5 in #ident.fst @ $x + (#ident.fst @ $x + ($y + $y)))
                             @ (##1, ##1))%expr).

  Eval compute in RewriteRules.RewriteNBE (fun var =>
                          ((\ x , expr_let y := ##5 in $y + ($y + (#ident.fst @ $x + #ident.snd @ $x)))
                             @ (##1, ##7))%expr).


  Eval cbv in partial.eval_with_bound partial.default_relax_zrange
                                      (RewriteRules.RewriteNBE (fun var =>
                (\z , ((\ x , expr_let y := ##5 in $y + ($z + (#ident.fst @ $x + #ident.snd @ $x)))
                         @ (##1, ##7)))%expr) _)
                (Datatypes.Some r[0~>100]%zrange, tt).
End testrewrite.
Module testpartial.
  Import expr.
  Import ident.

  Eval compute in partial.eval
                          (#ident.fst @ (expr_let x := ##10 in ($x, $x)))%expr.

  Notation "x + y" := (@expr.Ident base.type ident _ _ (ident.Z_add) @ x @ y)%expr : expr_scope.

  Eval compute in partial.eval
                          ((\ x , expr_let y := ##5 in #ident.fst @ $x + (#ident.fst @ $x + ($y + $y)))
                             @ (##1, ##1))%expr.

  Eval compute in partial.eval
                          ((\ x , expr_let y := ##5 in $y + ($y + (#ident.fst @ $x + #ident.snd @ $x)))
                             @ (##1, ##7))%expr.


  Eval cbv in partial.eval_with_bound
                partial.default_relax_zrange
                (\z , ((\ x , expr_let y := ##5 in $y + ($z + (#ident.fst @ $x + #ident.snd @ $x)))
                         @ (##1, ##7)))%expr
                (Datatypes.Some r[0~>100]%zrange, tt).
End testpartial.

Module test2.
  Example test2 : True.
  Proof.
    let v := Reify (fun y : Z
                    => (fun k : Z * Z -> Z * Z
                        => dlet_nd x := (y * y) in
                            dlet_nd z := (x * x) in
                            k (z, z))
                         (fun v => v)) in
    pose v as E.
    vm_compute in E.
    pose (partial.Eval E) as E'.
    vm_compute in E'.
    lazymatch (eval cbv delta [E'] in E') with
    | (fun var : type -> Type =>
         (λ x : var _,
                expr_let x0 := ($x * $x) in
              expr_let x1 := ($x0 * $x0) in
              ($x1, $x1))%expr) => idtac
    end.
    pose (partial.EvalWithBound partial.default_relax_zrange E' (Some r[0~>10]%zrange, tt)) as E''.
    lazy in E''.
     lazymatch (eval cbv delta [E''] in E'') with
     | (fun var : type -> Type =>
          (λ x : var _,
                 expr_let y := #(ident.Z_cast r[0 ~> 100]) @ ((#(ident.Z_cast r[0 ~> 10]) @ $x) * (#(ident.Z_cast r[0 ~> 10]) @ $x)) in
               expr_let y0 := #(ident.Z_cast r[0 ~> 10000]) @ ((#(ident.Z_cast r[0 ~> 100]) @ $y) * (#(ident.Z_cast r[0 ~> 100]) @ $y)) in
               (#(ident.Z_cast r[0 ~> 10000]) @ $y0, #(ident.Z_cast r[0 ~> 10000]) @ $y0))%expr)
      => idtac
    end.
    constructor.
  Qed.
End test2.
Module test3.
  Example test3 : True.
  Proof.
    let v := Reify (fun y : Z
                    => dlet_nd x := dlet_nd x := (y * y) in
                        (x * x) in
                        dlet_nd z := dlet_nd z := (x * x) in
                        (z * z) in
                        (z * z)) in
    pose v as E.
    vm_compute in E.
    pose (partial.Eval E) as E'.
    vm_compute in E'.
    lazymatch (eval cbv delta [E'] in E') with
    | (fun var : type -> Type =>
         (λ x : var _,
                expr_let x0 := $x * $x in
              expr_let x1 := $x0 * $x0 in
              expr_let x2 := $x1 * $x1 in
              expr_let x3 := $x2 * $x2 in
              $x3 * $x3)%expr)
      => idtac
    end.
    pose (partial.EvalWithBound partial.default_relax_zrange E' (Some r[0~>10]%zrange, tt)) as E'''.
    lazy in E'''.
    lazymatch (eval cbv delta [E'''] in E''') with
    | (fun var : type -> Type =>
          (λ x : var _,
           expr_let y := #(ident.Z_cast r[0 ~> 100]) @ ((#(ident.Z_cast r[0 ~> 10]) @ $x) * (#(ident.Z_cast r[0 ~> 10]) @ $x)) in
           expr_let y0 := #(ident.Z_cast r[0 ~> 10000]) @ ((#(ident.Z_cast r[0 ~> 100]) @ $y) * (#(ident.Z_cast r[0 ~> 100]) @ $y)) in
           expr_let y1 := #(ident.Z_cast r[0 ~> 100000000]) @ ((#(ident.Z_cast r[0 ~> 10000]) @ $y0) * (#(ident.Z_cast r[0 ~> 10000]) @ $y0)) in
           expr_let y2 := #(ident.Z_cast r[0 ~> 10000000000000000]) @ ((#(ident.Z_cast r[0 ~> 100000000]) @ $y1) * (#(ident.Z_cast r[0 ~> 100000000]) @ $y1)) in
           #(ident.Z_cast r[0 ~> 100000000000000000000000000000000]) @ ((#(ident.Z_cast r[0 ~> 10000000000000000]) @ $y2) * (#(ident.Z_cast r[0 ~> 10000000000000000]) @ $y2)))%expr)
      => idtac
    end.
    constructor.
  Qed.
End test3.
Module test3point5.
  Example test3point5 : True.
  Proof.
    let v := Reify (fun y : (list Z) => List.nth_default (-1) y 0) in
    pose v as E.
    vm_compute in E.
    pose (partial.EvalWithBound partial.default_relax_zrange E (Some [Some r[0~>10]%zrange], tt)) as E'.
    lazy in E'.
    clear E.
    lazymatch (eval cbv delta [E'] in E') with
    | (fun var : type -> Type =>
         (λ x : var _,
          #(ident.Z_cast r[0 ~> 10]) @ (#ident.List_nth_default @ #(ident.Literal (-1)%Z) @ $x @ #(ident.Literal 0%nat)))%expr)
      => idtac
    end.
    constructor.
  Qed.
End test3point5.
Module test4.
  Example test4 : True.
  Proof.
    let v := Reify (fun y : (list Z * list Z)
                    => dlet_nd x := List.nth_default (-1) (fst y) 0 in
                        dlet_nd z := List.nth_default (-1) (snd y) 0 in
                        dlet_nd xz := (x * z) in
                        (xz :: xz :: nil)) in
    pose v as E.
    vm_compute in E.
    pose (partial.Eval E) as E'.
    lazy in E'.
    clear E.
    pose (Some [Some r[0~>10]%zrange],Some [Some r[0~>10]%zrange], tt) as bound.
    pose (partial.EtaExpandWithListInfoFromBound E' bound) as E''.
    lazy in E''.
    clear E'.
    pose (PartialEvaluate E'') as E'''.
    lazy in E'''.
    pose (partial.EvalWithBound partial.default_relax_zrange E''' bound) as E''''.
    lazy in E''''.
    clear E'' E'''.
    lazymatch (eval cbv delta [E''''] in E'''') with
    | (fun var : type -> Type =>
         (λ x : var _,
          expr_let y := #(ident.Z_cast r[0 ~> 10]) @
                        (#ident.List_nth_default @ #(ident.Literal (-1)%Z) @ (#ident.fst @ $x) @ #(ident.Literal 0%nat)) in
          expr_let y0 := #(ident.Z_cast r[0 ~> 10]) @
                          (#ident.List_nth_default @ #(ident.Literal (-1)%Z) @ (#ident.snd @ $x) @ #(ident.Literal 0%nat)) in
          expr_let y1 := #(ident.Z_cast r[0 ~> 100]) @ ((#(ident.Z_cast r[0 ~> 10]) @ $y) * (#(ident.Z_cast r[0 ~> 10]) @ $y0)) in
          #(ident.Z_cast r[0 ~> 100]) @ $y1 :: #(ident.Z_cast r[0 ~> 100]) @ $y1 :: [])%expr)
      => idtac
    end.
    constructor.
  Qed.
End test4.
Module test5.
  Example test5 : True.
  Proof.
    let v := Reify (fun y : (Z * Z)
                    => dlet_nd x := (13 * (fst y * snd y)) in
                        x) in
    pose v as E.
    vm_compute in E.
    pose (RewriteRules.RewriteArith (2^8) (partial.Eval E)) as E'.
    lazy in E'.
    clear E.
    lazymatch (eval cbv delta [E'] in E') with
    | (fun var =>
         expr.Abs (fun v
              => (expr_let v0 := (#ident.Z_mul @ (#ident.fst @ $v) @ (#ident.Z_mul @ (#ident.snd @ $v) @ #(ident.Literal 13))) in
                      $v0)%expr))
      => idtac
    end.
    constructor.
  Qed.

  Example test5_2 : True.
  Proof.
    let v := Reify (fun y : (Z * Z)
                    => dlet_nd x := (2 * (19 * (fst y * snd y))) in
                        x) in
    pose v as E.
    vm_compute in E.
    pose (RewriteRules.RewriteArith (2^8) (partial.Eval E)) as E'.
    lazy in E'.
    clear E.
    lazymatch (eval cbv delta [E'] in E') with
    | (fun var =>
         expr.Abs (fun v
              => (expr_let v0 := (#ident.Z_mul @ (#ident.fst @ $v) @ (#ident.Z_mul @ (#ident.snd @ $v) @ (#ident.Z_mul @ #(ident.Literal 2) @ #(ident.Literal 19)))) in
                      $v0)%expr))
      => idtac
    end.
    constructor.
  Qed.
End test5.
Module test6.
  (* check for no dead code with if *)
  Example test6 : True.
  Proof.
    let v := Reify (fun y : Z
                    => if 0 =? 1
                       then dlet_nd x := (y * y) in
                                x
                       else y) in
    pose v as E.
    vm_compute in E.
    pose (PartialEvaluate E) as E''.
    lazy in E''.
    lazymatch eval cbv delta [E''] in E'' with
    | fun var : type -> Type => (λ x : var _, $x)%expr
      => idtac
    end.
    exact I.
  Qed.
End test6.
Module test7.
  Example test7 : True.
  Proof.
    let v := Reify (fun y : Z
                    => dlet_nd x := y + y in
                        dlet_nd z := x in
                        dlet_nd z' := z in
                        dlet_nd z'' := z in
                        z'' + z'') in
    pose v as E.
    vm_compute in E.
    pose (Subst01.Subst01 (DeadCodeElimination.EliminateDead E)) as E''.
    lazy in E''.
    lazymatch eval cbv delta [E''] in E'' with
    | fun var : type -> Type => (λ x : var _, expr_let v0 := $x + $x in $v0 + $v0)%expr
      => idtac
    end.
    exact I.
  Qed.
End test7.
Module test8.
  Example test8 : True.
  Proof.
    let v := Reify (fun y : Z
                    => dlet_nd x := y + y in
                        dlet_nd z := x in
                        dlet_nd z' := z in
                        dlet_nd z'' := z in
                        z'' + z'') in
    pose v as E.
    vm_compute in E.
    pose (GeneralizeVar.GeneralizeVar (E _)) as E''.
    lazy in E''.
    unify E E''.
    exact I.
  Qed.
End test8.
Module test9.
  Example test9 : True.
  Proof.
    let v := Reify (fun y : list Z => (hd 0%Z y, tl y)) in
    pose v as E.
    vm_compute in E.
    pose (PartialEvaluate E) as E'.
    lazy in E'.
    clear E.
    lazymatch (eval cbv delta [E'] in E') with
    | (fun var
       => (λ x,
           (#ident.list_case
              @ (λ _, #(ident.Literal 0%Z))
              @ (λ x0 _, $x0)
              @ $x,
            #ident.list_case
              @ (λ _, #ident.nil)
              @ (λ _ x0, $x0)
              @ $x))%expr)
      => idtac
    end.
    exact I.
  Qed.
End test9.
(*
Module test10.
  Example test10 : True.
  Proof.
    let v := Reify (fun (f : Z -> Z -> Z) x y => f (x + y) (x * y))%Z in
    pose v as E.
    vm_compute in E.
    pose (Uncurry.expr.Uncurry (partial.Eval true (canonicalize_list_recursion E))) as E'.
    lazy in E'.
    clear E.
    lazymatch (eval cbv delta [E'] in E') with
    | (fun var =>
         (λ v,
          ident.fst @@ $v @
                    (ident.fst @@ (ident.snd @@ $v) + ident.snd @@ (ident.snd @@ $v)) @
                    (ident.fst @@ (ident.snd @@ $v) * ident.snd @@ (ident.snd @@ $v)))%expr)
      => idtac
    end.
    constructor.
  Qed.
End test10.
 *)
(*
Module test11.
  Example test11 : True.
  Proof.
    let v := Reify (fun x y => (fun f a b => f a b) (fun a b => a + b) (x + y) (x * y))%Z in
    pose v as E.
    vm_compute in E.
    pose (Uncurry.expr.Uncurry (partial.Eval true (canonicalize_list_recursion E))) as E'.
    lazy in E'.
    clear E.
    lazymatch (eval cbv delta [E'] in E') with
    | (fun var =>
         (λ x,
          ident.fst @@ $x + ident.snd @@ $x + ident.fst @@ $x * ident.snd @@ $x)%expr)
      => idtac
    end.
    constructor.
  Qed.
End test11.
 *)
Module test12.
  Example test12 : True.
  Proof.
    let v := Reify (fun y : list Z => repeat y 2) in
    pose v as E.
    vm_compute in E.
    pose (Some (repeat (@None zrange) 3), tt) as bound.
    pose (PartialEvaluate (partial.EtaExpandWithListInfoFromBound E bound)) as E'.
    lazy in E'.
    clear E.
    lazymatch (eval cbv delta [E'] in E') with
    | (fun var
       => (λ x, [ [ $x[[0]] ; $x[[1]]; $x[[2]] ] ; [ $x[[0]] ; $x[[1]]; $x[[2]] ] ])%expr)
      => idtac
    end.
    exact I.
  Qed.
End test12.
Module test13.
  Example test13 : True.
  Proof.
    let v0 := constr:(nat_rect (fun _ => nat -> nat) (fun v => v) (fun n' rec v => (n' + rec (S v))%nat) 3 0%nat) in
    let v := Reify v0 in
    pose v as E;
      pose v0 as exp.
    vm_compute in E.
    vm_compute in exp.
    pose (PartialEvaluate E) as E'.
    vm_compute in E'.
    clear E.
    let r := Reify exp in
    unify r E'.
    exact I.
  Qed.
End test13.