(* Program logic (hoare logic + weakest preconditions) over
   Virtual Machine language. *)
module Compiler_logic

  use import int.Int
  use import list.List
  use import list.Length
  use import list.Append
  use import vm.Vm
  use import state.State

  function fst (p: ('a,'b)) : 'a = let (x,_) = p in x
  meta rewrite_def function fst

  function snd (p: ('a,'b)) : 'b = let (_,y) = p in y
  meta rewrite_def function snd

  predicate (-->) (x y:'a) = "rewrite" x = y
  meta rewrite_def predicate (-->)

  (* Unary predicates over machine states *)
  type pred  = machine_state -> bool

  (* Binary predicates over machine states *)
  type rel  = machine_state -> pred

  (* pre/post-conditions types, as parameterized unary/binary predicates.
     'a represents auxiliary variables
     pos is an auxiliary variable representing the absolute position at which
     the code is loaded. *)
  type pre 'a = 'a -> pos -> pred
  type post 'a = 'a -> pos -> rel

  (* Machine transition valid whatever the global code is. *)
  predicate contextual_irrelevance (c:code) (p:pos) (ms1 ms2:machine_state) =
    forall c_glob. codeseq_at c_glob p c -> transition_star c_glob ms1 ms2

  (* Hoare triples with explicit pre & post *)
  type hl 'a = { code: code; ghost pre : pre 'a; ghost post: post 'a }
    (* (Total) correctness for hoare triple. *)
    invariant { forall x:'a,p ms. pre x p ms ->
      exists ms'. post x p ms ms' /\ contextual_irrelevance code p ms ms' }
    by { code = Nil; pre = (fun _ _ _ -> false); post = fun _ _ _ _ -> true }

  (* Predicate transformer type. Same auxiliary variables as for
     Hoare triples. *)
  type wp_trans 'a = 'a -> pos -> pred -> pred

  (* Code with backward predicate transformer. *)
  type wp 'a = { wcode : code; ghost wp : wp_trans 'a }
    (* Similar invariant for backward predicate transformers *)
    invariant { forall x:'a,p post ms. wp x p post ms ->
      exists ms'. post ms' /\ contextual_irrelevance wcode p ms ms' }
    by { wcode = Nil; wp = fun _ _ _ _ -> false }

  (* WP combinator for sequence. Similar to the standard WP calculus
     for sequence. The initial machine state is memorized in auxiliary
     variables for potential use in the second code specification. *)
  function seq_wp
    (l1:int) (w1:wp_trans 'a) (w2:wp_trans ('a,machine_state)) : wp_trans 'a =
    fun x p q ms -> w1 x p (w2 (x,ms) (p+l1) q) ms

  lemma seq_wp_lemma : forall l1,w1: wp_trans 'a,w2 x p q ms.
      seq_wp l1 w1 w2 x p q ms = w1 x p (w2 (x,ms) (p+l1) q) ms

  meta rewrite lemma seq_wp_lemma

  (* Code combinator for sequence, with wp. *)
  let (--) (s1 : wp 'a) (s2 : wp ('a, machine_state)) : wp 'a
    ensures  { result.wcode.length --> s1.wcode.length + s2.wcode.length }
    ensures  { result.wp --> seq_wp s1.wcode.length s1.wp s2.wp }
  = let code = s1.wcode ++ s2.wcode in
    let res = { wcode = code; wp = seq_wp s1.wcode.length s1.wp s2.wp } in
    assert { forall x: 'a, p post ms. res.wp x p post ms ->
      not (exists ms'. post ms' /\ contextual_irrelevance res.wcode p ms ms') ->
        (forall ms'. s2.wp (x,ms) (p + s1.wcode.length) post ms' /\
           contextual_irrelevance res.wcode p ms ms' -> false) && false };
    res

  function fork_wp (w:wp_trans 'a) (cond:pre 'a) : wp_trans 'a =
    fun x p q ms -> if cond x p ms then w x p q ms else q ms

  lemma fork_wp_lemma: forall w:wp_trans 'a,cond x p q ms.
    fork_wp w cond x p q ms =
      ((not cond x p ms -> q ms) /\ (cond x p ms -> w x p q ms))

  meta rewrite lemma fork_wp_lemma

  (* Code combinator for conditional execution.
     Similar to WP calculus for (if cond then s). *)

  let (%) (s:wp 'a) (ghost cond:pre 'a) : wp 'a
    ensures  { result.wp --> fork_wp s.wp cond }
    ensures  { result.wcode.length --> s.wcode.length }
  = { wcode = s.wcode; wp = fork_wp s.wp cond }

  (* WP transformer for hoare triples. *)
  function towp_wp (pr:pre 'a) (ps:post 'a) : wp_trans 'a =
    fun x p q ms -> pr x p ms && (forall ms'. ps x p ms ms' -> q ms')

  lemma towp_wp_lemma:
    forall pr ps, x:'a, p q ms. towp_wp pr ps x p q ms =
      (pr x p ms && (forall ms'. ps x p ms ms' -> q ms'))
  meta rewrite lemma towp_wp_lemma

  (* Unwrap code with hoare triple into code with wp.
     Analogous to procedure call/abstract block. *)
  let ($_) (c:hl 'a) : wp 'a
    ensures  { result.wcode.length --> c.code.length }
    ensures  { result.wp --> towp_wp c.pre c.post }
  = { wcode = c.code; wp = towp_wp c.pre c.post }

  (* Equip code with pre/post-condition. That is here that proof happen.
     (P -> wp (c,Q)). Anologous to checking function/abstract block
     specification. *)
  let hoare (ghost pre:pre 'a) (c:wp 'a) (ghost post:post 'a) : hl 'a
    requires { forall x p ms. pre x p ms -> (c.wp x p (post x p ms)) ms }
    ensures { result.pre --> pre }
    ensures { result.post --> post }
    ensures  { result.code.length --> c.wcode.length }
  = { code = c.wcode ; pre = pre; post = post }

  function trivial_pre : pre 'a = fun _ p ms -> let VMS p' _ _ = ms in p = p'
  meta rewrite_def function trivial_pre

  (* Accessibility predicate. *)
  inductive acc ('a -> 'a -> bool) 'a =
    | Acc : forall r, x:'a. (forall y. r y x -> acc r y) -> acc r x

  (* Utility: some flavor of conjonction. *)
  function pconj (p1:pred) (x:machine_state)
                 (p2:machine_state -> pred) : pred =
                 fun y -> p1 y && p2 y x
  lemma pconj_lemma : forall p1 x p2 y. pconj p1 x p2 y <-> p1 y && p2 y x
  meta rewrite lemma pconj_lemma

  (* WP combinator for looping construction. Similar to weakest precondition
     for while loops. *)

  function loop_wp (w:wp_trans 'a) (inv cont:pre 'a)
                                   (var:post 'a) : wp_trans 'a =
    fun x p q ms -> inv x p ms && acc (var x p) ms && forall ms'. inv x p ms' ->
      if cont x p ms'
      then w x p (pconj (inv x p) ms' (var x p)) ms'
      else w x p q ms'

  lemma loop_wp_lemma : forall w:wp_trans 'a,inv cont var x p q ms.
    loop_wp w inv cont var x p q ms <->
      inv x p ms && acc (var x p) ms && forall ms'. inv x p ms' ->
        (cont x p ms' -> w x p (pconj (inv x p) ms' (var x p)) ms')
        /\ (not cont x p ms' -> w x p q ms')

  meta rewrite lemma loop_wp_lemma

  (* Code combinator for looping construct. *)
  let make_loop (c:wp 'a) (ghost inv cont:pre 'a)
                          (ghost var:post 'a) : wp 'a
    ensures { result.wp --> loop_wp c.wp inv cont var }
    ensures { result.wcode.length --> c.wcode.length }
  = let ghost wpt = loop_wp c.wp inv cont var in
    assert { forall x p q ms0. wpt x p q ms0 ->
      forall ms. inv x p ms -> acc (var x p) ms ->
        exists ms'. contextual_irrelevance c.wcode p ms ms' /\ q ms'
    };
    { wcode = c.wcode; wp = wpt }

end