theory RawFunctor := sig
  f : Type -> Type
  fmap : (a : Type) -> (b : Type) -> (a -> b) -> f a -> f b
end

theory HFunctor := sig
  f : (Type -> Type) -> Type -> Type
  hfmap : (m : RawFunctor) -> (a : Type) -> (b : Type) -> (a -> b) -> f m.f a -> f m.f b
  hmap : (m : RawFunctor) -> (n : RawFunctor) -> ((x : Type) -> m.f x -> n.f x) -> (a : Type) -> f m.f a -> f n.f a
end

theory Handler (h : HFunctor) := sig
  M : Type -> Type
  ret : (a : Type) -> a -> M a
  bind : (a : Type) -> (b : Type) -> M a -> (a -> M b) -> M b
  malg : (a : Type) -> h.f M a -> M a
end

def Co (h : HFunctor) (a : Type) : Type := (hd : Handler h) -> (r : Type) -> (a -> hd.M r) -> hd.M r

def val (h : HFunctor) (a : Type) (x : a) : Co h a := hd => r => k => k x

def bind (h : HFunctor) (a : Type) (b : Type) (c : Co h a) (f : a -> Co h b) : Co h b := hd => r => k => c hd r (x => f x hd r k)

def coFunctor (h : HFunctor) : RawFunctor := struct
  f := a => Co h a
  fmap := a => b => g => c => bind h a b c (x => val h b (g x))
end

def handlerFunctor (h : HFunctor) (hd : Handler h) : RawFunctor := struct
  f := a => hd.M a
  fmap := a => b => g => m => hd.bind a b m (x => hd.ret b (g x))
end

def handle (h : HFunctor) (hd : Handler h) (a : Type) (c : Co h a) : hd.M a := c hd a (x => hd.ret a x)

def op (h : HFunctor) (a : Type) (b : Type) (p : h.f (Co h) a) (k : a -> Co h b) : Co h b := hd => r => k2 => (
  hd.bind a r
    (hd.malg a (h.hmap (coFunctor h) (handlerFunctor h hd) (x => c => handle h hd x c) a p))
    (x => k x hd r k2))

def Sum (a : Type) (b : Type) : Type := sum
  'left a
  'right b
end

def plus (h1 : HFunctor) (h2 : HFunctor) : HFunctor := struct
  f := m => a => Sum (h1.f m a) (h2.f m a)
  hfmap := m => a => b => g => p => match p
    'left x => 'left (h1.hfmap m a b g x)
    'right y => 'right (h2.hfmap m a b g y)
  end
  hmap := m => n => nt => a => p => match p
    'left x => 'left (h1.hmap m n nt a x)
    'right y => 'right (h2.hmap m n nt a y)
  end
end

def Sub (h : HFunctor) (e : HFunctor) : Type := (m : Type -> Type) -> (x : Type) -> h.f m x -> e.f m x

def inl (h : HFunctor) (k : HFunctor) : Sub h (plus h k) := m => x => p => 'left p

def inr (h : HFunctor) (k : HFunctor) : Sub k (plus h k) := m => x => p => 'right p

def through (h : HFunctor) (k : HFunctor) (inj : Sub h k) (hd : Handler k) : Handler h := struct
  M := hd.M
  ret := hd.ret
  bind := hd.bind
  malg := x => p => hd.malg x (inj hd.M x p)
end

def lift (h : HFunctor) (k : HFunctor) (a : Type) (inj : Sub h k) (c : Co h a) : Co k a := hd => r => kont => c (through h k inj hd) r kont

def Unit : Type := sum
  'tt
end

def Pair (a : Type) (b : Type) : Type := sum
  'pair a b
end

def StateF (s : Type) : HFunctor := struct
  f := m => a => sum
    'get (s -> a)
    'put s a
  end
  hfmap := m => a => b => g => p => match p
    'get k => 'get (st => g (k st))
    'put st v => 'put st (g v)
  end
  hmap := m => n => t => a => p => match p
    'get k => 'get k
    'put st v => 'put st v
  end
