Require Import Computation.Monad.
Require Import Coq.Logic.JMeq.

Section Termination.

     (* an inductive predicate proving that a given computation terminates with a particular value *)  
  Inductive TerminatesWith (A:Set) : #A -> A -> Prop :=
  | TerminatesReturnWith :
    forall (a:A),
      TerminatesWith A (Return a) a
      
  | TerminatesBindWith :
    forall (B:Set) (b:B) (a:A) (f:B->#A) (c:#B),
      (TerminatesWith B c b)
      -> TerminatesWith A (f b) a
      -> TerminatesWith A (@Bind A B f c) a
      .
  
     (* an inductive predicate proving that a given computation terminates with /some/ value *)  
  Inductive Terminates (A:Set) : #A -> Prop :=
  | TerminatesReturn :
    forall (a:A),
      Terminates A (Return a)
      
  | TerminatesBind :
    forall (B:Set) (f:B->#A) (c:#B),
      Terminates B c
      -> (forall (b:B), (TerminatesWith B c b) -> Terminates A (f b))
      -> Terminates A (@Bind A B f c)
      .

  Lemma terminating_computations_must_produce_a_value
    : forall (A:Set) (c:#A), Terminates A c -> ex (fun a:A => TerminatesWith A c a).
    induction 1.
    
    exists a.
    constructor.
    
    elim IHTerminates.
    intros.
    clear IHTerminates.
    
    assert (exists a : A, TerminatesWith A (f x) a).
    apply H1.
    assumption.
    
    elim H3.
    intros.
    
    exists x0.
    apply (TerminatesBindWith A B x x0 f c).
    assumption.
    assumption.
  Defined.

  Theorem coerce : forall (A B:Set)(pf:A=B)(a:#A), #B.
    intros.
    subst.
    exact a.
  Defined.
  
  Inductive InvokedBy (A B C:Set) : #A -> #B -> #C -> Prop :=
  | invokesPrev : forall
    (z:#C)
    (c:#B)
    (f:B->#A),
    InvokedBy A B C (@Bind A B f c) c z
  | invokesFunc : forall
    (c:#C)
    (b:C)
    (f:C->#A)
    (pf:#A->#B)
    (eqpf:A=B)
    (_:JMeq (pf (f b)) (f b))
    (_:TerminatesWith C c b),
    InvokedBy A B C (@Bind A C f c) (pf (f b)) c.

  Inductive Safe : forall (A:Set) (c:#A), Prop :=
    Safe_intro :
    forall (A:Set) (c:#A),
      (forall (B C:Set) (c':#B)(z:#C), InvokedBy A B C c c' z -> Safe B c')
      -> Safe A c.

  Definition Safe_inv
    : forall (A B C:Set)(c:#A)(_:Safe A c)(c':#B)(z:#C)(_:InvokedBy A B C c c' z), Safe B c'.
    destruct 1.
    apply (H B).
  Defined.

  Notation "! c :: A" := {a:A|TerminatesWith A c a} (at level 5).
  Notation "'!Let' x := y 'in' z" := ((fun x => z)y)(at level 100).
  Definition eval' CC cc (Z:Set) (z:#Z) (s:Safe CC cc) : (!cc::CC).
    refine(
      !Let eval_one_step :=
        fun C c Z z =>
          match c return (forall PRED pred Z z, InvokedBy C PRED Z c pred z -> !pred::PRED) -> !c::C with
            | Return x => fun _ => exist _ x _
            | Bind CN f cn =>
              fun eval_pred =>
                match eval_pred CN cn Z z (invokesPrev C CN Z z cn f) with
                  | exist b pf =>
                    match eval_pred C (f b) CN cn _ with
                      | exist a' pf' => exist _ a' _
                    end
                end
          end
          in
          fix eval_all C c Z z (s:Safe C c) {struct s} : !c::C :=
          eval_one_step C c Z z (fun C' c' Z z icc => eval_all C' c' Z z (Safe_inv C C' Z c s c' z icc))
    ).
    
    constructor.

    refine (invokesFunc C C CN cn b f (fun x:#C=>x) _ _ _).
    auto.
    auto.
    assumption.
    
    apply (TerminatesBindWith C CN b a' f cn).
    assumption.
    assumption.
  Defined.

  Theorem jmeq_lemma : forall (A1 A2 B:Set)(c1:#A1)(c2:#A2)(f1:A1->#B)(f2:A2->#B),
    ((c1>>=f1)=(c2>>=f2))
    -> (JMeq c1 c2) /\ (JMeq f1 f2) /\ (A1=A2).
    intros.
    inversion H.
    split.
    dependent rewrite H3.
    simpl.
    auto.
    split.
    dependent rewrite H2.
    simpl.
    auto.
    auto.
  Qed.

  Theorem termination_is_safe : forall (A:Set) (c:#A) (t:Terminates A c), Safe A c.
    intros.
    induction t.
    apply Safe_intro.
    intros.
    inversion H.

    apply Safe_intro.
    intros.
    simple inversion H1.
    subst.

    apply jmeq_lemma in H2.
    destruct H2.
    destruct H3.
    subst.
    apply JMeq_eq in H2.
    subst.
    auto.

    intros.
    apply jmeq_lemma in H4.
    destruct H4.
    destruct H7.
    rewrite <- H5.
    clear H5.
    generalize H2.
    clear H2.
    subst.
    apply JMeq_eq in H7.
    apply JMeq_eq in H4.
    subst.
    intros.
    apply JMeq_eq in H2.
    rewrite H2.
    apply H0.
    auto.
  Defined.

End Termination.

Implicit Arguments Terminates [A].
Implicit Arguments TerminatesReturn [A].
Implicit Arguments TerminatesBind [A].
Implicit Arguments TerminatesReturnWith [A].
Implicit Arguments TerminatesBindWith [A].
Implicit Arguments eval' [CC].