mad/byg/src/DependencyRunner.hs

{-# LANGUAGE MonoLocalBinds #-}
module DependencyRunner
  ( DepRunM
  , runDeps
  , runDepRunMIO
  ) where

import Types (evalFunctionIO)
import Types.Value
import Types.Token
import Types.Dependency

import Type.Reflection (Typeable)
import Data.Map (Map)
import qualified Data.Map as M
import Control.Monad (void, forM)
import Control.Monad.State (MonadState, MonadIO, StateT, evalStateT, get, modify, liftIO)

data ValueExistence = Evaluated Value
                    | NotEvaluated (DepRunM Value)

newtype DepRunM a = DepRunM { unDepRunM :: StateT (Map Int ValueExistence) IO a }
  deriving (Functor, Applicative, Monad, MonadIO, MonadState (Map Int ValueExistence))

runDeps :: [Dependency] -> DepRunM ()
runDeps = mapM_ runDep

runDepRunMIO :: DepRunM a -> IO a
runDepRunMIO m = evalStateT (unDepRunM m) M.empty

evaluate :: ValueExistence -> DepRunM Value
evaluate = \case
  Evaluated v -> pure v
  NotEvaluated m -> m

runDep :: Dependency -> DepRunM ()
runDep (Dependency _ a action _ b) =
  if actionWritesAny action
  then void m
  else putTokenValue b $ NotEvaluated m
  where m :: DepRunM Value
        m = do
          input <- getTokenValue a
          result <- runAction action input
          putTokenValue b $ Evaluated result
          liftIO $ do
            putStrLn ("input: " ++ show input)
            putStrLn ("action: " ++ show action)
            putStrLn ("output: " ++ show result)
            putStrLn "----------"
          pure result

getTokenValueByIndex :: Int -> DepRunM Value
getTokenValueByIndex i = do
  m <- get
  evaluate (m M.! i)

getTokenValueRaw :: Token a -> DepRunM a
getTokenValueRaw token = case token of
  Token i ->
    fromValue <$> getTokenValueByIndex i
  TupleToken a b -> do
    (,) <$> getTokenValueRaw a <*> getTokenValueRaw b
  ZipToken a b -> do
    zip <$> getTokenValueRaw a <*> getTokenValueRaw b
  ListToken ts -> do
    mapM getTokenValueRaw ts
  NoToken ->
    pure ()

getTokenValue :: Token a -> DepRunM Value
getTokenValue token = case token of
  Token i ->
    getTokenValueByIndex i
  TupleToken _ _ ->
    fromRaw token
  ZipToken _ _ ->
    fromRaw token
  ListToken _ ->
    fromRaw token
  NoToken ->
    fromRaw token
  where fromRaw :: (Typeable a, Show a) => Token a -> DepRunM Value
        fromRaw = fmap toValue . getTokenValueRaw

putTokenValue :: Token a -> ValueExistence -> DepRunM ()
putTokenValue t e = case t of
  Token i ->
    modify $ M.insert i e
  NoToken ->
    pure ()
  _ ->
    error "unexpected"

runAction :: forall a b. Action a b -> Value -> DepRunM Value
runAction action input = case action of
  Function (F f) ->
    calc f
  FunctionIO f ->
    calcM (liftIO . evalFunctionIO f)
  Inject x ->
    pure $ toValue x
  FilterComp ->
    calc (\(vs, mask) -> map fst $ filter snd $ zip vs mask)
  UntupleFst ->
    calc fst
  UntupleSnd ->
    calc snd
  UnzipFst ->
    calc (map fst)
  UnzipSnd ->
    calc (map snd)
  MapComp subDeps innerInput innerOutput ->
    (toValueRep (actionTargetType action) <$>)
    $ forM (fromValueRep (actionSourceType action) input) $ \x -> do
      putTokenValue innerInput $ Evaluated $ toValue x
      runDeps subDeps
      fromValue <$> getTokenValue innerOutput
  where calcM :: (Typeable a, Typeable b, Show b) => (a -> DepRunM b) -> DepRunM Value
        calcM f = toValue <$> (f $ fromValueRep (actionSourceType action) input)

        calc :: (Typeable a, Typeable b, Show b) => (a -> b) -> DepRunM Value
        calc f = calcM (pure . f)