mad/byg/src/DependencyRunner.hs

module DependencyRunner
  ( DepRunM
  , runDeps
  , runDepRunM
  ) where

import Types (Value(..), Valuable(..))
import Types.Dependency
import Evaluation.Function
import Evaluation.FunctionIO

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

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

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

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

runDep :: Dependency -> DepRunM ()
runDep (Dependency a action b) = do
  input <- getTokenValue a
  result <- runAction action input
  putTokenValue b result

getTokenValue :: UToken -> DepRunM Value
getTokenValue = \case
  UToken i -> do
    m <- get
    pure (m M.! i)
  UTupleToken (a, b) -> do
    va <- getTokenValue a
    vb <- getTokenValue b
    pure $ Tuple (va, vb)
  UZipToken (a, b) -> do
    va <- getTokenValue a
    vb <- getTokenValue b
    case (va, vb) of
      (List as, List bs) ->
        pure $ List $ zipWith (curry Tuple) as bs
      _ ->
        error "unexpected"
  UNoToken ->
    pure Empty

putTokenValue :: UToken -> Value -> DepRunM ()
putTokenValue t v = case t of
  UToken i ->
    modify $ M.insert i v
  UNoToken ->
    pure ()
  _ ->
    error "unexpected"

runAction :: Action -> Value -> DepRunM Value
runAction action input = case action of
  Function f ->
    pure $ evalFunction f input
  FunctionIO f ->
    liftIO $ evalFunctionIO f input
  Inject v ->
    pure v
  FilterComp ->
    case input of
      Tuple (List vs, List mask) ->
        pure $ List $ map fst $ filter (fromValue . snd) $ zip vs mask
      _ ->
        error "unexpected"
  UntupleFst ->
    case input of
      Tuple (v, _) ->
        pure v
      _ ->
        error "unexpected"
  UntupleSnd ->
    case input of
      Tuple (_, v) ->
        pure v
      _ ->
        error "unexpected"
  UnzipFst ->
    case input of
      List vs ->
        List <$> mapM (runAction UntupleFst) vs
      _ ->
        error "unexpected"
  UnzipSnd ->
    case input of
      List vs ->
        List <$> mapM (runAction UntupleSnd) vs
      _ ->
        error "unexpected"
  MapComp subDeps innerInput innerOutput ->
    case input of
      List vs ->
        (List <$>) $ flip mapM vs $ \v -> do
          putTokenValue innerInput v
          runDeps subDeps
          getTokenValue innerOutput
      _ ->
        error "unexpected"