Require Export ClassicalEpsilon.
Require Export Bool.
Require Export SetoidDec.
Require Export SMT.
Require Export Programs.
Require Import String.
Require Export List.
Require Export Outcomes.

Definition sstore := StringMap.t term.
Inductive schunk := SChunk(p: pred)(ts: list term).
Definition spointsto l v := SChunk pointsto_pred [l, v].
Definition smalloc_block l v := SChunk malloc_block_pred [l, v].
Definition sheap := list schunk.
Inductive sstate := SState(pc: formula)(ss: sstore)(sh: sheap).

Definition sstore0: sstore := StringMap.empty _.
Definition sstate0 := SState f_true sstore0 nil.

Definition sstore_lookup s x :=
  match StringMap.find x s with Some t => t | None => t_lit 0 end.
Definition sstore_update (s: sstore) x t := StringMap.add x t s.
Fixpoint sstore_updates (s: sstore) xs ts :=
  match xs, ts with
    x::xs, t::ts => sstore_updates (sstore_update s x t) xs ts
  | _, _ => s
  end.

Fixpoint seval(s: sstore)(e: expr) :=
  match e with
    e_lit z => t_lit z
  | e_var x => sstore_lookup s x
  | e_add e1 e2 => t_add (seval s e1) (seval s e2)
  end.

Fixpoint sbeval(s: sstore)(b: bexpr) :=
  match b with
    b_eq e1 e2 => f_eq (seval s e1) (seval s e2)
  | b_lt e1 e2 => f_lt (seval s e1) (seval s e2)
  | b_not b => f_not (sbeval s b)
  end.

Definition soutcome A := outcome sstate A.
Definition smutatora A := sstate -> soutcome A.
Definition smutator := smutatora unit.

Definition fresh_mut: smutatora term :=
  fun st => let (pc, s, h) := st in
  let t := t_symb (fresh_symbol (formula_symbols pc)) in
  single (SState (f_and pc (f_eq t t)) s h) t.

Definition assume_formula(phi: formula): smutator :=
  fun st => let (pc, s, h) := st in
  if follows pc (f_not phi) then
    oblock
  else
    single (SState (f_and pc phi) s h) tt.

Definition update_sstore(x: var)(t: term): smutator :=
  fun st => let (ph, s, h) := st in
  single (SState ph (sstore_update s x t) h) tt.

Definition update_sstore_n(xs: list var)(ts: list term): smutator :=
  fun st => let (ph, s, h) := st in
  single (SState ph (sstore_updates s xs ts) h) tt.

Definition shavoc1(x: var): smutator :=
  v <- fresh_mut;
  update_sstore x v.

Fixpoint shavoc(xs: list var): smutator :=
  match xs with
    [] => noop
  | x::xs => shavoc1 x; shavoc xs
  end.

Definition swith_store{A}(s: sstore)(op: smutatora A): smutatora A :=
  fun st => let (pc, s0, h) := st in
  bindf
    (op (SState pc s h))
    (fun a st => let (pc, _, h) := st in single (SState pc s0 h) a).

Definition swith_store'(s: sstore)(op: smutator): smutatora sstore :=
  fun st => let (pc, s0, h) := st in
  bindf
    (op (SState pc s h))
    (fun _ st => let (pc, s, h) := st in single (SState pc s0 h) s).

Definition swith_local_store(op: smutator): smutator :=
  fun st => let (pc, s0, h) := st in
  bindf
    (op (SState pc s0 h))
    (fun _ st => let (pc, s, h) := st in single (SState pc s0 h) tt).

Definition seval_mut(e: expr): smutatora term :=
  fun st => let (pc, s, h) := st in
  single st (seval s e).

Definition sevals_mut(es: list expr) :=
  fun st: sstate => let (pc, s, h) := st in
  single st (map (seval s) es).

Definition sbeval_mut(b: bexpr): smutatora formula :=
  fun st => let (pc, s, h) := st in
  single st (sbeval s b).

Definition sassume_bexpr(b: bexpr): smutator :=
  seqf (sbeval_mut b) assume_formula.

Definition sassert_bexpr(b: bexpr): smutator :=
  fun st => let (pc, s, h) := st in
  if follows pc (sbeval s b) then
    single st tt
  else
    ofail.

Definition sclear_heap: smutator :=
  fun st => let (pc, s, h) := st in
  single (SState pc s nil) tt.

Definition sleakcheck: smutator :=
  fun st => let (pc, s, h) := st in
  match h with
    nil => oblock
  | SChunk p _::_ => message ("Leaking chunk " ++ p) ofail
  end.

Fixpoint extract0{A}{B}(f: A -> option B)(todo done: list A): option (list A * B) :=
  match todo with
    nil => None
  | x::xs =>
    match f x with
      None => extract0 f xs (x::done)
    | Some y => Some (done ++ xs, y)
    end
  end.

