module Cont where

-- resetST に対応するための
-- Continuation + State Transformer Monad

import MyState
import MyStream

newtype K r a = K ((a -> r) -> r)

unK :: K r a -> (a -> r) -> r
unK (K c) = c

instance Monad (K r) where
  return a    = K (\ c -> c a)
  (K m) >>= k = K (\ c -> m (\ a -> unK (k a) c))

instance Functor (K r) where
  fmap f m = m >>= \ x -> return (f x)

instance Applicative (K r) where
  pure = return
  g <*> m = g >>= \ f -> m >>= \ x -> return (f x)

evalK :: K r r -> r
evalK m = unK m (\ a -> a) 

{-
type K r q a = (a -> r) -> q

returnK :: a -> K r r a  
returnK a = \ c -> c a

bindK :: K r q a -> (a -> K p r b) -> K p q b
m `bindK` k = \ c -> m (\ a -> k a c)

callccK :: ((a -> K p r b) ->  K r q a)  -> K r q a
callccK h = \ c -> let k a = \ d -> c a in h k c

-- callccK2 :: (a -> a -> r) -> a -> r
callccK2 h c = h c c

throw c a _ = c a

type ST s t a = s -> (a, t)

callKST :: y -> K ((x, u) -> r) ((w, y) -> q) a
             -> K ((x, z) -> r) ((w, z) -> q) a
callKST y m = \ c (x,z) -> m (\ a (x1,_) -> c a (x1,z)) (x,y)
--}

newtype KS r s a = KS ((a -> s -> r) -> s -> r)

unKS :: KS r s a -> (a -> s -> r) -> s -> r
unKS (KS m) = m 

instance Monad (KS r s) where
  return a     = KS (\ c -> c a)
  (KS m) >>= k = KS (\ c -> m (\ a -> unKS (k a) c))

instance Functor (KS r s) where
  fmap f m = m >>= \ x -> return (f x)

instance Applicative (KS r s) where
  pure = return
  g <*> m = g >>= \ f -> m >>= \ x -> return (f x)


abortKS :: (s -> r) -> KS r s a
abortKS f = KS (\ c -> f)
  
callccKS :: ((a -> KS r s b) ->  KS r s a)  -> KS r s a
callccKS h = KS (\ c -> let k a = KS (\ d -> c a) in unKS (h k) c)

callcc2KS :: ((a -> s -> r) -> KS r s a) -> KS r s a
callcc2KS h = KS (\ c -> unKS (h c) c)

throwKS :: (a -> s -> r) -> a -> KS r s b
throwKS c a = KS (\ _ -> c a)

extendKS :: Pos s t -> KS r t a -> KS r s a
extendKS p (KS m) = KS (\ c s ->
                         let (t, rest) = p s
                           in m (\ a t1 -> c a (rest t1)) t)

shrinkKS :: (s -> t) -> (t -> s) -> KS r s a -> KS r t a
shrinkKS extract rest (KS m) = KS (\ c t -> m (\ a s1 -> c a (extract s1)) (rest t))

instance MyState (KS r) where
  get p    = KS (\ c s -> c (fst (p s)) s)
  set p v  = KS (\ c s -> c () (snd (p s) v))
  extend   = extendKS
  shrink   = shrinkKS

-- reset' :: y -> KS r (x, y) a -> KS r (x, z) a
-- reset' y (KS m) = KS (\ c (x,z) -> m (\ a (x1,_) -> c a (x1,z)) (x,y))

callKS :: y -> KS r (x, y) a -> KS r (x, z) a
callKS y m = extendKS xP $ shrinkKS (\ (x, _) -> x) (\ x -> (x, y)) m

startKS :: x -> KS r (x, y) a -> KS r (w, y) a
startKS x m = extendKS yP $ shrinkKS (\ (_, y) -> y) (\ y -> (x, y)) m 

evalKS :: KS r s r -> s -> r
evalKS m s = unKS m (\ a s -> a) s

newtype KIO r s a = KIO ((a -> WithIO s -> r) -> WithIO s -> r)

unKIO :: KIO r s a -> (a -> WithIO s -> r) -> WithIO s -> r
unKIO (KIO m) = m

instance Monad (KIO r s) where
  return a      = KIO (\ c -> c a)
  (KIO m) >>= k = KIO (\ c -> m (\ a -> unKIO (k a) c))

instance Functor (KIO r s) where
  fmap f m = m >>= \ x -> return (f x)

instance Applicative (KIO r s) where
  pure = return
  g <*> m = g >>= \ f -> m >>= \ x -> return (f x)

