Monadic robot acting in a comonadic world…

{-# LANGUAGE FlexibleInstances, DeriveFunctor, TypeFamilies #-}
module Main where

import Control.Comonad.Identity
import Control.Comonad.Trans.Class
import Control.Comonad.Trans.Cofree
import Control.Monad.Trans.Free
import Control.Monad (void)

import Control.Monad.State

import Data.Maybe (fromMaybe)

-- ==============================================================
-- Helpers
-- ==============================================================

-- | Try to apply a function.
try :: (a -> Maybe a) -> (a -> a)
try f w = fromMaybe w $ f w

-- | Unfold CofreeT structure using iteration.
-- Should be in Control.Comonad.Trans.Cofree
coiterT :: (Functor f, Comonad w) => (w a -> f (w a)) -> w a -> CofreeT f w a
coiterT psi = CofreeT . (extend $ \w -> extract w :< fmap (coiterT psi) (psi w))

-- | Tear down through a free monad transformer using iteration.
-- Should be in Control.Monad.Trans.Free
iterT :: (Functor f, Monad m) => (f (m a) -> m a) -> FreeT f m a -> m a
iterT psi (FreeT m) = do
    val <- m
    case fmap (iterT psi) val of
        Pure x -> return x
        Free y -> psi y

-- ==============================================================

-- | World base functor.
data WorldF pos x = WorldF
    { _worldPos     :: pos          -- position (cell ID)
    , _worldLeft    :: (Maybe x)    -- move to the left cell, if possible
    , _worldRight   :: (Maybe x)    -- move to the right cell, if possible
    } deriving (Functor)

-- | World cofree comonad transformer
type WorldT pos = CofreeT (WorldF pos)

-- | Interface of a comonadic world
class Comonad w => ComonadWorld w where
    type WPos w :: *
    wPos        :: w a -> WPos w
    wMoveLeft   :: w a -> Maybe (w a)
    wMoveRight  :: w a -> Maybe (w a)

-- | Implementation for WorldT.
instance Comonad w => ComonadWorld (CofreeT (WorldF pos) w) where
    type WPos (CofreeT (WorldF pos) w) = pos
    wPos       = _worldPos   . unwrap
    wMoveLeft  = _worldLeft  . unwrap
    wMoveRight = _worldRight . unwrap

-- | Robot base functor.
data RobotF pos x
    = RGetPos (pos -> x)    -- ^ get current position
    | RMoveLeft  x          -- ^ move left
    | RMoveRight x          -- ^ move right
    deriving (Functor)

-- | Robot free monad transformer.
type RobotT pos = FreeT (RobotF pos)

-- | Robot API.
class Monad m => MonadRobot m where
    type RPos m :: *
    getPos      :: m (RPos m)
    moveLeft    :: m ()
    moveRight   :: m ()

-- | Implementation for RobotT.
instance Monad m => MonadRobot (FreeT (RobotF pos) m) where
    type RPos (FreeT (RobotF pos) m) = pos
    getPos      = liftF $ RGetPos id
    moveLeft    = liftF $ RMoveLeft  ()
    moveRight   = liftF $ RMoveRight ()

-- | Run robot in given environment.
runRobot :: (ComonadWorld w, Monad m, pos ~ WPos w) => w a -> RobotT pos m r -> m r
runRobot w m = evalStateT (iterT runRobotF $ hoistFreeT lift $ m) w
    where
        runRobotF (RGetPos       f)    = gets   wPos >>= f
        runRobotF (RMoveLeft  next) = modify (try wMoveLeft)  >> next
        runRobotF (RMoveRight next) = modify (try wMoveRight) >> next

-- | Infinite world represinting Z.
infiniteWorld :: (Comonad w, Num t) => w t -> WorldT t w ()
infiniteWorld = void . coiterT f
    where
        f w = WorldF
            { _worldPos   = extract w
            , _worldLeft  = Just $ fmap (subtract 1) w
            , _worldRight = Just $ fmap (+1) w }

-- XXX: is it possible to create abstract cell (with no neighbors) ?
cell :: (ComonadWorld w) => WPos w -> w ()
cell pos = undefined

-- XXX: is it reasonable and possible to combine (semi-)finite worlds?
-- For instance, should that be possible:
--      cell 0 |~> cell 1   -- two-cell world
(<~|), (|~>) :: (ComonadWorld w) => w a -> w a -> w a
(<~|) = undefined
(|~>) = undefined

-- XXX: is it possible to provide an API for modifying world?
-- For instance, is it possible for a robot to "build" new cells?

-- sample world
world :: (Num t) => WorldT t Identity ()
world = infiniteWorld $ Identity 1

-- sample robot
robot :: (MonadRobot m) => m [RPos m]
robot = do
    moveLeft
    moveLeft
    x <- getPos

    moveRight
    moveRight
    moveRight
    y <- getPos

    moveRight
    z <- getPos

    return [x, y, z]

-- main
main :: IO ()
main = do
    res <- runRobot world robot
    print res