defun.v

Require Import List PeanoNat.
Import ListNotations.
Set Universe Polymorphism.

Set Implicit Arguments.

Inductive t (eff : Type -> Type) (a : Type) :=
| Pure : a -> t eff a
| Eff : forall T, eff T -> (T -> t eff a) -> t eff a.

Inductive effect : Type -> Type :=
| get : unit -> effect nat
| put : nat -> effect unit.


(*
n <- get;
put (1 + n);
*)

Definition step1 (s : state1) : { t & (effect t * (t -> option state1))%type }:=
  match s with
  | BeforeGet => existT _ nat (get tt , fun n => Some (BeforePut (1+n)))
  | BeforePut n => existT _ unit (put n, fun _ => None)
  end.


(*
Run the original program and the corresponding state machine in parallel,
and check if they behave equivalently.
*)
Fixpoint check (state : Set)
         (step : state -> {ty & (effect ty * (ty -> option state))%type })
         (current : state)
         (a : Set)
         (prog : t effect a)
         (input : list nat) :=
  let (_, ef) := step current in
  let (es,fs) := ef in
  match prog with
  | Pure _ x => None
  | Eff e f =>
    match e in effect ty, es in effect tys return ((ty -> t effect a) -> (tys -> option state) -> option (list nat)) with
    | get _,get _ =>
      match input with
      | [] => fun _ _ => Some []
      | x::rest => fun f' f's =>
                     match f's x with
                     | None =>
                       match f' x with
                       | Pure _ _ => Some []
                       | _ => None
                       end
                     | Some next =>
                       check step next (f' x) rest
                     end
      end
    | put x,put xs =>
      fun f' f's =>
        if Nat.eq_dec x xs then
          match f's tt with
          | None =>
            match f' tt with
            | Pure _ _ => Some [x]
            | _ => None
            end
          | Some next =>
            option_map (cons x) (check step next (f' tt) input)
          end
        else
          None
    | _,_ => fun _ _ => None
    end
      f fs
  end.

Eval compute in check step1 BeforeGet ex1 [2].

Lemma ex1_ok : forall input, exists output,
      check step1 BeforeGet ex1 input = Some output.
Proof.
  unfold ex1. simpl.
  destruct input.
  - eauto.
  - eexists. simpl.
    destruct (Nat.eq_dec n n); congruence || eauto.
Qed.

(*
n <_ get;
m <- get;
put (n + m);
*)

Definition ex2 : t effect unit :=
  Eff (get tt) (fun n => Eff (get tt) (fun m => Eff (put (n + m)) (fun _ => Pure _ tt))).

Inductive state2 :=
| BeforeGet2_1 : state2
| BeforeGet2_2 : nat -> state2
| BeforePut2 : nat -> nat -> state2.

Definition step2 (s : state2) : { t & (effect t * (t -> option state2))%type } :=
  match s with
  | BeforeGet2_1 =>
    existT _ nat (get tt, fun n => Some (BeforeGet2_2 n))
  | BeforeGet2_2 n =>
    existT _ nat (get tt, fun m => Some (BeforePut2 n m))
  | BeforePut2 n m =>
    existT _ unit (put (n+m), fun _ => None)
  end.

Eval compute in check step2 BeforeGet2_1 ex2 [2;3].

Lemma ex2_ok : forall input, exists output,
      check step2 BeforeGet2_1 ex2 input = Some output.
Proof.
  intros. unfold ex2. simpl.
  destruct input.
  - eauto.
  - simpl. destruct input.
    eauto.
    simpl. eexists.
    destruct (Nat.eq_dec (n + n0) (n + n0)); congruence || eauto.
Qed.
  
(*
n <- get;
a = 0;
for 0..(n-1)
  x <- get;
  a = a + x;
put a;
 *)

Fixpoint ex3' (a n : nat) : t effect unit :=
  match n with
  | O => Eff (put a) (fun _ => Pure _ tt)
  | S n' => Eff (get tt) (fun x => ex3' (a + x) n')
  end.

Definition ex3 : t effect unit :=
  Eff (get tt) (fun n => ex3' 0 n).

Inductive state3 :=
| BeforeGet3 : state3
| BeforeGet3' : nat -> nat -> state3.

Definition step3 (s : state3) : {t & (effect t * (t -> option state3))%type } :=
  match s with
  | BeforeGet3 =>
    existT _ nat (get tt, fun n => Some (BeforeGet3' n 0))
  | BeforeGet3' n a =>
    match n with
    | O => existT _ unit (put a, fun _ => None)
    | S n' => existT _ nat (get tt, fun x => Some (BeforeGet3' n' (a + x)))
    end
  end.

Eval compute in check step3 BeforeGet3 ex3 [3;1;2;3].

Lemma ex3'_ok : forall n a input, exists output,
      check step3 (BeforeGet3' n a) (ex3' a n) input = Some output.
Proof.
  induction n; simpl; intros.
  - destruct (Nat.eq_dec a a); congruence || eauto.
  - destruct input.
    + eauto.
    + destruct (IHn (a + n0) input).
      rewrite H.
      eauto.
Qed.

Lemma ex3_ok : forall input, exists output,
      check step3 BeforeGet3 ex3 input = Some output.
Proof.
  unfold ex3.
  intros. simpl.
  destruct input.
  - eauto.
  - apply ex3'_ok.
Qed.