-- | Type-indexed fragment store for collecting and sampling fragments from.
module Test.Mutagen.Fragment.Store
  ( -- * Type-indexed fragment store
    FragmentStore (..)
  , emptyFragmentStore
  , fragmentStoreSize
  , storeFragments
  , sampleFragments
  , printFragmentStore
  , FragmentTypeFilter (..)
  , isFragmentTypeAllowed
  )
where

import Control.Monad (forM_)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (mapMaybe)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Typeable (TypeRep, Typeable, cast, typeOf)
import Test.Mutagen.Fragment (Fragment (..), Fragmentable (..), singleton)
import Test.QuickCheck (Gen, shuffle)

{-------------------------------------------------------------------------------
-- * Type-indexed fragment store
-------------------------------------------------------------------------------}

-- | A collection of fragments indexed by their type representation.
newtype FragmentStore = FragmentStore (Map TypeRep (Set Fragment))

instance Semigroup FragmentStore where
  FragmentStore Map TypeRep (Set Fragment)
fs1 <> :: FragmentStore -> FragmentStore -> FragmentStore
<> FragmentStore Map TypeRep (Set Fragment)
fs2 =
    Map TypeRep (Set Fragment) -> FragmentStore
FragmentStore ((Set Fragment -> Set Fragment -> Set Fragment)
-> Map TypeRep (Set Fragment)
-> Map TypeRep (Set Fragment)
-> Map TypeRep (Set Fragment)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith Set Fragment -> Set Fragment -> Set Fragment
forall a. Ord a => Set a -> Set a -> Set a
Set.union Map TypeRep (Set Fragment)
fs1 Map TypeRep (Set Fragment)
fs2)

instance Monoid FragmentStore where
  mempty :: FragmentStore
mempty = FragmentStore
emptyFragmentStore

-- | An empty fragment store.
emptyFragmentStore :: FragmentStore
emptyFragmentStore :: FragmentStore
emptyFragmentStore = Map TypeRep (Set Fragment) -> FragmentStore
FragmentStore Map TypeRep (Set Fragment)
forall a. Monoid a => a
mempty

-- | Get the number of fragments stored for each type.
fragmentStoreSize :: FragmentStore -> [(TypeRep, Int)]
fragmentStoreSize :: FragmentStore -> [(TypeRep, Int)]
fragmentStoreSize (FragmentStore Map TypeRep (Set Fragment)
fs) =
  [ (TypeRep
tyRep, Set Fragment -> Int
forall a. Set a -> Int
Set.size Set Fragment
frags)
  | (TypeRep
tyRep, Set Fragment
frags) <- Map TypeRep (Set Fragment) -> [(TypeRep, Set Fragment)]
forall k a. Map k a -> [(k, a)]
Map.toList Map TypeRep (Set Fragment)
fs
  ]

-- | Store fragments from a value into the fragment store.
storeFragments
  :: (Fragmentable a)
  => FragmentTypeFilter
  -> a
  -> FragmentStore
  -> FragmentStore
storeFragments :: forall a.
Fragmentable a =>
FragmentTypeFilter -> a -> FragmentStore -> FragmentStore
storeFragments FragmentTypeFilter
typeFilter a
a (FragmentStore Map TypeRep (Set Fragment)
store) =
  Map TypeRep (Set Fragment) -> FragmentStore
FragmentStore ((Set Fragment -> Set Fragment -> Set Fragment)
-> Map TypeRep (Set Fragment)
-> Map TypeRep (Set Fragment)
-> Map TypeRep (Set Fragment)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith Set Fragment -> Set Fragment -> Set Fragment
forall a. Ord a => Set a -> Set a -> Set a
Set.union Map TypeRep (Set Fragment)
store (a -> Map TypeRep (Set Fragment)
collect a
a))
  where
    collect :: a -> Map TypeRep (Set Fragment)
collect = (Fragment
 -> Map TypeRep (Set Fragment) -> Map TypeRep (Set Fragment))
-> Map TypeRep (Set Fragment)
-> Set Fragment
-> Map TypeRep (Set Fragment)
forall a b. (a -> b -> b) -> b -> Set a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Fragment
-> Map TypeRep (Set Fragment) -> Map TypeRep (Set Fragment)
insertIfAllowed Map TypeRep (Set Fragment)
forall k a. Map k a
Map.empty (Set Fragment -> Map TypeRep (Set Fragment))
-> (a -> Set Fragment) -> a -> Map TypeRep (Set Fragment)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Set Fragment
forall a. Fragmentable a => a -> Set Fragment
fragmentize

    insertIfAllowed :: Fragment
