Arrow.lean

class Arrow (A : Type u → Type u → Type w) where
  comp (outer : A β γ) (inner : A α β) : A α γ
  ofFunction : (α → β) → A α β
  mapFirst : A α β → A (α × γ) (β × γ)

instance : Arrow (· → ·) where
  comp := Function.comp
  ofFunction := id
  mapFirst arrow | (a, g) => (arrow a, g)

instance [Monad m] : Arrow (· → m ·) where
  comp outer inner := (inner · >>= outer)
  ofFunction f := pure ∘ f
  mapFirst arrow | (a, g) => do pure (← arrow a, g)

def Arrow.id [Arrow A] : A α α := ofFunction _root_.id

def Arrow.mapSecond {α β γ : Type u} [Arrow A] (arrow : A α β) : A (γ × α) (γ × β) := by
  apply comp (ofFunction Prod.swap)
  apply comp _ (ofFunction Prod.swap)
  exact mapFirst arrow

def Arrow.mapBoth {α β γ δ : Type u} [Arrow A] (a1 : A α γ) (a2 : A β δ) : A (α × β) (γ × δ)
  := comp (mapSecond a2) (mapFirst a1)

def Arrow.pair [Arrow A] (a1 : A α β) (a2 : A α γ) : A α (β × γ) := by
  apply comp
  apply mapBoth a1 a2
  apply ofFunction
  intro a
  exact (a, a)

class ArrowChoice (A) extends Arrow A where
  mapOk {α β γ : Type u} : A α β → A (Except γ α) (Except γ β)

section
open Arrow

def ArrowChoice.mapError {α β γ : Type u} [ArrowChoice A] (arrow : A α β) : A (Except α γ) (Except β γ) := by
  let exceptSwap {ε α'} : Except ε α' → Except α' ε
    | Except.ok a => Except.error a
    | Except.error e => Except.ok e
  apply comp (ofFunction exceptSwap)
  apply comp _ (ofFunction exceptSwap)
  exact mapOk arrow

def ArrowChoice.mapOkError {α β γ δ : Type u} [ArrowChoice A] (arrowOk : A α γ) (arrowError : A β δ) : A (Except β α) (Except δ γ) :=
  comp (mapOk arrowOk) (mapError arrowError)

def ArrowChoice.mapOrElse {α β γ : Type u} [ArrowChoice A] (arrowOk : A α γ) (arrowError : A β γ) : A (Except β α) γ := by
  apply comp _ (mapOkError arrowOk arrowError)
  apply ofFunction
  intro ex
  cases ex <;> assumption

end