{-|
Module      : Helpers
Description : Helping functions, sort of ExtraExtra
Copyright   : (c) John Maraist, 2022
License     : AllRightsReserved
Maintainer  : haskell-tms@maraist.org
Stability   : experimental
Portability : POSIX

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, for NON-COMMERCIAL use.  See the License for the specific
language governing permissions and limitations under the License.

-}

{-# LANGUAGE RankNTypes #-}

module Data.TMS.Helpers where

import Control.Monad.State
import Control.Monad.ST.Trans
import Control.Monad.Except
import Control.Monad.Extra
import Data.List

-- * Lists and monads

-- | Check whether a list contains a value which, when applied to a
-- computation, returns @True@.
anyByM :: (Monad m) => (a -> m Bool) -> [a] -> m Bool
anyByM :: (a -> m Bool) -> [a] -> m Bool
anyByM a -> m Bool
_ [] = Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
anyByM a -> m Bool
k (a
x : [a]
xs) = do
  Bool
b <- a -> m Bool
k a
x
  if Bool
b then Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True else (a -> m Bool) -> [a] -> m Bool
forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
anyByM a -> m Bool
k [a]
xs

-- | Check whether a returned list contains a value which satisfies
-- some monadic predicate.
anyMM :: (Monad m) => (a -> m Bool) -> m [a] -> m Bool
anyMM :: (a -> m Bool) -> m [a] -> m Bool
anyMM a -> m Bool
predM m [a]
srcM = do
  [a]
src <- m [a]
srcM
  (a -> m Bool) -> [a] -> m Bool
forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
anyByM a -> m Bool
predM [a]
src

-- | Check whether all of the values of a list, when applied to a
-- computation, return @True@.
allByM :: (Monad m) => (a -> m Bool) -> [a] -> m Bool
allByM :: (a -> m Bool) -> [a] -> m Bool
allByM a -> m Bool
_ [] = Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
allByM a -> m Bool
k (a
x : [a]
xs) = do
  Bool
b <- a -> m Bool
k a
x
  if Bool
b then (a -> m Bool) -> [a] -> m Bool
forall (m :: * -> *) a. Monad m => (a -> m Bool) -> [a] -> m Bool
allByM a -> m Bool
k [a]
xs else Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

-- * Ordered lists

-- | Determine whether one list is a subset of the other, under the
-- assumption that both lists are sorted in ascending order.
ordSubsetp :: Ord a => [a] -> [a] -> Bool
ordSubsetp :: [a] -> [a] -> Bool
ordSubsetp [] [a]
_ = Bool
True
ordSubsetp (a
_ : [a]
_) [] = Bool
False
ordSubsetp l1 :: [a]
l1@(a
n1 : [a]
ns1) l2 :: [a]
l2@(a
n2 : [a]
ns2) =
  case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
n1 a
n2 of
    Ordering
LT -> Bool
False
    Ordering
EQ -> [a] -> [a] -> Bool
forall a. Ord a => [a] -> [a] -> Bool
ordSubsetp [a]
ns1 [a]
ns2
    Ordering
GT -> [a] -> [a] -> Bool
forall a. Ord a => [a] -> [a] -> Bool
ordSubsetp [a]
l1 [a]
ns2

-- * Even more loops

-- | Convert a list to a string, where the converter for each element
-- is a monadic computation.
formatList :: Monad m => String -> (a -> m String) -> [a] -> m String
formatList :: String -> (a -> m String) -> [a] -> m String
formatList String
s a -> m String
f [a]
xs = (a -> m String) -> [a] -> m [String]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM a -> m String
f [a]
xs m [String] -> ([String] -> m String) -> m String
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= String -> m String
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> m String)
-> ([String] -> String) -> [String] -> m String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
s