-> Map TypeRep (Set Fragment) -> Map TypeRep (Set Fragment)
insertIfAllowed (Fragment a
a') Map TypeRep (Set Fragment)
store'
      | FragmentTypeFilter -> TypeRep -> Bool
isFragmentTypeAllowed FragmentTypeFilter
typeFilter (a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
a') =
          (Set Fragment -> Set Fragment -> Set Fragment)
-> TypeRep
-> Set Fragment
-> Map TypeRep (Set Fragment)
-> Map TypeRep (Set Fragment)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith Set Fragment -> Set Fragment -> Set Fragment
forall a. Ord a => Set a -> Set a -> Set a
Set.union (a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
a') (a -> Set Fragment
forall a. IsFragment a => a -> Set Fragment
singleton a
a') Map TypeRep (Set Fragment)
store'
      | Bool
otherwise =
          Map TypeRep (Set Fragment)
store'

-- | Sample fragments of the same type as the given value from a fragment store.
sampleFragments
  :: (Typeable a)
  => a
  -> FragmentStore
  -> Gen [a]
sampleFragments :: forall a. Typeable a => a -> FragmentStore -> Gen [a]
sampleFragments a
a (FragmentStore Map TypeRep (Set Fragment)
store) = do
  case TypeRep -> Map TypeRep (Set Fragment) -> Maybe (Set Fragment)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (a -> TypeRep
forall a. Typeable a => a -> TypeRep
typeOf a
a) Map TypeRep (Set Fragment)
store of
    Maybe (Set Fragment)
Nothing ->
      [a] -> Gen [a]
forall a. a -> Gen a
forall (m :: * -> *) a. Monad m => a -> m a
return []
    Just Set Fragment
frags ->
      (Fragment -> Maybe a) -> [Fragment] -> [a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(Fragment a
a') -> a -> Maybe a
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast a
a') ([Fragment] -> [a]) -> Gen [Fragment] -> Gen [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Fragment] -> Gen [Fragment]
forall a. [a] -> Gen [a]
shuffle (Set Fragment -> [Fragment]
forall a. Set a -> [a]
Set.toList Set Fragment
frags)

-- | Print the contents of a fragment store for debugging purposes.
printFragmentStore :: FragmentStore -> IO ()
printFragmentStore :: FragmentStore -> IO ()
printFragmentStore (FragmentStore Map TypeRep (Set Fragment)
fs) = do
  [(TypeRep, Set Fragment)]
-> ((TypeRep, Set Fragment) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Map TypeRep (Set Fragment) -> [(TypeRep, Set Fragment)]
forall k a. Map k a -> [(k, a)]
Map.assocs Map TypeRep (Set Fragment)
fs) (((TypeRep, Set Fragment) -> IO ()) -> IO ())
-> ((TypeRep, Set Fragment) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(TypeRep
tyRep, Set Fragment
frags) -> do
    String -> IO ()
putStrLn (String
"TypeRep: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> TypeRep -> String
forall a. Show a => a -> String
show TypeRep
tyRep)
    Set Fragment -> (Fragment -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Set Fragment
frags ((Fragment -> IO ()) -> IO ()) -> (Fragment -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Fragment
frag -> do
      String -> IO ()
putStrLn (String
"* " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Fragment -> String
forall a. Show a => a -> String
show Fragment
frag)

-- ** Fragment type filters

-- | Fragment type allow and deny lists.
data FragmentTypeFilter = FragmentTypeFilter
  { FragmentTypeFilter -> Set TypeRep
allowList :: Set TypeRep
  -- ^ List of allowed fragment types
  , FragmentTypeFilter -> Set TypeRep
denyList :: Set TypeRep
  -- ^ List of denied fragment types
  }
  deriving (FragmentTypeFilter -> FragmentTypeFilter -> Bool
(FragmentTypeFilter -> FragmentTypeFilter -> Bool)
-> (FragmentTypeFilter -> FragmentTypeFilter -> Bool)
-> Eq FragmentTypeFilter
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: FragmentTypeFilter -> FragmentTypeFilter -> Bool
== :: FragmentTypeFilter -> FragmentTypeFilter -> Bool
$c/= :: FragmentTypeFilter -> FragmentTypeFilter -> Bool
/= :: FragmentTypeFilter -> FragmentTypeFilter -> Bool
Eq, Int -> FragmentTypeFilter -> String -> String
[FragmentTypeFilter] -> String -> String
FragmentTypeFilter -> String
(Int -> FragmentTypeFilter -> String -> String)
-> (FragmentTypeFilter -> String)
-> ([FragmentTypeFilter] -> String -> String)
-> Show FragmentTypeFilter
forall a.
(Int -> a -> String -> String)
-> (a -> String) -> ([a] -> String -> String) -> Show a
$cshowsPrec :: Int -> FragmentTypeFilter -> String -> String
showsPrec :: Int -> FragmentTypeFilter -> String -> String
$cshow :: FragmentTypeFilter -> String
show :: FragmentTypeFilter -> String
$cshowList :: [FragmentTypeFilter] -> String -> String
showList :: [FragmentTypeFilter] -> String -> String
Show)

instance Semigroup FragmentTypeFilter where
  (FragmentTypeFilter Set TypeRep
a1 Set TypeRep
d1) <> :: FragmentTypeFilter -> FragmentTypeFilter -> FragmentTypeFilter
<> (FragmentTypeFilter Set TypeRep
a2 Set TypeRep
d2) =
    Set TypeRep -> Set TypeRep -> FragmentTypeFilter
FragmentTypeFilter (Set TypeRep
a1 Set TypeRep -> Set TypeRep -> Set TypeRep
forall a. Semigroup a => a -> a -> a
<> Set TypeRep
a2) (Set TypeRep
d1 Set TypeRep -> Set TypeRep -> Set TypeRep
forall a. Semigroup a => a -> a -> a
<> Set TypeRep
d2)

instance Monoid FragmentTypeFilter where
  mempty :: FragmentTypeFilter
mempty = Set TypeRep -> Set TypeRep -> FragmentTypeFilter
FragmentTypeFilter Set TypeRep
forall a. Monoid a => a
mempty Set TypeRep
forall a. Monoid a => a
mempty

-- | Check if a type is allowed by the fragment type filter.
isFragmentTypeAllowed :: FragmentTypeFilter -> TypeRep -> Bool
isFragmentTypeAllowed :: FragmentTypeFilter -> TypeRep -> Bool
isFragmentTypeAllowed (FragmentTypeFilter Set TypeRep
allow Set TypeRep
deny) TypeRep
tr =
  (TypeRep
tr TypeRep -> Set TypeRep -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set TypeRep
allow)
    Bool -> Bool -> Bool
&& Bool -> Bool
not (TypeRep
tr TypeRep -> Set TypeRep -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set TypeRep
deny)