{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-}
module Control.Arrow.Transformer.Cont where

import Control.Arrow
import Control.Arrow.Transformer (ArrowTransformer(), lift)

newtype ContArrow e a b c = CO (a c e -> a b e)

instance ArrowApply a => Arrow (ContArrow e a) where
 arr f = liftCont (arr f)
 CO f >>> CO g = CO (\k -> f (g k))
 first (CO f) = CO $ \k -> app <<^ \(b,d) ->
  (f (arr (\c -> (c,d)) >>> k), b)

instance ArrowApply a => ArrowChoice (ContArrow e a) where
 left (CO f) = CO $ \k -> app <<^ \x ->
  case x of
   Left  b -> ((\(Left a) -> a) ^>> f (arr (\c -> Left c) >>> k), Left b)
   Right d -> ((\(Right a) -> a) ^>> id (arr (\c -> Right c) >>> k), Right d)

instance ArrowApply a => ArrowApply (ContArrow e a) where
 app = CO $ \k -> first (\(CO a) -> a k) ^>> app

liftCont :: Arrow a => a b c -> ContArrow e a b c
liftCont f = CO (\k -> f >>> k)

instance ArrowApply a => ArrowTransformer (ContArrow e) a where
 lift = liftCont

runCont :: ArrowApply a => ContArrow e a b c -> a c e -> a b e
runCont (CO f) = f


jump :: ArrowApply a => ContArrow e a (a c e, c) z
jump = CO (\_ -> app)

callcc :: ArrowApply a
 => (a c e -> ContArrow e a b c) -> ContArrow e a b c
callcc f = CO (\k -> let CO g = f k in g k)