From Tealeaves Require Export
  Classes.Kleisli.Monad
  Functors.Compose.


(** * State Monad *)
(**********************************************************************)
Inductive State (S A: Type) :=
| mkState: (S -> S * A) -> State S A.

Arguments mkState {S A}%type_scope _.

Section state_monad.

  Context
    {S: Type}.

  Definition runState {A}: State S A -> S -> S * A :=
    fun st s => match st with
             | mkState run => run s
             end.

  Definition runStateValue {A}: State S A -> S -> A :=
    fun st s => snd (runState st s).

  Definition runStateState {A}: State S A -> S -> S :=
    fun st s => fst (runState st s).

  (** ** Functor Instance *)
  (********************************************************************)
  #[export] Instance Map_State: Map (State S) :=
    fun A B (f: A -> B) (st: State S A) =>
      match st with
      | mkState r =>
          mkState (fun s => match (r s) with (s', a) => (s', f a) end)
      end.

  #[export] Instance Functor_State: Functor (State S).
S: Type
Functor (State S)
  Proof.
S: Type
Functor (State S)
    constructor; try typeclasses eauto.
S: Type
forall A : Type, map id = id
S: Type
forall (A B C : Type) (f : A -> B) (g : B -> C),
map g ∘ map f = map (g ∘ f)
    -
S: Type
forall A : Type, map id = id intros.
S, A: Type
map id = id ext [st].
S, A: Type
st: S -> S * A
map id (mkState st) = id (mkState st) unfold id.
S, A: Type
st: S -> S * A
map (fun x : A => x) (mkState st) = mkState st
      cbn.
S, A: Type
st: S -> S * A
mkState (fun s : S => let (s', a) := st s in (s', a)) =
mkState st fequal.
S, A: Type
st: S -> S * A
(fun s : S => let (s', a) := st s in (s', a)) = st ext s.
S, A: Type
st: S -> S * A
s: S
(let (s', a) := st s in (s', a)) = st s now destruct (st s).
    -
S: Type
forall (A B C : Type) (f : A -> B) (g : B -> C),
map g ∘ map f = map (g ∘ f) intros.
S, A, B, C: Type
f: A -> B
g: B -> C
map g ∘ map f = map (g ∘ f) ext [st].
S, A, B, C: Type
f: A -> B
g: B -> C
st: S -> S * A
(map g ∘ map f) (mkState st) =
map (g ∘ f) (mkState st) unfold compose.
S, A, B, C: Type
f: A -> B
g: B -> C
st: S -> S * A
map g (map f (mkState st)) = map (g ○ f) (mkState st)
      cbn.