extendKIO :: Pos s t -> KIO r t a -> KIO r s a
extendKIO p (KIO m) = KIO (\ c (s, i, o) ->
                          let (t, rest) = p s
                            in m (\ a (t1, i1, o1) -> c a (rest t1, i1, o1)) (t, i, o))

shrinkKIO :: (s -> t) -> (t -> s) -> KIO r s a -> KIO r t a
shrinkKIO extract rest (KIO m) = KIO (\ c (t, i, o) ->
    m (\ a (s1, i1, o1) -> c a (extract s1, i1, o1))
      (rest t, i, o))

instance MyState (KIO r) where
  get p    = KIO (\ c (s, i, o) -> c (fst (p s)) (s, i, o))
  set p v  = KIO (\ c (s, i, o) -> c () (snd (p s) v, i, o))
  extend   = extendKIO
  shrink   = shrinkKIO

instance MyStream (KIO r s) where
  readChar   = KIO (\ c (s, ch : i, o) -> c ch (s, i, o))
  eof        = KIO (\ c (s, i, o) -> c (null i) (s, i, o))
  writeStr v = KIO (\ c (s, i, o) -> c () (s, i, o ++ v))

abort :: (WithIO s -> r) -> KIO r s a
abort f = KIO (\ c -> f)
  
callcc :: ((a -> KIO r s b) ->  KIO r s a)  -> KIO r s a
callcc h = KIO (\ c -> let k a = KIO (\ d -> c a) in unKIO (h k) c)

callcc2 :: ((a -> WithIO s -> r) -> KIO r s a) -> KIO r s a
callcc2 h = KIO (\ c -> unKIO (h c) c)

throwKIO :: (a -> WithIO s -> r) -> a -> KIO r s b
throwKIO c a = KIO (\ _ -> c a)

resetKIO :: y -> KIO r (x, y) a -> KIO r (x, z) a
resetKIO y (KIO m) = KIO (\ c ((x,z),i,o) -> m (\ a ((x1,_),i1,o1) -> c a ((x1,z),i1,o1)) ((x,y),i,o))

{-
extendKIO :: Pos s t -> KIO r t a -> KIO r s a
extendKIO p (KIO m) = KIO (\ c (s, i, o) ->
                        let (t, rest) = p s
                          in m (\ a (t1, i1, o1) -> c a (rest t1, i1, o1)) (t, i, o))

shrinkKIO :: Pos s t -> (t -> s) -> KIO r s a -> KIO r t a
shrinkKIO p rest (KIO m) = KIO (\ c (t, i, o) -> m (\ a (s1, i1, o1) -> c a (fst (p s1), i1, o1)) (rest t, i, o))
-}

evalKIO :: KIO r s r -> s -> r
evalKIO m s  = unKIO m (\ a s -> a) (s, "", "")

evalKIO2 :: KIO String s a -> s -> String
evalKIO2 m s = unKIO m (\ a (_,_,o) -> o) (s,"","")

data Trick v s = Trick { unTrick :: Trick v s -> KIO v s (Trick v s)}
-- newtype Trick v s = Trick { unTrick :: Trick v s -> KIO v s (Trick v s)}
-- or, newtype Mu0 m = Mu0 { unMu0 :: Mu0 m -> m (Mu0 m) }

newtype Mu r = Mu { unMu :: Mu r -> r }

-- Mu   :: (Mu r -> r) -> Mu r
-- unMu :: Mu r -> (Mu r -> r)

{-
foo = \ y -> begin
  set xP 1; set yP y;
  while get yP > 0 do begin
    val x = get xP in val y = get yP in
    begin
      if y == 10 then break
      else if y == 3 then begin
        set yP (y - 1); continue
      end else ();
      set xP (x * y);
      set yP (y - 1)
    end
  end;
  get xP
end
-}

{-
foo :: Integer -> KIO a (Integer,Integer) Integer
foo  = \ n ->
         set xP 1 >>= \ _ ->
         set yP n >>= \ _ ->
         KIO (\ _break ->
           let KIO _while
                 = get yP >>= \ _x ->
                   if _x > 0 then
                     get xP >>= \ r ->
                     get yP >>= \ n ->
                     (if n == 10 then abort (_break ()) else
                        if n == 3 then
                          set yP (n - 1) >>= \ _ ->
                          abort (_while _break)
                        else return ()) >>= \ _ ->
                     set xP (r * n) >>= \ _ ->
                     set yP (n - 1) >>= \ _ ->
                     KIO _while
                   else return ()
             in _while _break) >>= \ _ ->
         get xP
-}