end

def get (s : Type) : Co (StateF s) s := op (StateF s) s s ('get (st => st)) (x => val (StateF s) s x)

def put (s : Type) (st : s) : Co (StateF s) Unit := op (StateF s) Unit Unit ('put st 'tt) (x => val (StateF s) Unit x)

def ExnF (e : Type) : HFunctor := struct
  f := m => a => sum
    'throw e
    'catch (m a) (e -> m a)
  end
  hfmap := m => a => b => g => p => match p
    'throw err => 'throw err
    'catch c h => 'catch (m.fmap a b g c) (err => m.fmap a b g (h err))
  end
  hmap := m => n => t => a => p => match p
    'throw err => 'throw err
    'catch c h => 'catch (t a c) (err => t a (h err))
  end
end

def throw (e : Type) (a : Type) (err : e) : Co (ExnF e) a := op (ExnF e) a a ('throw err) (x => val (ExnF e) a x)

def catch (e : Type) (a : Type) (c : Co (ExnF e) a) (h : e -> Co (ExnF e) a) : Co (ExnF e) a := op (ExnF e) a a ('catch c h) (x => val (ExnF e) a x)

def StExn (s : Type) (e : Type) : Handler (plus (StateF s) (ExnF e)) := struct
  M := a => s -> Sum e (Pair a s)
  ret := a => x => st => 'right ('pair x st)
  bind := a => b => mx => f => st => match mx st
    'right p => match p
      'pair x st1 => f x st1
    end
    'left err => 'left err
  end
  malg := a => p => match p
    'left sop => match sop
      'get k => st => 'right ('pair (k st) st)
      'put st1 v => st => 'right ('pair v st1)
    end
    'right eop => match eop
      'throw err => st => 'left err
      'catch mx h => st => match mx st
        'right q => 'right q
        'left err => h err st
      end
    end
  end
end

def ExnSt (s : Type) (e : Type) : Handler (plus (StateF s) (ExnF e)) := struct
  M := a => s -> Pair (Sum e a) s
  ret := a => x => st => 'pair ('right x) st
  bind := a => b => mx => f => st => match mx st
    'pair v st1 => match v
      'right x => f x st1
      'left err => 'pair ('left err) st1
    end
  end
  malg := a => p => match p
    'left sop => match sop
      'get k => st => 'pair ('right (k st)) st
      'put st1 v => st => 'pair ('right v) st1
    end
    'right eop => match eop
      'throw err => st => 'pair ('left err) st
      'catch mx h => st => match mx st
        'pair v st1 => match v
          'right x => 'pair ('right x) st1
          'left err => h err st1
        end
      end
    end
  end
end

def Flag : Type := sum
  'off
  'on
end

def Err : Type := sum
  'boom
end

def StF : HFunctor := StateF Flag
def ExF : HFunctor := ExnF Err
def St : HFunctor := plus StF ExF

def getF : Co St Flag := lift StF St Flag (inl StF ExF) (get Flag)

def putF (st : Flag) : Co St Unit := lift StF St Unit (inl StF ExF) (put Flag st)

def throwF (a : Type) (err : Err) : Co St a := lift ExF St a (inr StF ExF) (throw Err a err)

def catchF (a : Type) (c : Co St a) (h : Err -> Co St a) : Co St a := op St a a ('right ('catch c h)) (x => val St a x)

def prog : Co St Flag := (
  catchF Flag
    (bind St Unit Flag (putF 'on) (u => throwF Flag 'boom))
    (err => getF))

def testStExn : Sum Err (Pair Flag Flag) := handle St (StExn Flag Err) Flag prog 'off

def testExnSt : Pair (Sum Err Flag) Flag := handle St (ExnSt Flag Err) Flag prog 'off