Translate.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.Machine.Translate where

import Control.Lens
import Control.Monad
import Data.Machine

type FPrism f = forall g g' j. Prism (f g j) (f g' j) (g j) (g' j)

type Translate m b c = forall d i (f :: (* -> *) -> * -> *). FPrism f -> (i -> MachineT m (f b) d) -> f b i -> MachineT m (f b) d -> MachineT m (f c) d

newtype Id f a = Id {runId :: f a}

newtype Compose (f :: (* -> *) -> * -> *) (g :: (* -> *) -> * -> *) (a :: * -> *) (i :: *) = Compose {runCompose :: f (g a) i}

translateWith :: forall m b c d f. Monad m => FPrism f -> Translate m b c -> MachineT m (f b) d -> MachineT m (f c) d
translateWith p t m = MachineT $ runMachineT m >>= \u -> case u of
  Stop -> return Stop
  Yield o k -> return $ Yield o $ translateWith p t k
  Await f rq ff -> runMachineT $ t p f rq ff

translate :: forall m b c d. Monad m => Translate m b c -> MachineT m b d -> MachineT m c d
translate t = fit runId . translateWith (iso runId Id) t . fit Id 

translator :: forall m b c. Monad m => MachineT m b c -> Translate m (Is c) b
translator m = go
  where
    go :: forall d i (f :: (* -> *) -> * -> *). FPrism f
       -> (i -> MachineT m (f (Is c)) d) -> f (Is c) i
       -> MachineT m (f (Is c)) d -> MachineT m (f b) d
    go p f rq ff = case matching p rq of
      Left (rq' :: f b i)-> MachineT $ return $
        Await (\i -> translateWith p go (f i)) rq' (translateWith p go ff)
      Right Refl -> MachineT $ runMachineT m >>= \u -> case u of
        Stop -> runMachineT $ translateWith p (translator stopped) ff
        Yield o k -> runMachineT $ translateWith p (translator k) $ f o
        Await f' rq' ff' -> return $
          Await (\i -> (translator (f' i) p f (review p Refl) ff)) (review p rq') $ translator ff' p f (review p Refl) ff

tgPrism :: Prism (TG a b d) (TG a c d) (b d) (c d)
tgPrism = prism R $ \u -> case u of
  L rq -> Left $ L rq
  R rq -> Right rq

mapR :: forall m a b c. Monad m => Translate m b c -> Translate m (TG a b) (TG a c)
mapR t = go
  where
    go :: forall d i (f :: (* -> *) -> * -> *). FPrism f
       -> (i -> MachineT m (f (TG a b)) d) -> f (TG a b) i
       -> MachineT m (f (TG a b)) d -> MachineT m (f (TG a c)) d
    go p f rq ff = fit runCompose $ t p' (\i -> fit Compose (f i)) (Compose rq) (fit Compose ff)
      where
        p' :: FPrism (Compose f (TG a))
        p' = iso runCompose Compose . p . tgPrism

compose :: forall m a b c. Monad m => Translate m b c -> Translate m a b -> Translate m a c
compose t1 t2 = go
  where
    go :: forall d i (f :: (* -> *) -> * -> *). FPrism f
       -> (i -> MachineT m (f a) d) -> f a i
       -> MachineT m (f a) d -> MachineT m (f c) d
    go p f rq ff = translateWith p t1 $ t2 p f rq ff

foo :: Machine (TG (Is Int) (Is Int)) Int
foo = construct $ replicateM 5 $ do
    x <- awaits ar
    yield x

bar :: Monad m => Translate m (TG (Is Int) (Is Int)) (TG (Is Int) (Is Int))
bar = mapR $ translator $ repeatedly $ yield 1