{-# LANGUAGE FlexibleContexts #-}

module Test.Tasty.Sugar.Iterations
where

import           Control.Monad ( mplus, mzero )
import           Control.Monad.Logic
import           Control.Monad.State ( StateT, runStateT, modify )
import           Data.Function ( on )
import           Data.Functor.Identity ( Identity, runIdentity )
import qualified Data.List as DL
import qualified Data.Map as Map
import           Data.Text ( Text )


type IterStat = Map.Map Text Int

emptyStats :: IterStat
emptyStats :: IterStat
emptyStats = forall a. Monoid a => a
mempty

joinStats :: IterStat -> IterStat -> IterStat
joinStats :: IterStat -> IterStat -> IterStat
joinStats = forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith forall a. Num a => a -> a -> a
(+)


----------------------------------------------------------------------

type LogicI a = LogicT (StateT IterStat Identity) a

-- Note: stats collection can increase the runtime if the amount of backtracking
-- becomes significant.  It can also increase runtime because a stats report
-- (from --showsearch) will force the evaluation of branches that might otherwise
-- have been lazily ignored.  To disable the dilatory effects (and disable stats
-- collection), disable the modify statements in addSubLogicStats and eachFrom.

addSubLogicStats :: (a, IterStat) -> LogicI a
addSubLogicStats :: forall a. (a, IterStat) -> LogicI a
addSubLogicStats (a
r, IterStat
stats) = do forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ IterStat -> IterStat -> IterStat
joinStats IterStat
stats
                                 forall (m :: * -> *) a. Monad m => a -> m a
return a
r

observeIAll :: LogicI a -> ([a], IterStat)
observeIAll :: forall a. LogicI a -> ([a], IterStat)
observeIAll LogicI a
op = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall (m :: * -> *) a. Applicative m => LogicT m a -> m [a]
observeAllT LogicI a
op) IterStat
emptyStats

observeIT :: LogicI a -> ([a], IterStat)
observeIT :: forall a. LogicI a -> ([a], IterStat)
observeIT LogicI a
op = forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall (m :: * -> *) a. Monad m => Int -> LogicT m a -> m [a]
observeManyT Int
1 LogicI a
op) IterStat
emptyStats

----------------------------------------------------------------------

-- | Core Logic function to iteratively return elements of a list via
-- backtracking.
eachFrom :: Text -> [a] -> LogicI a
eachFrom :: forall a. Text -> [a] -> LogicI a
eachFrom Text
location =
  let attempt :: b -> m b -> m b
attempt b
c m b
a = do forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith forall a. Num a => a -> a -> a
(+) Text
location a
1
                       forall (m :: * -> *) a. Monad m => a -> m a
return b
c forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` m b
a
  in forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall {m :: * -> *} {a} {b}.
(MonadState (Map Text a) m, Num a, MonadPlus m) =>
b -> m b -> m b
attempt forall (m :: * -> *) a. MonadPlus m => m a
mzero


-- | Given a list, return the list of lists representing all permutations of the
-- same length or shorter, removing any duplications, from longest to shortest
-- (shortest being the empty list).
combosLongToShort :: Eq a => [a] -> [ [a] ]
combosLongToShort :: forall a. Eq a => [a] -> [[a]]
combosLongToShort = forall a. [a] -> [a]
reverse
                    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Ordering) -> [a] -> [a]
DL.sortBy (forall a. Ord a => a -> a -> Ordering
compare forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall (t :: * -> *) a. Foldable t => t a -> Int
length)
                    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Eq a => [a] -> [a]
DL.nub
                    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. [a] -> [[a]]
DL.inits
                    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [[a]]
DL.permutations