{-# 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 = IterStat
forall a. Monoid a => a
mempty

joinStats :: IterStat -> IterStat -> IterStat
joinStats :: IterStat -> IterStat -> IterStat
joinStats = (Int -> Int -> Int) -> IterStat -> IterStat -> IterStat
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith Int -> Int -> Int
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 (IterStat -> IterStat) -> LogicT (StateT IterStat Identity) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((IterStat -> IterStat) -> LogicT (StateT IterStat Identity) ())
-> (IterStat -> IterStat) -> LogicT (StateT IterStat Identity) ()
forall a b. (a -> b) -> a -> b
$ IterStat -> IterStat -> IterStat
joinStats IterStat
stats
                                 a -> LogicI a
forall a. a -> LogicT (StateT IterStat Identity) a
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 = Identity ([a], IterStat) -> ([a], IterStat)
forall a. Identity a -> a
runIdentity (Identity ([a], IterStat) -> ([a], IterStat))
-> Identity ([a], IterStat) -> ([a], IterStat)
forall a b. (a -> b) -> a -> b
$ StateT IterStat Identity [a]
-> IterStat -> Identity ([a], IterStat)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (LogicI a -> StateT IterStat Identity [a]
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 = Identity ([a], IterStat) -> ([a], IterStat)
forall a. Identity a -> a
runIdentity (Identity ([a], IterStat) -> ([a], IterStat))
-> Identity ([a], IterStat) -> ([a], IterStat)
forall a b. (a -> b) -> a -> b
$ StateT IterStat Identity [a]
-> IterStat -> Identity ([a], IterStat)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (Int -> LogicI a -> StateT IterStat Identity [a]
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 (Map Text a -> Map Text a) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((Map Text a -> Map Text a) -> m ())
-> (Map Text a -> Map Text a) -> m ()
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> Text -> a -> Map Text a -> Map Text a
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith a -> a -> a
forall a. Num a => a -> a -> a
(+) Text
location a
1
                       b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return b
c m b -> m b -> m b
forall a. m a -> m a -> m a
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` m b
a
  in (a -> LogicI a -> LogicI a) -> LogicI a -> [a] -> LogicI a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> LogicI a -> LogicI a
forall {m :: * -> *} {a} {b}.
(MonadState (Map Text a) m, Num a, MonadPlus m) =>
b -> m b -> m b
attempt LogicI a
forall a. LogicT (StateT IterStat Identity) a
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 = [[a]] -> [[a]]
forall a. [a] -> [a]
reverse
                    ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> [a] -> Ordering) -> [[a]] -> [[a]]
forall a. (a -> a -> Ordering) -> [a] -> [a]
DL.sortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering) -> ([a] -> Int) -> [a] -> [a] -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length)
                    ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> [[a]]
forall a. Eq a => [a] -> [a]
DL.nub
                    ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> [[a]]) -> [[a]] -> [[a]]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap [a] -> [[a]]
forall a. [a] -> [[a]]
DL.inits
                    ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [[a]]
forall a. [a] -> [[a]]
DL.permutations