{-
bar = \ _ -> begin
    set xP 1;
  label1:
    if get xP > 100 then goto label2 else ();
    set xP (get xP * 2);
    goto label1;
  label2:
    get xP
  end
-}

{-
bar :: t -> KIO a (Integer, y) Integer
bar = \ _ ->
       set xP 1     >>= \ _ ->
       KIO (\ _end ->
          let label1 = \ _ -> unKIO
                         (get xP                             >>= \ _x ->
                          (if _x > 100 then abort (label2 ())
                                       else return ())         >>= \ _ ->
                          get xP                             >>= \ _x ->
                          set xP (_x * 2)                    >>= \ _ ->
                          abort (label1 ())) label2
              label2 = \ _ -> unKIO (get xP) _end
            in label1 ())
-}

{-
baz = \ x ->
  callcc (\ k ->
    100 + (if x == 0 then 1 else k x))
-}

{-
baz
  = \ x ->
      callcc
        (\ k ->
           (if x == 0 then return 1 else k x) >>= \ _x -> return (100 + _x))
-}

{-
multlist = \ list ->
  let aux = \ xs -> \ k -> begin
        set xP 1; set yP xs;
        while not (null (get yP)) do begin
          val y = get yP in val n = head y in
          if n == 0 then k 0 else
            begin set xP (get xP * n); set yP (tail y);
                  writeStr " ";
                  write n
            end
        end;
        get xP
    end in
  val result = callcc (\ k -> aux list k) in begin
    writeStr ";\n result = ";
    write result
  end
-}

{-
multlist = \ list ->
      let aux = \ xs ->
            return (\ k ->
               set xP 1              >>= \ _ ->
               set yP xs             >>= \ _ ->
               KIO (\ _break ->
                  let KIO _while
                        = get yP >>= \ _x ->
                          if not (null _x) then
                               get yP                 >>= \ y ->
                               (if head y == 0 then k 0 else
                                   get xP               >>= \ _x ->
                                   set xP (_x * head y) >>= \ _ ->
                                   set yP (tail y)      >>= \ _ ->
                                   writeStr " "             >>= \ _ ->
                                   write (head y))      >>= \ _ ->
                               KIO _while
                             else return ()
                    in _while _break) >>= \ _ ->
                get xP)
        in callcc (\ k -> aux list >>= \ _f ->
                           _f k)   >>= \ result ->
           writeStr "; result="        >>= \ _ ->
           write result
-}

{-
-- Util
increase = \ n -> \ k ->
     if n > 10 then ()
     else begin writeStr " i:"; write n;
     increase (n + 1) (callcc k) end;
decrease = \ n -> \ k ->
     if n < 0 then ()
     else begin writeStr " d:"; write n;
     decrease (n - 1) (callcc k) end;
coroutine = \ _ -> increase 0 (decrease 10)

-- The Util code above translates into the following code,
-- which is untypeable due to the need of an infinite type.
-- r = r -> KIO v s r

increase n = return (\ k ->
               if n > 10 then return () else
                 writeStr " i:"   >>= \ _ ->
                 write n          >>= \ _ ->
                 increase (n + 1) >>= \ _f ->
                 callcc k         >>= \ _x ->
                 _f _x)

decrease n = return (\ k ->
               if n < 0 then return () else
                 writeStr " d:"   >>= \ _ ->
                 write n          >>= \ _ ->
                 decrease (n - 1) >>= \ f ->
                 callcc k         >>= \ x ->
                 f x)

-- which is hand-optimized to the following

increase n k = if n > 10 then return () else
                 writeStr " i:"   >>= \ _ ->
                 write n          >>= \ _ ->
                 callcc k         >>= \ x ->
                 increase (n + 1) x
decrease n k = if n < 0 then return () else
                 writeStr " d:"   >>= \ _ ->
                 write n          >>= \ _ ->
                 callcc k         >>= \ x ->
                 decrease (n - 1) x);

coroutine = increase 0 (decrease 10)

-- Here is a possible trick.

increase :: Integer -> MMu (KIO v s) -> KIO v s (MMu (KIO v s))
increase n k = if n > 10 then return k -- instead of (), to fit the type
                 else
                   writeStr " i:"                      >>= \ _ ->
                   write n                             >>= \ _ ->
                   callcc (\ c -> unTrick k (Trick c)) >>= \ x ->
                   increase (n + 1) x