Definition extract{A B}(f: A -> option B)(xs: list A): option (list A * B) :=
  extract0 f xs nil.

Fixpoint forall2b{A B}(xs: list A)(ys: list B)(f: A -> B -> bool): bool :=
  match xs, ys with
    x::xs, y::ys => f x y && forall2b xs ys f
  | _, _ => true
  end.

Definition match_chunk pc p ts n (c: schunk) :=
  let (p', ts') := c in
  if string_dec p p' then
    if (length ts' == length ts + n)%nat then
      if forall2b ts (firstn (length ts) ts') (fun t t' => follows pc (f_eq t t')) then
        Some (skipn (length ts) ts')
      else
        None
    else
      None
  else
    None.

Definition extract_chunk p ts n: smutatora (list term) :=
  fun st => let (pc, s, h) := st in
  match extract (match_chunk pc p ts n) h with
    None => ofail
  | Some (h, ts') => single (SState pc s h) ts'
  end.

Definition extract_chunk_1_1 p l :=
  ts <- extract_chunk p [l] 1;
  yield (hd (t_lit 0) ts).

Definition extract_pointsto l := extract_chunk_1_1 pointsto_pred l.
Definition extract_malloc_block l := extract_chunk_1_1 malloc_block_pred l.

Fixpoint extract_pointstos l n :=
  match n with
    O => noop
  | S n =>
    extract_pointsto l;
    extract_pointstos (t_add l (t_lit 1)) n
  end.

Definition sprod_chunk(c: schunk): smutator :=
  fun st => let (pc, s, h) := st in
  single (SState pc s (c::h)) tt.

Fixpoint sprod_pointstos l n: smutator :=
  match n with
    O => noop
  | S n =>
    v <- fresh_mut;
    sprod_chunk (spointsto l v);
    sprod_pointstos (t_add l (t_lit 1)) n
  end.

Definition constant_term_value t: smutatora Z :=
  match t with
    t_lit z => yield z
  | _ => fail
  end.

Fixpoint sconsume(a: asn): smutator :=
  match a with
    BAsn b => sassert_bexpr b
  | PAsn p es xs =>
    message_mut "sconsume PAsn";
    ts <- sevals_mut es;
    ts' <- extract_chunk p ts (length xs);
    update_sstore_n xs ts'
  | IfAsn b a1 a2 =>
    fork (
      sassume_bexpr b;
      message_mut "sconsume IfAsn true";
      sconsume a1
    ) (
      sassume_bexpr (b_not b);
      message_mut "sconsume IfAsn false";
      sconsume a2
    )
  | SepAsn a1 a2 =>
    message_mut "sconsume SepAsn a1";
    sconsume a1;
    message_mut "sconsume SepAsn a2";
    sconsume a2
  end.

Fixpoint sproduce(a: asn): smutator :=
  match a with
    BAsn b => sassume_bexpr b
  | PAsn p es xs =>
    ts <- sevals_mut es;
    ts' <- iter fresh_mut (length xs);
    update_sstore_n xs ts';
    sprod_chunk (SChunk p (ts ++ ts'))
  | IfAsn b a1 a2 =>
    fork (
      sassume_bexpr b;
      sproduce a1
    ) (
      sassume_bexpr (b_not b);
      sproduce a2
    )
  | SepAsn a1 a2 =>
    sproduce a1;
    sproduce a2
  end.

Section SymbolicExecution.

Variable routine_specs: spec_table.
Variable pred_defs: pred_table.

Fixpoint s_exec(c: cmd): smutator :=
  match c with
  | Assign x e =>
    t <- seval_mut e;
    update_sstore x t
  | Malloc x n =>
    message_mut "Malloc";
    l <- fresh_mut;
    update_sstore x l;
    assume_formula (f_lt (t_lit 0) l);
    sprod_chunk (smalloc_block l (t_lit (Z.of_nat n)));
    sprod_pointstos l n
  | Free e =>
    message_mut "Free";
    tl <- seval_mut e;
    tn <- extract_malloc_block tl;
    n <- constant_term_value tn;
    extract_pointstos tl (Z.to_nat n)
  | Read x e =>
    message_mut "Read";
    t1 <- seval_mut e;
    t2 <- extract_pointsto t1;
    sprod_chunk (spointsto t1 t2);
    update_sstore x t2
  | Write e1 e2 =>
    message_mut "Write";
    t1 <- seval_mut e1;
    t2 <- seval_mut e2;
    _ <- extract_pointsto t1;
    sprod_chunk (spointsto t1 t2)
  | Call x r es =>
    message_mut "Call";
    match StringMap.find r routine_specs with
      None => fail
    | Some (RoutineSpec xs pre post) =>
      if length es == length xs then
        ts <- sevals_mut es;
        s1 <- swith_store' (sstore_updates sstore0 xs ts) (sconsume pre);
        t <- fresh_mut;
        swith_store (sstore_update s1 result t) (sproduce post);
        update_sstore x t
      else
        fail
    end
  | IfCmd b c1 c2 =>
    fork (
      sassume_bexpr b;
      message_mut "IfCmd true";
      s_exec c1
    ) (
      sassume_bexpr (b_not b);
      message_mut "IfCmd false";
      s_exec c2
    )
  | While b a c =>
    message_mut "While";
    swith_local_store (sconsume a);
    shavoc (targets c);
    fork (
      sclear_heap; (
        swith_local_store (sproduce a);
        sassume_bexpr b;
        s_exec c;
        swith_local_store (sconsume a)
      );
      sleakcheck
    ) (
      swith_local_store (sproduce a);
      sassume_bexpr (b_not b)
    )
  | Seq c1 c2 =>
    s_exec c1;
    s_exec c2
  | Open p es =>
    message_mut "Open";
    match pred_defs p with
      None => fail
    | Some (PredDef xs a) =>
      if le_dec (length es) (length xs) then
        ts1 <- sevals_mut es;
        ts2 <- extract_chunk p ts1 (length xs - length es);
        swith_store (sstore_updates sstore0 xs (ts1 ++ ts2)) (
          sproduce a
        )
      else
        fail
    end
  | Close p es =>
    message_mut "Close";
    match pred_defs p with
      None => fail
    | Some (PredDef xs a) =>
      if length es == length xs then
        ts <- sevals_mut es;
        swith_store (sstore_updates sstore0 xs ts) (
          sconsume a
        );
        sprod_chunk (SChunk p ts)
      else
        fail
    end
  | Skip =>
    message_mut "Skip"
  | Message m =>
    message_mut m
  end.

Inductive eresult A := ok | error: A -> eresult A.

Global Implicit Arguments ok [A].
Global Implicit Arguments error [A].

Fixpoint eresult_bind {A} (r: eresult A) (op: eresult A): eresult A :=
  match r with
    ok => op
  | error msg => error msg
  end.

Notation "op1 &&> op2" := (eresult_bind op1 op2) (at level 90).

Definition eresult_annotate(msg: string)(r: eresult string): eresult string :=
  match r with
    ok => ok
  | error msg' => error (msg ++ msg')%string
  end.

Fixpoint sresult_ok0{A}(prefix: string)(r: soutcome A): eresult string :=
  match r with
    single st a => ok
  | demonic (set_ n_Empty_set _) => ok
  | demonic (set_ n_bool o) => sresult_ok0 prefix (o true) '_x">&&> sresult_ok0 prefix (o false)
  | demonic _ => if excluded_middle_informative (safe r) then ok else error prefix
  | angelic (set_ n_Empty_set _) => error (prefix ++ " -> " ++ "failure")%string
  | angelic (set_ n_bool o) => error (prefix ++ " -> " ++ "angelic fork")%string
  | angelic _ => if excluded_middle_informative (safe r) then ok else error prefix
  | message msg o => sresult_ok0 (prefix ++ "; " ++ msg)%string o
  end.

Definition sresult_ok{A} := sresult_ok0 (A:=A) ""%string.

Section SRoutineDefs.

Variable rdefs: routine_table.

Definition svalid_routine r (rspec: routine_spec) :=
  eresult_annotate ("Routine " ++ r ++ ": ")%string (
  match rdefs r with
    None => error "routine not implemented"%string
  | Some (RoutineDef xs c) =>
    let (xs', pre, post) := rspec in
    sresult_ok (
      sstate0 '_x">|>
      vs <- iter fresh_mut (length xs');
      swith_store sstore0 (
        s1 <- swith_store' (sstore_updates sstore0 xs' vs) (sproduce pre);
        t <- swith_store (sstore_updates sstore0 xs vs) (s_exec c; seval_mut result);
        message_mut "Checking postcondition";
        swith_store (sstore_update s1 result t) (sconsume post)
      );
      sleakcheck
    )
  end
  ).

Fixpoint foralle{A B}(f: A -> eresult B)(xs: list A): eresult B :=
  match xs with
    nil => ok
  | x::xs => f x '_x">&&> foralle f xs
  end.

Definition svalid_routines :=
  foralle
    (fun el => svalid_routine (fst el) (snd el))
    (StringMap.elements routine_specs).

Definition svalid_command c :=
  sresult_ok (
    sstate0 '_x">|> s_exec c
  ).

Definition pdefs_check :=
  match pred_defs pointsto_pred with
    None => ok
  | _ => error "|-> should not be user-defined"%string
  end '_x">&&>
  match pred_defs malloc_block_pred with
    None => ok
  | _ => error "mb should not be user-defined"%string
  end.

Definition svalid_program c :=
  pdefs_check '_x">&&>
  svalid_routines '_x">&&>
  svalid_command c.

End SRoutineDefs.

End SymbolicExecution.

Open Local Scope string_scope.

Local Notation l := "l".
Local Notation v := "v".
Local Notation next := "next".
Local Notation list := "list".
Local Notation mb := malloc_block_pred.
Local Notation n := "n".

Definition list_pred := "list".
Definition list_pred_def :=
  PredDef [l] (
    If l == 0 Then
      b_true
    Else (
      #mb(l, 2) &*& l '_'''_'''">|-> '_'''_'''">_ &*& '_'??'_x">(l + 1'_'??'_x">) '_'??'_x">|-> '_'??'_x">??next &*& #list(next)
    )
  ).

Local Notation a := "a".
Local Notation b := "b".
Local Notation tail := "tail".

Definition make_range_spec :=
  RoutineSpec [a, b]
    (b_true)
    (#list(result)).

Definition make_range :=
  RoutineDef [a, b] (
    If a == b Then (
      result ':=' 0
    ) Else (
      next ':=' call "make_range"(a + 1, b);
      result ':=' malloc(2); [result] ':=' a; [result + 1] ':=' next
    );
    close list(result)
  ).

Definition reverse_spec :=
  RoutineSpec [a, b]
    (#list(a) &*& #list(b))
    (#list(result)).

Definition reverse :=
  RoutineDef [a, b] (
    open list(a);
    If a == 0 Then
      result ':=' b
    Else (
      next ':=' [a + 1];
      [a + 1] ':=' b;
      close list(a);
      result ':=' call "reverse"(next, a)
    )
  ).

Definition dispose_spec :=
  RoutineSpec [a]
    (#list(a))
    (b_true).

Definition dispose :=
  RoutineDef [a] (
    open list(a);
    If a == 0 Then
      Skip
    Else (
      next ':=' [a + 1];
      free(a);
      call "dispose"(next)
    )
  ).

Definition reverse_loop_spec :=
  RoutineSpec [a]
    (#list(a))
    (#list(result)).

Definition reverse_loop :=
  RoutineDef [a] (
    close list(b);
    while ~ (a == 0) invariant #list(a) &*& #list(b) do (
      open list(a);
      tail ':=' [a + 1];
      [a + 1] ':=' b;
      b ':=' a;
      a ':=' tail;
      close list(b)
    );
    open list(a);
    result ':=' b
  ).

Definition dispose_loop :=
  RoutineDef [a] (
    while ~ (a == 0) invariant #list(a) do (
      open list(a);
      tail ':=' [a + 1];
      free(a);
      a ':=' tail
    );
    open list(a)
  ).

Definition my_rspecs :=
  StringMap.empty _
    '_x">|> StringMap.add "make_range" make_range_spec
    '_x">|> StringMap.add "reverse" reverse_spec
    '_x">|> StringMap.add "reverse_loop" reverse_loop_spec
    '_x">|> StringMap.add "dispose" dispose_spec
    '_x">|> StringMap.add "dispose_loop" dispose_spec
    .

Definition my_rdefs :=
  (fun r =>
     if string_dec r "make_range" then Some make_range else
     if string_dec r "reverse" then Some reverse else
     if string_dec r "reverse_loop" then Some reverse_loop else
     if string_dec r "dispose" then Some dispose else
     if string_dec r "dispose_loop" then Some dispose_loop else
     None).

Definition my_preddefs :=
  fun p =>
  if string_dec p list_pred then Some list_pred_def else
  None.

Definition reverse_test :=
  (
    a ':=' call "make_range"(0, 5);
    close list(0);
    b ':=' call "reverse"(a, 0);
    close list(0);
    a ':=' call "reverse"(b, 0);
    call "dispose"(a)
  )%cmd.

Goal
  svalid_command my_rspecs my_preddefs reverse_test = ok.

Goal
  svalid_routine my_rspecs my_preddefs my_rdefs "make_range" make_range_spec = ok.

Goal
  svalid_program my_rspecs my_preddefs my_rdefs reverse_test = ok.

Definition reverse_test_broken :=
  (
    a ':=' call "make_range"(0, 5);
    close list(0, 0);
    b ':=' call "reverse"(a, 0);
    close list(0, 0);
    call "reverse"(b + 10, 0)
  )%cmd.

Goal
  svalid_program my_rspecs my_preddefs my_rdefs reverse_test_broken '_x"><> ok.

Definition reverse_loop_test :=
  (
    a ':=' call "make_range"(0, 5);
    b ':=' call "reverse_loop"(a);
    call "reverse_loop"(b)
  )%cmd.

Goal
  svalid_program my_rspecs my_preddefs my_rdefs reverse_loop_test = ok.