spawn-yield.hs

{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE NoFieldSelectors #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

import Control.Monad.State
import Control.Monad.Trans.Cont
import GHC.Generics
import Control.Lens
import Data.Generics.Labels
import Control.Concurrent.MVar
import Data.Foldable (for_)
import Prelude hiding (log)

newtype PID = PID Int
  deriving newtype (Show)

data SchedulerCtx = MkSchedulerCtx
  { queue :: [(PID, M ())]
  , fuel :: Int
  , gen :: Int
  , current :: ~PID
  } deriving stock (Generic)

_INITIAL_FUEL :: Int
_INITIAL_FUEL = 2

type M = ContT () (StateT SchedulerCtx IO)

push :: (PID, M ()) -> M ()
push p = #queue <>= [p]

pop :: M (Maybe (PID, M ()))
pop = do
  use #queue >>= \case
    [] -> pure Nothing
    p:ps -> do
      #queue .= ps
      pure (Just p)

self :: M PID
self = use #current

fresh :: M Int
fresh = #gen <<+= 1

spawn :: M () -> M PID
spawn action = do
  pid <- PID <$> fresh
  push (pid, do action; switch)
  pure pid

yield :: M ()
yield = do
  fuel <- #fuel <-= 1
  when (fuel == 0) do
    shiftT \next -> do
      pid <- use #current
      push (pid, lift (next ()))
      switch

switch :: M ()
switch = do
  pop >>= \case
    Nothing -> pure ()
    Just (pid, task) -> do
      #fuel .= _INITIAL_FUEL
      #current .= pid
      task

runM :: M a -> IO a
runM action = do
  result <- newEmptyMVar
  let action' = liftIO . putMVar result =<< action
  evalStateT (runContT (do spawn action'; switch) pure) initial
  readMVar result
  where
    initial = MkSchedulerCtx
      { queue = []
      , fuel = 0
      , current = undefined
      , gen = 0
      }

log :: String -> M ()
log msg = do
  pid <- self
  liftIO do putStrLn ("<" ++ show pid ++ "> " ++ msg)

printer :: M ()
printer = do
  log "start"
  for_ [1..5] \value -> do
    log (show value)
    yield
  log "end"

main :: IO ()
main = runM do
  log "main start"
  replicateM_ 5 do
    spawn printer
    yield
  log "main end"