module Quox.ST

import Data.IORef
import Control.MonadRec


export
Tag : Type
Tag = () -- shhh don't tell anyone

export
record ST (s : Tag) a where
  constructor MkST
  action : IO a
%name ST st

export
runST : (forall s. ST s a) -> a
runST st = unsafePerformIO (st {s = ()}).action

export %inline Functor (ST s) where map f st = MkST $ map f st.action

export %inline
Applicative (ST s) where
  pure = MkST . pure
  f <*> x = MkST $ f.action <*> x.action

export %inline Monad (ST s) where m >>= k = MkST $ m.action >>= action . k

export %inline
MonadRec (ST s) where
  tailRecM s (Access r) x k = MkST $ do
    let MkST yy = k s x
    case !yy of
      Done y      => pure y
      Cont s2 p y => let MkST z = tailRecM s2 (r s2 p) y k in z


public export
interface HasST (0 m : Tag -> Type -> Type) where
  liftST : ST s a -> m s a

export %inline HasST ST where liftST = id

export
record STRef (s : Tag) a where
  constructor MkSTRef
  ref : IORef a
%name STRef r

export %inline
newRef : HasST m => a -> m s (STRef s a)
newRef x = liftST $ MkST $ MkSTRef <$> newIORef x

export %inline
readRef : HasST m => STRef s a -> m s a
readRef r = liftST $ MkST $ readIORef r.ref

export %inline
writeRef : HasST m => STRef s a -> a -> m s ()
writeRef r x = liftST $ MkST $ writeIORef r.ref x

export %inline
modifyRef : HasST m => STRef s a -> (a -> a) -> m s ()
modifyRef r f = liftST $ MkST $ modifyIORef r.ref f


public export
record STErr e (s : Tag) a where
  constructor STE
  fromSTErr : ST s (Either e a)

export
Functor (STErr e s) where
  map f (STE e) = STE $ map f <$> e

export
Applicative (STErr e s) where
  pure x = STE $ pure $ pure x
  STE f <*> STE x = STE [|f <*> x|]

export
Monad (STErr e s) where
  STE m >>= k = STE $ do
    case !m of
      Left  err => pure $ Left err
      Right x   => fromSTErr $ k x

export
MonadRec (STErr e s) where
  tailRecM s (Access r) x k = STE $ do
    let STE m = k s x
    case !m of
      Left  err           => pure $ Left err
      Right (Cont s' p y) => fromSTErr $ tailRecM s' (r s' p) y k
      Right (Done y)      => pure $ Right y

export
runSTErr : (forall s. STErr e s a) -> Either e a
runSTErr ste = runST $ fromSTErr ste

export %inline HasST (STErr e) where liftST = STE . map Right

export
stLeft : e -> STErr e s a
stLeft e = STE $ pure $ Left e