decrease :: Integer -> MMu (KIO v s) -> KIO v s (MMu (KIO v s))
decrease n k = if n < 0 then return k  -- instead of (), to fit the type
                 else
                   writeStr " d:"                      >>= \ _ ->
                   write n                             >>= \ _ ->
                   callcc (\ c -> unTrick k (Trick c)) >>= \ x ->
                   decrease (n - 1) x

coroutine =  increase 0 (MMu (decrease 10))
main = print (unKIO (coroutine ()) (\ _ s -> ((), s)) ((),"",""))
-- e.g.) snd $ unKIO coroutine (\ _ s -> ((),s)) ((),"","")

-- which is translated back into Util

increase = \ n -> \ k ->
     if n > 10 then k
     else begin writeStr " i:"; write n;
     increase (n + 1) (callcc (\ c -> unMMu k (MMu c))) end;
decrease = \ n -> \ k ->
     if n < 0 then k
     else begin writeStr " d:"; write n;
     decrease (n - 1) (callcc (\ c -> unMMu k (MMu c))) end;
coroutine = \ _ -> increase 0 (MMu (decrease 10))

-}

{-
-- version for callcc2 and throw
-- Util
increase2 = \ n -> \ k ->
     if n > 10 then ()
     else begin writeStr " i:"; write n;
     increase2 (n + 1) (callcc2 (\ c -> throw k c)) end;

decrease2 = \ n -> \ k ->
     if n < 0 then ()
     else begin writeStr " d:"; write n;
     decrease2 (n - 1) (callcc2 (\ c -> throw k c)) end;

coroutine2 = \ _ -> callcc2 (\ u -> decrease2 10 (callcc2 (\ k -> throw u (increase2 0 k))))

-- The Util code above translates into the following code,
-- which is untypeable due to the need of an infinite type.
increase2
  = \ n -> return (\ k ->
       if n > 10 then return () else
          writeStr " i:"             >>= \ _ ->
          write n                    >>= \ _ ->
          increase2 (n + 1)          >>= \ _f ->
          callcc2 (\ c -> throw k c) >>= \ _x ->
           _f _x)
decrease2
  = \ n -> return (\ k ->
       if n < 0 then return () else
          writeStr " d:"             >>= \ _ ->
          write n                    >>= \ _ ->
          decrease2 (n - 1)          >>= \ _f ->
          callcc2 (\ c -> throw k c) >>= \ _x ->
          _f _x)

coroutine2 = \ _ -> callcc2 (\ u ->
                      decrease2 10                 >>= \ f ->
                      callcc2 (\ k -> increase2 0     >>= \ g ->
                                      g k             >>= \ y ->
                                      throw u y)   >>= \ x ->
                      f x)

-- which is hand-optimized to

increase2 n k = if n > 10 then return () else
                 writeStr " i:"             >>= \ _ ->
                 write n                    >>= \ _ ->
                 callcc2 (\ c -> throw k c) >>= \ x ->
                 increase2 (n + 1) x

decrease2 n k = if n < 0 then return () else
                 writeStr " d:"             >>= \ _ ->
                 write n                    >>= \ _ ->
                 callcc2 (\ c -> throw k c) >>= \ x ->
                 decrease2 (n - 1) x

coroutine2 = callcc2 (\ c -> callcc2 (\ d -> increase2 0 d >>= \ y ->
                                            throw c y) >>= \ x ->
                            decrease2 10 x)

-- This definition will cause the GHC simplifier panic!
-- y f = (\ x -> f (unMu x x)) (Mu (\ x -> f (unMu x x)))

increase2 :: Integer -> Mu (ST0 (WithIO s) v)
                     -> KIO v s (Mu (ST0 (WithIO s) v))
increase2 n k = if n > 10 then return k -- instead of (), to fit the type
                  else
                    writeStr " i:"               >>= \ _ ->
                    write n                      >>= \ _ ->
                    callcc2 (\ c -> throw (unMu k) (Mu c)) >>= \ x ->
                    increase2 (n + 1) x

decrease2 :: Integer -> Mu (ST0 (WithIO s) v)
                     -> KIO v s (Mu (ST0 (WithIO s) v))
decrease2 n k = if n < 0 then return k -- instead of (), to fit the type
                  else
                    writeStr " d:"               >>= \ _ ->
                    write n                      >>= \ _ ->
                    callcc2 (\ c -> throw (unMu k) (Mu c)) >>= \ x ->
                    decrease2 (n - 1) x

coroutine2 = callcc2 (\ c -> increase2 0 (Mu c)) >>= \ d ->
             decrease2 10 d

-- e.g.)  snd $ unKIO coroutine2 (\ _ s -> ((),s)) ((),"","")
-}