{-# OPTIONS_GHC -Wunused-imports #-}

{-# LANGUAGE RebindableSyntax #-}

module Agda.Utils.Function
  ( module Agda.Utils.Function
  , module Data.Function
  ) where

import Prelude hiding ( not, (&&), (||) )

import Data.Function  ( on )
import Data.String    ( fromString )       -- for RebindableSyntax, somehow not covered by Prelude

import Agda.Utils.Boolean

-- | Repeat a state transition @f :: a -> (b, a)@ with output @b@
--   while condition @cond@ on the output is true.
--   Return all intermediate results and the final result
--   where @cond@ is @False@.
--
--   Postconditions (when it terminates):
--   @fst (last (iterWhile cond f a)) == False@.
--   @all fst (init (interWhile cond f a))@.

iterWhile :: (b -> Bool) -> (a -> (b, a)) -> a -> [(b,a)]
iterWhile :: forall b a. (b -> Bool) -> (a -> (b, a)) -> a -> [(b, a)]
iterWhile b -> Bool
cond a -> (b, a)
f = a -> [(b, a)]
loop where
  loop :: a -> [(b, a)]
loop a
a = (b, a)
r (b, a) -> [(b, a)] -> [(b, a)]
forall a. a -> [a] -> [a]
: if b -> Bool
cond b
b then a -> [(b, a)]
loop a
a' else []
    where r :: (b, a)
r@(b
b, a
a') = a -> (b, a)
f a
a

-- | Repeat something while a condition on some state is true.
--   Return the last state (including the changes of the last
--   transition, even if the condition became false then).

repeatWhile :: (a -> (Bool, a)) -> a -> a
repeatWhile :: forall a. (a -> (Bool, a)) -> a -> a
repeatWhile a -> (Bool, a)
f = a -> a
loop where
  loop :: a -> a
loop a
a = if Bool
again then a -> a
loop a
a' else a
a'
    where (Bool
again, a
a') = a -> (Bool, a)
f a
a

-- | Monadic version of 'repeatWhile'.
repeatWhileM :: (Monad m) => (a -> m (Bool, a)) -> a -> m a
repeatWhileM :: forall (m :: * -> *) a. Monad m => (a -> m (Bool, a)) -> a -> m a
repeatWhileM a -> m (Bool, a)
f = a -> m a
loop where
  loop :: a -> m a
loop a
a = do
    (Bool
again, a
a') <- a -> m (Bool, a)
f a
a
    if Bool
again then a -> m a
loop a
a' else a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a'

-- | A version of the trampoline function.
--
--   The usual function iterates @f :: a -> Maybe a@ as long
--   as @Just{}@ is returned, and returns the last value of @a@
--   upon @Nothing@.
--
--   @usualTrampoline f = trampolineWhile $ \ a -> maybe (False,a) (True,) (f a)@.
--
--   @trampolineWhile@ is very similar to @repeatWhile@, only that
--   it discards the state on which the condition went @False@,
--   and returns the last state on which the condition was @True@.
trampolineWhile :: (a -> (Bool, a)) -> a -> a
trampolineWhile :: forall a. (a -> (Bool, a)) -> a -> a
trampolineWhile a -> (Bool, a)
f = (a -> (Bool, a)) -> a -> a
forall a. (a -> (Bool, a)) -> a -> a
repeatWhile ((a -> (Bool, a)) -> a -> a) -> (a -> (Bool, a)) -> a -> a
forall a b. (a -> b) -> a -> b
$ \ a
a ->
  let (Bool
again, a
a') = a -> (Bool, a)
f a
a
  in (Bool
again,) (a -> (Bool, a)) -> a -> (Bool, a)
forall a b. (a -> b) -> a -> b
$ if Bool
again then a
a' else a
a

-- | Monadic version of 'trampolineWhile'.
trampolineWhileM :: (Monad m) => (a -> m (Bool, a)) -> a -> m a
trampolineWhileM :: forall (m :: * -> *) a. Monad m => (a -> m (Bool, a)) -> a -> m a
trampolineWhileM a -> m (Bool, a)
f = (a -> m (Bool, a)) -> a -> m a
forall (m :: * -> *) a. Monad m => (a -> m (Bool, a)) -> a -> m a
repeatWhileM ((a -> m (Bool, a)) -> a -> m a) -> (a -> m (Bool, a)) -> a -> m a
forall a b. (a -> b) -> a -> b
$ \ a
a -> do
  (Bool
again, a
a') <- a -> m (Bool, a)
f a
a
  (Bool, a) -> m (Bool, a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Bool, a) -> m (Bool, a)) -> (Bool, a) -> m (Bool, a)
forall a b. (a -> b) -> a -> b
$ (Bool
again,) (a -> (Bool, a)) -> a -> (Bool, a)
forall a b. (a -> b) -> a -> b
$ if Bool
again then a
a' else a
a

-- | More general trampoline, which allows some final computation
--   from iteration state @a@ into result type @b@.
trampoline :: (a -> Either b a) -> a -> b
trampoline :: forall a b. (a -> Either b a) -> a -> b
trampoline a -> Either b a
f = a -> b
loop where
  loop :: a -> b
loop a
a = (b -> b) -> (a -> b) -> Either b a -> b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either b -> b
forall a. a -> a
id a -> b
loop (Either b a -> b) -> Either b a -> b
forall a b. (a -> b) -> a -> b
$ a -> Either b a
f a
a

-- | Monadic version of 'trampoline'.
trampolineM :: Monad m => (a -> m (Either b a)) -> a -> m b
trampolineM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m (Either b a)) -> a -> m b
trampolineM a -> m (Either b a)
f = a -> m b
loop where
  loop :: a -> m b
loop a
a = (b -> m b) -> (a -> m b) -> Either b a -> m b
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a -> m b
loop (Either b a -> m b) -> m (Either b a) -> m b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< a -> m (Either b a)
f a
a

-- | Iteration to fixed-point.
--
--   @iterateUntil r f a0@ iterates endofunction @f@, starting with @a0@,
--   until @r@ relates its result to its input, i.e., @f a `r` a@.
--
--   This is the generic pattern behind saturation algorithms.
--
--   If @f@ is monotone with regard to @r@,
--   meaning @a `r` b@ implies @f a `r` f b@,
--   and @f@-chains starting with @a0@ are finite
--   then iteration is guaranteed to terminate.
--
--   A typical instance will work on sets, and @r@ could be set inclusion,
--   and @a0@ the empty set, and @f@ the step function of a saturation algorithm.
iterateUntil :: (a -> a -> Bool) -> (a -> a) -> a -> a
iterateUntil :: forall a. (a -> a -> Bool) -> (a -> a) -> a -> a
iterateUntil a -> a -> Bool
r a -> a
f = a -> a
loop where
  loop :: a -> a
loop a
a = if a -> a -> Bool
r a
a' a
a then a
a' else a -> a
loop a
a'
    where a' :: a
a' = a -> a
f a
a

-- | Monadic version of 'iterateUntil'.
iterateUntilM :: Monad m => (a -> a -> Bool) -> (a -> m a) -> a -> m a
iterateUntilM :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> Bool) -> (a -> m a) -> a -> m a
iterateUntilM a -> a -> Bool
r a -> m a
f = a -> m a
loop where
  loop :: a -> m a
loop a
a = do
    a
a' <- a -> m a
f a
a
    if a -> a -> Bool
r a
a' a
a then a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a' else a -> m a
loop a
a'

-- | @'iterate'' n f x@ applies @f@ to @x@ @n@ times and returns the
-- result.
--
-- The applications are calculated strictly.

iterate' :: Integral i => i -> (a -> a) -> a -> a
iterate' :: forall i a. Integral i => i -> (a -> a) -> a -> a
iterate' i
0 a -> a
_ a
x             = a
x
iterate' i
n a -> a
f a
x | i
n i -> i -> Bool
forall a. Ord a => a -> a -> Bool
> i
0     = i -> (a -> a) -> a -> a
forall i a. Integral i => i -> (a -> a) -> a -> a
iterate' (i
n i -> i -> i
forall a. Num a => a -> a -> a
- i
1) a -> a
f (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$! a -> a
f a
x
               | Bool
otherwise = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"iterate': Negative input."

-- * Iteration over Booleans.

-- | @applyWhen b f a@ applies @f@ to @a@ when @b@.
{-# SPECIALIZE applyWhen :: Bool -> (a -> a) -> (a -> a) #-}
{-# INLINE applyWhen #-}
applyWhen :: IsBool b => b -> (a -> a) -> a -> a
applyWhen :: forall b a. IsBool b => b -> (a -> a) -> a -> a
applyWhen b
b a -> a
f = if b
b then a -> a
f else a -> a
forall a. a -> a
id
  -- Note: RebindableSyntax translates this if-then-else to ifThenElse of IsBool.

-- | @applyUnless b f a@ applies @f@ to @a@ unless @b@.
{-# SPECIALIZE applyUnless :: Bool -> (a -> a) -> (a -> a) #-}
{-# INLINE applyUnless #-}
applyUnless :: IsBool b => b -> (a -> a) -> a -> a
applyUnless :: forall b a. IsBool b => b -> (a -> a) -> a -> a
applyUnless b
b a -> a
f = if b
b then a -> a
forall a. a -> a
id else a -> a
f

-- | Monadic version of @applyWhen@
{-# SPECIALIZE applyWhenM :: Monad m => m Bool -> (m a -> m a) -> m a -> m a #-}
{-# INLINE applyWhenM #-}
applyWhenM :: (IsBool b, Monad m) => m b -> (m a -> m a) -> m a -> m a
applyWhenM :: forall b (m :: * -> *) a.
(IsBool b, Monad m) =>
m b -> (m a -> m a) -> m a -> m a
applyWhenM m b
mb m a -> m a
f m a
x = m b
mb m b -> (b -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ b
b -> b -> (m a -> m a) -> m a -> m a
forall b a. IsBool b => b -> (a -> a) -> a -> a
applyWhen b
b m a -> m a
f m a
x

-- | Monadic version of @applyUnless@
{-# SPECIALIZE applyUnlessM :: Monad m => m Bool -> (m a -> m a) -> m a -> m a #-}
{-# INLINE applyUnlessM #-}
applyUnlessM :: (IsBool b, Monad m) => m b -> (m a -> m a) -> m a -> m a
applyUnlessM :: forall b (m :: * -> *) a.
(IsBool b, Monad m) =>
m b -> (m a -> m a) -> m a -> m a
applyUnlessM m b
mb m a -> m a
f m a
x = m b
mb m b -> (b -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ b
b -> b -> (m a -> m a) -> m a -> m a
forall b a. IsBool b => b -> (a -> a) -> a -> a
applyUnless b
b m a -> m a
f m a
x

-- | 'Maybe' version of 'applyWhen'.
{-# INLINE applyWhenJust #-}
applyWhenJust :: Maybe b -> (b -> a -> a) -> a -> a
applyWhenJust :: forall b a. Maybe b -> (b -> a -> a) -> a -> a
applyWhenJust Maybe b
m b -> a -> a
f = (a -> a) -> (b -> a -> a) -> Maybe b -> a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a -> a
forall a. a -> a
id b -> a -> a
f Maybe b
m

-- | 'Maybe' version of 'applyUnless'.
{-# INLINE applyWhenNothing #-}
applyWhenNothing :: Maybe b -> (a -> a) -> a -> a
applyWhenNothing :: forall b a. Maybe b -> (a -> a) -> a -> a
applyWhenNothing Maybe b
m a -> a
f = (a -> a) -> (b -> a -> a) -> Maybe b -> a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a -> a
f ((a -> a) -> b -> a -> a
forall a b. a -> b -> a
const a -> a
forall a. a -> a
id) Maybe b
m