module Quox.EffExtra

import public Control.Eff

import Control.Monad.ST.Extra
import Data.IORef


export
localAt : (0 lbl : tag) -> Has (StateL lbl s) fs =>
          (s -> s) -> Eff fs a -> Eff fs a
localAt lbl f act = do
  old <- getAt lbl
  modifyAt lbl f *> act <* putAt lbl old

export %inline
localAt_ : (0 lbl : tag) -> Has (StateL lbl s) fs =>
           s -> Eff fs a -> Eff fs a
localAt_ lbl x = localAt lbl $ const x

export %inline
local : Has (State s) fs => (s -> s) -> Eff fs a -> Eff fs a
local = localAt ()

export %inline
local_ : Has (State s) fs => s -> Eff fs a -> Eff fs a
local_ = localAt_ ()


export %inline
getsAt : (0 lbl : tag) -> Has (StateL lbl s) fs => (s -> a) -> Eff fs a
getsAt lbl f = f <$> getAt lbl

export %inline
gets : Has (State s) fs => (s -> a) -> Eff fs a
gets = getsAt ()


export
handleStateIORef : HasIO m => IORef st -> StateL lbl st a -> m a
handleStateIORef r Get     = readIORef r
handleStateIORef r (Put s) = writeIORef r s

export
handleStateSTRef : HasST m => STRef s st -> StateL lbl st a -> m s a
handleStateSTRef r Get     = liftST $ readSTRef r
handleStateSTRef r (Put s) = liftST $ writeSTRef r s



public export
data Length : List a -> Type where
  Z : Length []
  S : Length xs -> Length (x :: xs)
%builtin Natural Length

export
subsetWith : Length xs => (forall z. Has z xs -> Has z ys) ->
             Subset xs ys
subsetWith @{Z}     f = []
subsetWith @{S len} f = f Z :: subsetWith (f . S)

export
subsetSelf : Length xs => Subset xs xs
subsetSelf = subsetWith id

export
subsetTail : Length xs => (0 x : a) -> Subset xs (x :: xs)
subsetTail _ = subsetWith S



export
rethrowAtWith : (0 lbl : tag) -> Has (ExceptL lbl e') fs =>
                (e -> e') -> Either e a -> Eff fs a
rethrowAtWith lbl f = rethrowAt lbl . mapFst f

export
rethrowWith : Has (Except e') fs => (e -> e') -> Either e a -> Eff fs a
rethrowWith = rethrowAtWith ()

export
wrapErr : Length fs => (e -> e') ->
          Eff (ExceptL lbl e  :: fs) a ->
          Eff (ExceptL lbl e' :: fs) a
wrapErr f act =
  catchAt lbl (throwAt lbl . f) @{S Z} $
  lift @{subsetTail _} act


export
handleExcept : Functor m => (forall c. e -> m c) -> ExceptL lbl e a -> m a
handleExcept thr (Err e) = thr e


export
handleReaderConst : Applicative m => r -> ReaderL lbl r a -> m a
handleReaderConst x Ask = pure x

export
handleWriterST : HasST m => STRef s (SnocList w) -> WriterL lbl w a -> m s a
handleWriterST ref (Tell w) = liftST $ modifySTRef ref (:< w)


public export
record IOErr e a where
  constructor IOE
  fromIOErr : IO (Either e a)

export
Functor (IOErr e) where
  map f (IOE e) = IOE $ map f <$> e

export
Applicative (IOErr e) where
  pure x = IOE $ pure $ pure x
  IOE f <*> IOE x = IOE [|f <*> x|]

export
Monad (IOErr e) where
  IOE m >>= k = IOE $ do
    case !m of
      Left  err => pure $ Left err
      Right x   => fromIOErr $ k x

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

export
HasIO (IOErr e) where
  liftIO = IOE . map Right

export %inline
ioLeft : e -> IOErr e a
ioLeft = IOE . pure . Left