Cell.lean

inductive Cell (m : Type (u + 1) → Type (u + 1)) [Monad m] (a b : Type (u + 1)) : Type (u+1) where
  | cell {s : Type u} : (s × (s → a → m (b × s))) → Cell m a b 
  | arrm : (a → m b) → Cell m a b


def cellCatId {m : Type _ -> Type _} [Monad m]: Cell m a a := Cell.arrm pure

def cellCatCompose  {m : Type (u+1) -> Type (u+1)} [Monad m] : Cell m b c → Cell m a b → Cell m a c
  | Cell.arrm f, Cell.arrm g => Cell.arrm $ Bind.kleisliLeft f g
  | Cell.arrm f, Cell.cell (s, cellStep) =>
    let f' := λ (b, s) => (λ c => (c, s)) <$> (f b)
    let cellStep' := λ state => Bind.kleisliLeft  f' (cellStep state)
    Cell.cell (s, cellStep')
  | Cell.cell (s, cellStep), Cell.arrm f =>
    let f' := λs => Bind.kleisliLeft (cellStep s) f
    Cell.cell (s, f')
  | Cell.cell (state2, step2), Cell.cell (state1, step1) => 
    let cellStep' := λ (state1, state2) a => do
        let (b, state1') ← step1 state1 a
        let (c, state2') ← step2 state2 b
        return (c, (state1', state2'))
    Cell.cell ((state1, state2), cellStep')

theorem cell_id_comp {m : Type (u+1) -> Type (u+1)} [Monad m] [LawfulMonad m] (cell : Cell m b c) :  cellCatCompose cellCatId cell = cell := by
    have h {α β : Type} : ∀x : Prod α β, (x.fst, x.snd) = x := by
      simp 
    cases cell with
      | arrm f => 
        simp [cellCatCompose, cellCatId, Bind.kleisliLeft]
      | cell p => 
        cases p with
          | mk s f => 
            simp [cellCatCompose, cellCatId]
            rw [h]