S, A, B, C: Type
f: A -> B
g: B -> C
st: S -> S * A
mkState
  (fun s : S =>
   let (s', a) := let (s', a) := st s in (s', f a) in
   (s', g a)) =
mkState
  (fun s : S => let (s', a) := st s in (s', g (f a))) fequal.
S, A, B, C: Type
f: A -> B
g: B -> C
st: S -> S * A
(fun s : S =>
 let (s', a) := let (s', a) := st s in (s', f a) in
 (s', g a)) =
(fun s : S => let (s', a) := st s in (s', g (f a))) ext s.
S, A, B, C: Type
f: A -> B
g: B -> C
st: S -> S * A
s: S
(let (s', a) := let (s', a) := st s in (s', f a) in
 (s', g a)) = (let (s', a) := st s in (s', g (f a))) now destruct (st s).
  Qed.

  (** ** Monad Instance (Kleisli) *)
  (********************************************************************)
  #[export] Instance Return_State: Return (State S) :=
    fun A (a: A) => mkState (fun s => (s, a)).

  #[export] Instance Bind_State: Bind (State S) (State S) :=
    fun A B (f: A -> State S B) (st1: State S A) =>
      match st1 with
      | mkState action =>
          mkState
            (fun s => match (action s) with
                     (s', a) => runState (f a) s'
                   end)
      end.

  #[export] Instance Monad_State: Monad (State S).
S: Type
Monad (State S)
  Proof.
S: Type
Monad (State S)
    constructor.
S: Type
forall (A B : Type) (f : A -> State S B),
bind f ∘ ret = f
S: Type
RightPreModule (State S) (State S)
    -
S: Type
forall (A B : Type) (f : A -> State S B),
bind f ∘ ret = f intros.
S, A, B: Type
f: A -> State S B
bind f ∘ ret = f ext a.
S, A, B: Type
f: A -> State S B
a: A
(bind f ∘ ret) a = f a cbv.
S, A, B: Type
f: A -> State S B
a: A
mkState
  (fun s : S =>
   match f a with
   | mkState run => run s
   end) = f a
      now destruct (f a).
    -
S: Type
RightPreModule (State S) (State S) constructor.
S: Type
forall A : Type, bind ret = id
S: Type
forall (A B C : Type) (g : B -> State S C)
  (f : A -> State S B),
bind g ∘ bind f = bind (kc g f)
      +
S: Type
forall A : Type, bind ret = id intros.
S, A: Type
bind ret = id ext s.
S, A: Type
s: State S A
bind ret s = id s
        cbv.
S, A: Type
s: State S A
match s with
| mkState action =>
    mkState
      (fun s : S => let (s', a) := action s in (s', a))
end = s destruct s.
S, A: Type
p: S -> S * A
mkState (fun s : S => let (s', a) := p s in (s', a)) =
mkState p
        fequal.
S, A: Type
p: S -> S * A
(fun s : S => let (s', a) := p s in (s', a)) = p ext s.
S, A: Type
p: S -> S * A
s: S
(let (s', a) := p s in (s', a)) = p s now destruct (p s).
      +
S: Type
forall (A B C : Type) (g : B -> State S C)
  (f : A -> State S B),
bind g ∘ bind f = bind (kc g f) intros.
S, A, B, C: Type
g: B -> State S C
f: A -> State S B
bind g ∘ bind f = bind (kc g f) ext s.
S, A, B, C: Type
g: B -> State S C
f: A -> State S B
s: State S A
(bind g ∘ bind f) s = bind (kc g f) s
        cbv.
S, A, B, C: Type
g: B -> State S C
f: A -> State S B
s: State S A
match
  match s with
  | mkState action =>
      mkState
        (fun s : S =>
         let (s', a) := action s in
         match f a with
         | mkState run => run s'
         end)
  end
with
| mkState action =>
    mkState
      (fun s : S =>
       let (s', a) := action s in
       match g a with
       | mkState run => run s'
       end)
end =
match s with
| mkState action =>
    mkState
      (fun s : S =>
       let (s', a) := action s in
       match
         match f a with
         | mkState action0 =>
             mkState
               (fun s0 : S =>
                let (s'0, a0) := action0 s0 in
                match g a0 with
                | mkState run => run s'0
                end)
         end
       with
       | mkState run => run s'
       end)
end destruct s.
S, A, B, C: Type
g: B -> State S C
f: A -> State S B
p: S -> S * A
mkState
  (fun s : S =>
   let (s', a) :=
     let (s', a) := p s in
     match f a with
     | mkState run => run s'
     end in
   match g a with
   | mkState run => run s'
   end) =
mkState
  (fun s : S =>
   let (s', a) := p s in
   match
     match f a with
     | mkState action =>
         mkState
           (fun s0 : S =>
            let (s'0, a0) := action s0 in
            match g a0 with
            | mkState run => run s'0
            end)
     end
   with
   | mkState run => run s'
   end) fequal.
S, A, B, C: Type
g: B -> State S C
f: A -> State S B
p: S -> S * A
(fun s : S =>
 let (s', a) :=
   let (s', a) := p s in
   match f a with
   | mkState run => run s'
   end in
 match g a with
 | mkState run => run s'
 end) =
(fun s : S =>
 let (s', a) := p s in
 match
   match f a with
   | mkState action =>
       mkState
         (fun s0 : S =>
          let (s'0, a0) := action s0 in
          match g a0 with
          | mkState run => run s'0
          end)
   end
 with
 | mkState run => run s'
 end)
        ext s.
S, A, B, C: Type
g: B -> State S C
f: A -> State S B
p: S -> S * A
s: S
(let (s', a) :=
   let (s', a) := p s in
   match f a with
   | mkState run => run s'
   end in
 match g a with
 | mkState run => run s'
 end) =
(let (s', a) := p s in
 match
   match f a with
   | mkState action =>
       mkState
         (fun s : S =>
          let (s'0, a0) := action s in
          match g a0 with
          | mkState run => run s'0
          end)
   end
 with
 | mkState run => run s'
 end) destruct (p s).
S, A, B, C: Type
g: B -> State S C
f: A -> State S B
p: S -> S * A
s, s0: S
a: A
(let (s', a) :=
   match f a with
   | mkState run => run s0
   end in
 match g a with
 | mkState run => run s'
 end) =
match
  match f a with
  | mkState action =>
      mkState
        (fun s : S =>
         let (s', a) := action s in
         match g a with
         | mkState run => run s'
         end)
  end
with
| mkState run => run s0
end
        now destruct (f a).
  Qed.

  (** ** Monad Instance (Categorical) *)
  (********************************************************************)
  #[local] Instance Join_State: Monad.Join (State S) :=
    fun A (st: State S (State S A)) =>
      match st with
      | mkState r =>
          mkState (fun s => match (r s) with (s', st') => runState st' s' end)
      end.

  Instance Natural_ret: Natural (@ret (State S) _).
S: Type
Natural (@ret (State S) Return_State)
  Proof.
S: Type
Natural (@ret (State S) Return_State)
    ltac:(constructor; try typeclasses eauto).
S: Type
forall (A B : Type) (f : A -> B),
map f ∘ ret = ret ∘ map f
    reflexivity.
  Qed.

  Instance Natural_join: Natural (@Monad.join (State S) _).
S: Type
Natural (@Monad.join (State S) Join_State)
  Proof.
S: Type
Natural (@Monad.join (State S) Join_State)
    ltac:(constructor; try typeclasses eauto).
S: Type
forall (A B : Type) (f : A -> B),
map f ∘ Monad.join = Monad.join ∘ map f
    intros.
S, A, B: Type
f: A -> B
map f ∘ Monad.join = Monad.join ∘ map f ext [st].
S, A, B: Type
f: A -> B
st: S -> S * State S A
(map f ∘ Monad.join) (mkState st) =
(Monad.join ∘ map f) (mkState st) cbn.
S, A, B: Type
f: A -> B
st: S -> S * State S A
mkState
  (fun s : S =>
   let (s', a) :=
     let (s', st') := st s in runState st' s' in
   (s', f a)) =
mkState
  (fun s : S =>
   let (s', st') :=
     let (s', a) := st s in (s', map f a) in
   runState st' s') fequal.
S, A, B: Type
f: A -> B
st: S -> S * State S A
(fun s : S =>
 let (s', a) :=
   let (s', st') := st s in runState st' s' in
 (s', f a)) =
(fun s : S =>
 let (s', st') :=
   let (s', a) := st s in (s', map f a) in
 runState st' s')
    ext s.
S, A, B: Type
f: A -> B
st: S -> S * State S A
s: S
(let (s', a) :=
   let (s', st') := st s in runState st' s' in
 (s', f a)) =
(let (s', st') :=
   let (s', a) := st s in (s', map f a) in
 runState st' s') destruct (st s); cbn.
S, A, B: Type
f: A -> B
st: S -> S * State S A
s, s0: S
s1: State S A
(let (s', a) := runState s1 s0 in (s', f a)) =
runState (map f s1) s0 now (destruct s1).
  Qed.

  #[export] Instance Categorical_Monad_State: Categorical.Monad.Monad (State S).
S: Type
Categorical.Monad.Monad (State S)
  Proof.
S: Type
Categorical.Monad.Monad (State S)
    constructor; try typeclasses eauto; unfold compose.
S: Type
forall A : Type, Monad.join ○ ret = id
S: Type
forall A : Type, Monad.join ○ map ret = id
S: Type
forall A : Type,
Monad.join ○ Monad.join = Monad.join ○ map Monad.join
    -
S: Type
forall A : Type, Monad.join ○ ret = id intros.
S, A: Type
Monad.join ○ ret = id
      ext [st].
S, A: Type
st: S -> S * A
Monad.join (ret (mkState st)) = id (mkState st)
      reflexivity.
    -
S: Type
forall A : Type, Monad.join ○ map ret = id intros.
S, A: Type
Monad.join ○ map ret = id ext [st].
S, A: Type
st: S -> S * A
Monad.join (map ret (mkState st)) = id (mkState st) unfold id.
S, A: Type
st: S -> S * A
Monad.join (map ret (mkState st)) = mkState st
      cbn.
S, A: Type
st: S -> S * A
mkState
  (fun s : S =>
   let (s', st') :=
     let (s', a) := st s in (s', ret a) in
   runState st' s') = mkState st fequal.
S, A: Type
st: S -> S * A
(fun s : S =>
 let (s', st') :=
   let (s', a) := st s in (s', ret a) in
 runState st' s') = st ext s.
S, A: Type
st: S -> S * A
s: S
(let (s', st') :=
   let (s', a) := st s in (s', ret a) in
 runState st' s') = st s now destruct (st s).
    -
S: Type
forall A : Type,
Monad.join ○ Monad.join = Monad.join ○ map Monad.join intros.
S, A: Type
Monad.join ○ Monad.join = Monad.join ○ map Monad.join ext [st].
S, A: Type
st: S -> S * State S (State S A)
Monad.join (Monad.join (mkState st)) =
Monad.join (map Monad.join (mkState st)) unfold compose; cbn.
S, A: Type
st: S -> S * State S (State S A)
mkState
  (fun s : S =>
   let (s', st') :=
     let (s', st') := st s in runState st' s' in
   runState st' s') =
mkState
  (fun s : S =>
   let (s', st') :=
     let (s', a) := st s in (s', Monad.join a) in
   runState st' s')
      fequal.
S, A: Type
st: S -> S * State S (State S A)
(fun s : S =>
 let (s', st') :=
   let (s', st') := st s in runState st' s' in
 runState st' s') =
(fun s : S =>
 let (s', st') :=
   let (s', a) := st s in (s', Monad.join a) in
 runState st' s') ext s.
S, A: Type
st: S -> S * State S (State S A)
s: S
(let (s', st') :=
   let (s', st') := st s in runState st' s' in
 runState st' s') =
(let (s', st') :=
   let (s', a) := st s in (s', Monad.join a) in
 runState st' s') destruct (st s).
S, A: Type
st: S -> S * State S (State S A)
s, s0: S
s1: State S (State S A)
(let (s', st') := runState s1 s0 in runState st' s') =
runState (Monad.join s1) s0 now (destruct s1).
  Qed.

End state_monad.