Excercises from https://haskell-lang.org/tutorial/primitive-haskell

Prim.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

import Control.Monad.Reader.Class
import GHC.ST
import Control.Monad.ST
import GHC.Prim
import GHC.Types
import Control.Monad.IO.Class

ioToSt :: IO a -> ST RealWorld a
ioToSt (IO a) = ST $ \s -> a s

stToIO :: ST RealWorld a -> IO a
stToIO (ST a) = IO $ \s -> a s

data STRef s a = STRef (MutVar# s a)

newSTRef :: a -> ST s (STRef s a)
newSTRef a = ST $ \s -> case newMutVar# a s of
                          (# s', ma #) -> (# s', STRef ma #)

readSTRef :: STRef s a -> ST s a
readSTRef (STRef mv) = ST $ \s -> readMutVar# mv s

writeSTRef :: STRef s a -> a -> ST s ()
writeSTRef (STRef mv) a = ST $ \s -> case writeMutVar# mv a s of
                                       s' -> (# s', () #)

class Monad m => PrimMonad m where
    type PrimState m
    primitive :: (State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a

instance PrimMonad IO where
  type PrimState IO = RealWorld
  primitive f = IO $ \rw -> f rw

instance PrimMonad (ST s) where
  type PrimState (ST s) = s
  primitive f = ST $ \s -> f s

data PrimRef m a = PrimRef (MutVar# (PrimState m) a)

newPrimRef :: PrimMonad m => a -> m (PrimRef m a)
newPrimRef a = primitive $ \s -> case newMutVar# a s of
                                   (# s', mv #) -> (# s', PrimRef mv #)

readPrimRef :: PrimMonad m => PrimRef m a -> m a
readPrimRef (PrimRef mv) = primitive (readMutVar# mv)

writePrimRef :: PrimMonad m => PrimRef m a -> a -> m ()
writePrimRef (PrimRef mv) a = primitive $ \s -> case writeMutVar# mv a s of
                                                  s' -> (# s', () #)


newtype ReaderIO r a = ReaderIO (r -> State# RealWorld -> (# State# RealWorld, a #))

runRIO :: r -> ReaderIO r a -> IO a
runRIO r (ReaderIO ra) = IO $ ra r

instance Functor (ReaderIO r) where
  fmap :: (a -> b) -> ReaderIO r a -> ReaderIO r b
  fmap f (ReaderIO ra) = ReaderIO $ \r s -> case ra r s of
                                     (# s', a #) -> (# s', f a #)

instance Applicative (ReaderIO r) where
  pure :: a -> ReaderIO r a
  pure a = ReaderIO $ \_ s -> (# s, a #)

  (<*>) :: ReaderIO r (a -> b) -> ReaderIO r a -> ReaderIO r b
  (ReaderIO rf) <*> (ReaderIO ra) = ReaderIO $ \r s ->
    case rf r s of
      (# s', f #) -> case ra r s' of
                       (# s'', a #) -> (# s'', f a #)

instance Monad (ReaderIO r) where
  return = pure
  (>>=) :: ReaderIO r a -> (a -> ReaderIO r b) -> ReaderIO r b
  (ReaderIO ra) >>= f = ReaderIO $ \r s ->
    case ra r s of
      (# s', a #) ->
        case f a of
          ReaderIO rb -> rb r s'

instance MonadIO (ReaderIO r) where
  liftIO :: IO a -> ReaderIO r a
  liftIO (IO ia) = ReaderIO $ \_ s -> ia s

instance MonadReader r (ReaderIO r) where
  ask :: ReaderIO r r
  ask = ReaderIO $ \r s -> (# s, r #)

  local :: (r -> r) -> ReaderIO r a -> ReaderIO r a
  local f (ReaderIO ra) = ReaderIO $ \r s ->
    ra (f r) s

instance PrimMonad (ReaderIO r) where
  type PrimState (ReaderIO r) = RealWorld
  primitive :: (State# RealWorld -> (# State# RealWorld, a #)) -> ReaderIO r a
  primitive f = ReaderIO $ \_ s -> f s

testRIO :: IO ()
testRIO = runRIO (42 :: Int) $ do
  r <- ask
  liftIO $ putStrLn $ show r
  return ()