module Control.Monad.ST.Extra

import public Control.Monad.ST
import Data.IORef
import Control.MonadRec

%default total

export %inline
MonadRec (ST s) where
  tailRecM seed (Access rec) st f = MkST $ do
    let MkST io = f seed st
    case !io of
      Done res           => pure res
      Cont seed2 prf vst =>
        let MkST io = tailRecM seed2 (rec seed2 prf) vst f in io


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

export %inline HasST ST where liftST = id


public export
record STErr e s 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


parameters {auto _ : HasST m}
  export %inline
  newSTRef' : a -> m s (STRef s a)
  newSTRef' x = liftST $ newSTRef x

  export %inline
  readSTRef' : STRef s a -> m s a
  readSTRef' r = liftST $ readSTRef r

  export %inline
  writeSTRef' : STRef s a -> a -> m s ()
  writeSTRef' r x = liftST $ writeSTRef r x

  export %inline
  modifySTRef' : STRef s a -> (a -> a) -> m s ()
  modifySTRef' r f = liftST $ modifySTRef r f