package.yaml

name: hspg

dependencies:
  - base
  - lens
  - transformers
  - generic-lens
  - ghc-prim

executables:
  hspg-exe:
    source-dirs: app
    main: Main.hs

spawn_yield_control0.hs

{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE NoFieldSelectors #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE GeneralizedNewtypeDeriving, ImplicitParams, RankNTypes, ConstraintKinds #-}
{-# LANGUAGE MagicHash, ScopedTypeVariables, UnboxedTuples #-}

import Control.Monad.Trans.State (State, runState)
import GHC.Generics (Generic)
import Control.Lens (use, (.=), (<-=), (<<+=), (<>=))
import Data.Generics.Labels ()
import Control.Monad (replicateM_, when)
import Data.Foldable (for_)
import Prelude hiding (log)

import Data.IORef (IORef, newIORef, atomicModifyIORef')

import GHC.Prim (control0#, newPromptTag#, prompt#, PromptTag#)
import GHC.IO (IO(..))
import Data.Coerce (coerce)


-- cont

data PromptTag a = MkPromptTag (PromptTag# a)

newPromptTag :: forall a. IO (PromptTag a)
newPromptTag = IO \s -> case newPromptTag# s of (# s, tag #) -> (# s, MkPromptTag tag #)

prompt :: forall a. PromptTag a -> IO a -> IO a
prompt (MkPromptTag tag) = coerce (prompt# @a tag)

control0 :: forall a b. PromptTag a -> ((IO b -> IO a) -> IO a) -> IO b
control0 (MkPromptTag tag) = coerce (control0# @a @b tag)



-- spawn+yield

newtype PID = PID Int
  deriving newtype (Show)

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

_INITIAL_FUEL :: Int
_INITIAL_FUEL = 2

state :: State SchedulerCtx a -> ((?ctx :: IORef SchedulerCtx) => IO a)
state action = atomicModifyIORef' ?ctx \s ->
  let (s', a) = runState action s in (a, s')

type WithCtx = (?ctx :: IORef SchedulerCtx, ?tag :: PromptTag ())

push :: WithCtx => (PID, IO ()) -> IO ()
push p = state do #queue <>= [p]

pop :: WithCtx => IO (Maybe (PID, IO ()))
pop = do
  queue <- state do use #queue
  case queue of
    [] -> pure Nothing
    p:ps -> do
      state do #queue .= ps
      pure (Just p)

self :: WithCtx => IO PID
self = state do use #current

fresh :: WithCtx => IO Int
fresh = state do #gen <<+= 1

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

yield :: WithCtx => IO ()
yield = do
  fuel <- state do #fuel <-= 1
  when (fuel <= 0) do
    control0 ?tag \next -> do
      pid <- state do use #current
      push (pid, next (pure ()))
      switch

switch :: WithCtx => IO ()
switch = do
  pop >>= \case
    Nothing -> pure ()
    Just (pid, task) -> do
      state do
        #fuel .= _INITIAL_FUEL
        #current .= pid
      prompt ?tag task

runM :: (WithCtx => IO ()) -> IO ()
runM action = do
  ctx <- newIORef initial; let ?ctx = ctx
  tag <- newPromptTag; let ?tag = tag
  spawn action
  switch
  where
    initial = MkSchedulerCtx
      { queue = []
      , fuel = 0
      , current = undefined
      , gen = 0
      }



-- Example

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

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

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

main :: IO ()
main = do
  runM example
  putStrLn "done"