-- | Like `forM_`, but with both the elements source as well as the
-- loop body as computations over the monad.
forMM_ :: (Monad m, Foldable t) => m (t a) -> (a -> m ()) -> m ()
forMM_ :: m (t a) -> (a -> m ()) -> m ()
forMM_ m (t a)
srcM a -> m ()
f = do
  t a
src <- m (t a)
srcM
  t a -> (a -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ t a
src a -> m ()
f

-- | A @while@ loop, guard at the top.
whileDo :: Monad m => m Bool -> m () -> m ()
whileDo :: m Bool -> m () -> m ()
whileDo m Bool
cond m ()
body =
  m Bool
cond m Bool -> (Bool -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
b -> if Bool
b then m ()
body m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m Bool -> m () -> m ()
forall (m :: * -> *). Monad m => m Bool -> m () -> m ()
whileDo m Bool
cond m ()
body else () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | A @while@ loop based on stuff, guard at the top.
whileDoWith :: Monad m => m a -> (a -> Bool) -> (a -> m ()) -> m ()
whileDoWith :: m a -> (a -> Bool) -> (a -> m ()) -> m ()
whileDoWith m a
src a -> Bool
predicate a -> m ()
body = do
  a
val <- m a
src
  if a -> Bool
predicate a
val then (do a -> m ()
body a
val
                            m a -> (a -> Bool) -> (a -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
m a -> (a -> Bool) -> (a -> m ()) -> m ()
whileDoWith m a
src a -> Bool
predicate a -> m ()
body)
  else () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Like `forM_`, but with an extra check run after the body of the
-- loop.  If the check fails, the loop exits early.
forMwhile_ :: Monad m => [a] -> m Bool -> (a -> m ()) -> m ()
forMwhile_ :: [a] -> m Bool -> (a -> m ()) -> m ()
forMwhile_ [] m Bool
_ a -> m ()
_ = () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
forMwhile_ (a
x : [a]
xs) m Bool
pred a -> m ()
bodyf = do
  m Bool -> m () -> m ()
forall (m :: * -> *). Monad m => m Bool -> m () -> m ()
whenM m Bool
pred (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    a -> m ()
bodyf a
x
    [a] -> m Bool -> (a -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
[a] -> m Bool -> (a -> m ()) -> m ()
forMwhile_ [a]
xs m Bool
pred a -> m ()
bodyf

-- | Like `forMwhile_`, but the source list is also the result of a
-- monadic computation.
forMMwhile_ :: Monad m => m [a] -> m Bool -> (a -> m ()) -> m ()
forMMwhile_ :: m [a] -> m Bool -> (a -> m ()) -> m ()
forMMwhile_ m [a]
xsM m Bool
condM a -> m ()
bodyf = do
  [a]
xs <- m [a]
xsM
  [a] -> m Bool -> (a -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
[a] -> m Bool -> (a -> m ()) -> m ()
forMwhile_ [a]
xs m Bool
condM a -> m ()
bodyf

-- | Like `forMM_`, except instead of a fixed list, loop over `Maybe`
-- values returned from a subcomputation, until that subcomputation
-- returns `Nothing`.
whileReturnJust :: Monad m => m (Maybe a) -> (a -> m ()) -> m ()
whileReturnJust :: m (Maybe a) -> (a -> m ()) -> m ()
whileReturnJust m (Maybe a)
gen a -> m ()
f = do
  Maybe a
res <- m (Maybe a)
gen
  case Maybe a
res of
    Maybe a
Nothing -> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    Just a
x  -> do
      a -> m ()
f a
x
      m (Maybe a) -> (a -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
m (Maybe a) -> (a -> m ()) -> m ()
whileReturnJust m (Maybe a)
gen a -> m ()
f

-- | Like `unless`, expect both the tested value and the body are
-- returned from a computation in a monad.
unlessMM :: Monad m => m Bool -> m () -> m ()
unlessMM :: m Bool -> m () -> m ()
unlessMM m Bool
cnd m ()
body = m Bool -> m () -> m ()
forall (m :: * -> *). Monad m => m Bool -> m () -> m ()
whenM (m Bool -> m Bool
forall (m :: * -> *). Functor m => m Bool -> m Bool
notM m Bool
cnd) m ()
body

-- * Lists under references in the `STT` monad transformer

-- |Monadic version of @null@ for a list stored in an `STRef`: returns
-- `True` when the list is empty.
nullR :: Monad m => STRef s [a] -> STT s m Bool
nullR :: STRef s [a] -> STT s m Bool
nullR STRef s [a]
ref = do
  [a]
xs <- STRef s [a] -> STT s m [a]
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s [a]
ref
  Bool -> STT s m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> STT s m Bool) -> Bool -> STT s m Bool
forall a b. (a -> b) -> a -> b
$ [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs

-- |Opposite of `nullR`, returning `False` when the referenced list is
-- empty.
nonnullR :: Monad m => STRef s [a] -> STT s m Bool
nonnullR :: STRef s [a] -> STT s m Bool
nonnullR STRef s [a]
ref = do
  [a]
xs <- STRef s [a] -> STT s m [a]
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s [a]
ref
  Bool -> STT s m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> STT s m Bool) -> Bool -> STT s m Bool
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs

-- |Like a combination of `whenM` and `nonnullR`, where the body
-- receives the (pure) non-null list as an argument.
whenNonnullR :: (Monad m0, Monad m) =>
  (forall r . STT s m0 r -> m r) -> STRef s [a] -> ([a] -> m ()) -> m ()
whenNonnullR :: (forall r. STT s m0 r -> m r)
-> STRef s [a] -> ([a] -> m ()) -> m ()
whenNonnullR forall r. STT s m0 r -> m r
lifter STRef s [a]
ref [a] -> m ()
bodyf = do
  [a]
xs <- STT s m0 [a] -> m [a]
forall r. STT s m0 r -> m r
lifter (STT s m0 [a] -> m [a]) -> STT s m0 [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ STRef s [a] -> STT s m0 [a]
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s [a]
ref
  if ([a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs) then () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return () else [a] -> m ()
bodyf [a]
xs

-- |Map over the values contained within a list of references.
mapRefs :: Monad m => (a -> b) -> [STRef s a] -> STT s m [b]
mapRefs :: (a -> b) -> [STRef s a] -> STT s m [b]
mapRefs a -> b
f [] = [b] -> STT s m [b]
forall (m :: * -> *) a. Monad m => a -> m a
return []
mapRefs a -> b
f (STRef s a
xr : [STRef s a]
xrs) = do
  a
x <- STRef s a -> STT s m a
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s a
xr
  [b]
xs' <- (a -> b) -> [STRef s a] -> STT s m [b]
forall (m :: * -> *) a b s.
Monad m =>
(a -> b) -> [STRef s a] -> STT s m [b]
mapRefs a -> b
f [STRef s a]
xrs
  [b] -> STT s m [b]
forall (m :: * -> *) a. Monad m => a -> m a
return ([b] -> STT s m [b]) -> [b] -> STT s m [b]
forall a b. (a -> b) -> a -> b
$ a -> b
f a
x b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [b]
xs'

-- |Fold (right-associatively) the values contained within a list of
-- references.
foldrRefs :: Monad m => (a -> b -> b) -> b -> [STRef s a] -> STT s m b
foldrRefs :: (a -> b -> b) -> b -> [STRef s a] -> STT s m b
foldrRefs a -> b -> b
f b
z [] = b -> STT s m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
z
foldrRefs a -> b -> b
f b
z (STRef s a
xr : [STRef s a]
xrs) = do
  a
x <- STRef s a -> STT s m a
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s a
xr
  b
z' <- (a -> b -> b) -> b -> [STRef s a] -> STT s m b
forall (m :: * -> *) a b s.
Monad m =>
(a -> b -> b) -> b -> [STRef s a] -> STT s m b
foldrRefs a -> b -> b
f b
z [STRef s a]
xrs
  b -> STT s m b
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> STT s m b) -> b -> STT s m b
forall a b. (a -> b) -> a -> b
$ a -> b -> b
f a
x b
z'

-- |Fold (left-associatively) the values contained within a list of
-- references.
foldlRefs :: Monad m => (b -> a -> b) -> b -> [STRef s a] -> STT s m b
foldlRefs :: (b -> a -> b) -> b -> [STRef s a] -> STT s m b
foldlRefs b -> a -> b
f b
z [] = b -> STT s m b
forall (m :: * -> *) a. Monad m => a -> m a
return b
z
foldlRefs b -> a -> b
f b
z (STRef s a
xr : [STRef s a]
xrs) = do
  a
x <- STRef s a -> STT s m a
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s a
xr
  (b -> a -> b) -> b -> [STRef s a] -> STT s m b
forall (m :: * -> *) b a s.
Monad m =>
(b -> a -> b) -> b -> [STRef s a] -> STT s m b
foldlRefs b -> a -> b
f (b -> a -> b
f b
z a
x) [STRef s a]
xrs

-- | Like `forM_`, but with the list under an `STRef`.  The first
-- argument lifts an `STT` operation into @m@.
forRM_ ::
  (Monad m, Monad m0, Foldable t) =>
    (STT s m0 (t a) -> m (t a)) ->
      STRef s (t a) -> (a -> m ()) -> m ()
forRM_ :: (STT s m0 (t a) -> m (t a)) -> STRef s (t a) -> (a -> m ()) -> m ()
forRM_ STT s m0 (t a) -> m (t a)
liftSTT STRef s (t a)
srcR a -> m ()
f = do
  t a
src <- STT s m0 (t a) -> m (t a)
liftSTT (STT s m0 (t a) -> m (t a)) -> STT s m0 (t a) -> m (t a)
forall a b. (a -> b) -> a -> b
$ STRef s (t a) -> STT s m0 (t a)
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s (t a)
srcR
  t a -> (a -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ t a
src a -> m ()
f

-- ** Stack-like operations

-- |Push a value onto the front of the list at the given `STT`
-- reference.
push :: Monad m => a -> STRef s [a] -> STT s m ()
push :: a -> STRef s [a] -> STT s m ()
push a
v STRef s [a]
r = do
  [a]
prev <- STRef s [a] -> STT s m [a]
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s [a]
r
  STRef s [a] -> [a] -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s [a]
r ([a] -> STT s m ()) -> [a] -> STT s m ()
forall a b. (a -> b) -> a -> b
$ a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
prev

-- |Push the result of a computation onto the front of the list at the
-- given `STT` reference.
pushM :: Monad m => m a -> STRef s [a] -> STT s m ()
pushM :: m a -> STRef s [a] -> STT s m ()
pushM m a
m STRef s [a]
r = do
  a
v <- m a -> STT s m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m a
m
  a -> STRef s [a] -> STT s m ()
forall (m :: * -> *) a s. Monad m => a -> STRef s [a] -> STT s m ()
push a
v STRef s [a]
r

-- |Push every value in a collection onto the front of the list at the
-- given `STT` reference.
pushAll :: (Monad m, Traversable t) => t a -> STRef s [a] -> STT s m ()
pushAll :: t a -> STRef s [a] -> STT s m ()
pushAll t a
vs STRef s [a]
r = t a -> (a -> STT s m ()) -> STT s m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ t a
vs ((a -> STT s m ()) -> STT s m ())
-> (a -> STT s m ()) -> STT s m ()
forall a b. (a -> b) -> a -> b
$ \a
v -> a -> STRef s [a] -> STT s m ()
forall (m :: * -> *) a s. Monad m => a -> STRef s [a] -> STT s m ()
push a
v STRef s [a]
r

-- |Push every value in a collection returned from a computation onto
-- the front of the list at the given `STT` reference.
pushAllM :: (Monad m, Traversable t) => m (t a) -> STRef s [a] -> STT s m ()
pushAllM :: m (t a) -> STRef s [a] -> STT s m ()
pushAllM m (t a)
m STRef s [a]
r = do
  t a
vs <- m (t a) -> STT s m (t a)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (t a)
m
  t a -> STRef s [a] -> STT s m ()
forall (m :: * -> *) (t :: * -> *) a s.
(Monad m, Traversable t) =>
t a -> STRef s [a] -> STT s m ()
pushAll t a
vs STRef s [a]
r

-- |Pop a value from the given reference to a list if one exists.
pop :: Monad m => STRef s [a] -> STT s m (Maybe a)
pop :: STRef s [a] -> STT s m (Maybe a)
pop STRef s [a]
queue = do
  [a]
queueList <- STRef s [a] -> STT s m [a]
forall (m :: * -> *) s a. Applicative m => STRef s a -> STT s m a
readSTRef STRef s [a]
queue
  case [a]
queueList of
    [] -> Maybe a -> STT s m (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
    (a
x : [a]
xs) -> do
      STRef s [a] -> [a] -> STT s m ()
forall (m :: * -> *) s a.
Applicative m =>
STRef s a -> a -> STT s m ()
writeSTRef STRef s [a]
queue [a]
xs
      Maybe a -> STT s m (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> STT s m (Maybe a)) -> Maybe a -> STT s m (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
x

-- |Consumes the elements of a referenced list, one at a time, until
-- the list is empty.  The first argument is a @lift@-style function
-- which brings `STT` operations into the top-level monad of interest.
-- Intended to be compatible with stack-like behavior (such as with
-- `push`; this function does use `pop`) where the body of the loop
-- may add elements.
whileListM_ :: (Monad m0, Monad m) =>
  (forall r . STT s m0 r -> m r) -> STRef s [a] -> (a -> m ()) -> m ()
whileListM_ :: (forall r. STT s m0 r -> m r) -> STRef s [a] -> (a -> m ()) -> m ()
whileListM_ forall r. STT s m0 r -> m r
lifter STRef s [a]
listRef a -> m ()
bodyf = m ()
whileListM_'
  where whileListM_' :: m ()
whileListM_' = do
          Maybe a
top <- STT s m0 (Maybe a) -> m (Maybe a)
forall r. STT s m0 r -> m r
lifter (STT s m0 (Maybe a) -> m (Maybe a))
-> STT s m0 (Maybe a) -> m (Maybe a)
forall a b. (a -> b) -> a -> b
$ STRef s [a] -> STT s m0 (Maybe a)
forall (m :: * -> *) s a.
Monad m =>
STRef s [a] -> STT s m (Maybe a)
pop STRef s [a]
listRef
          case Maybe a
top of
            Maybe a
Nothing -> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            Just a
x -> do
              a -> m ()
bodyf a
x
              m ()
whileListM_'

-- * Strings

-- |Form a comma-separated string from a list.
commaList :: (a -> String) -> [a] -> String
commaList :: (a -> String) -> [a] -> String
commaList a -> String
f [] = String
""
commaList a -> String
f [a]
xs = (String -> String -> String) -> [String] -> String
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 (\ String
x String
y -> String
x String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
", " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
y) ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (a -> String) -> [a] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map a -> String
f [a]
xs

-- |Remove the `Just` constructors from the elements of a list,
-- discarding elements which are `Nothing`.
unmaybe :: [Maybe a] -> [a]
unmaybe [] = []
unmaybe (Just a
a : [Maybe a]
xs) = a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [Maybe a] -> [a]
unmaybe [Maybe a]
xs
unmaybe (Maybe a
_ : [Maybe a]
xs) = [Maybe a] -> [a]
unmaybe [Maybe a]
xs