{-#LANGUAGE DeriveDataTypeable, TemplateHaskell #-}

-- | Basic combinators for building enumerations
-- most users will want to use the type class
-- based combinators in "Test.Feat.Class" instead.

module Test.Feat.Enumerate (


  Index,
  Enumerate(..),
  parts,
  fromParts,

  -- ** Reversed lists
  RevList(..),
  toRev,

  -- ** Finite ordered sets
  Finite(..),
  fromFinite,


  -- ** Combinators for building enumerations
  module Data.Monoid,
  union,
  module Control.Applicative,
  cartesian,
  singleton,
  pay,
  ) where

-- testing-feat
-- import Control.Monad.TagShare(Sharing, runSharing, share)
-- import Test.Feat.Internals.Tag(Tag(Source))
-- base
import Control.Sized
import Control.Applicative
import Data.Semigroup
import Data.Monoid hiding ((<>))
import Data.Typeable
import Data.List(transpose)
import Test.Feat.Finite

-- | A functional enumeration of type @t@ is a partition of
-- @t@ into finite numbered sets. Each part contains values
-- of a certain cost (typically the size of the value).
data Enumerate a = Enumerate
   { Enumerate a -> RevList (Finite a)
revParts   ::  RevList (Finite a)
   } deriving Typeable

parts :: Enumerate a -> [Finite a]
parts :: Enumerate a -> [Finite a]
parts = RevList (Finite a) -> [Finite a]
forall a. RevList a -> [a]
fromRev (RevList (Finite a) -> [Finite a])
-> (Enumerate a -> RevList (Finite a)) -> Enumerate a -> [Finite a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Enumerate a -> RevList (Finite a)
forall a. Enumerate a -> RevList (Finite a)
revParts

fromParts :: [Finite a] -> Enumerate a
fromParts :: [Finite a] -> Enumerate a
fromParts [Finite a]
ps = RevList (Finite a) -> Enumerate a
forall a. RevList (Finite a) -> Enumerate a
Enumerate ([Finite a] -> RevList (Finite a)
forall a. [a] -> RevList a
toRev [Finite a]
ps)

-- | Only use fmap with bijective functions (e.g. data constructors)
instance Functor Enumerate where
  fmap :: (a -> b) -> Enumerate a -> Enumerate b
fmap a -> b
f Enumerate a
e = RevList (Finite b) -> Enumerate b
forall a. RevList (Finite a) -> Enumerate a
Enumerate ((Finite a -> Finite b) -> RevList (Finite a) -> RevList (Finite b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> Finite a -> Finite b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) (RevList (Finite a) -> RevList (Finite b))
-> RevList (Finite a) -> RevList (Finite b)
forall a b. (a -> b) -> a -> b
$ Enumerate a -> RevList (Finite a)
forall a. Enumerate a -> RevList (Finite a)
revParts Enumerate a
e)

-- | Pure is 'singleton' and '<*>' corresponds to cartesian product (as with lists)
instance Applicative Enumerate where
  pure :: a -> Enumerate a
pure     = a -> Enumerate a
forall a. a -> Enumerate a
singleton
  Enumerate (a -> b)
f <*> :: Enumerate (a -> b) -> Enumerate a -> Enumerate b
<*> Enumerate a
a  = ((a -> b, a) -> b) -> Enumerate (a -> b, a) -> Enumerate b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((a -> b) -> a -> b) -> (a -> b, a) -> b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($)) (Enumerate (a -> b) -> Enumerate a -> Enumerate (a -> b, a)
forall a b. Enumerate a -> Enumerate b -> Enumerate (a, b)
cartesian Enumerate (a -> b)
f Enumerate a
a)

instance Alternative Enumerate where
  empty :: Enumerate a
empty = RevList (Finite a) -> Enumerate a
forall a. RevList (Finite a) -> Enumerate a
Enumerate RevList (Finite a)
forall a. Monoid a => a
mempty
  <|> :: Enumerate a -> Enumerate a -> Enumerate a
(<|>) = Enumerate a -> Enumerate a -> Enumerate a
forall a. Enumerate a -> Enumerate a -> Enumerate a
union

instance Sized Enumerate where
  pay :: Enumerate a -> Enumerate a
pay Enumerate a
e    = RevList (Finite a) -> Enumerate a
forall a. RevList (Finite a) -> Enumerate a
Enumerate (Finite a -> RevList (Finite a) -> RevList (Finite a)
forall a. a -> RevList a -> RevList a
revCons Finite a
forall a. Monoid a => a
mempty (RevList (Finite a) -> RevList (Finite a))
-> RevList (Finite a) -> RevList (Finite a)
forall a b. (a -> b) -> a -> b
$ Enumerate a -> RevList (Finite a)
forall a. Enumerate a -> RevList (Finite a)
revParts Enumerate a
e)
  aconcat :: [Enumerate a] -> Enumerate a
aconcat  = [Enumerate a] -> Enumerate a
forall a. Monoid a => [a] -> a
mconcat
  pair :: Enumerate a -> Enumerate b -> Enumerate (a, b)
pair     = Enumerate a -> Enumerate b -> Enumerate (a, b)
forall a b. Enumerate a -> Enumerate b -> Enumerate (a, b)
cartesian
  fin :: Integer -> Enumerate Integer
fin Integer
k    = [Finite Integer] -> Enumerate Integer
forall a. [Finite a] -> Enumerate a
fromParts [Integer -> Finite Integer
finFin Integer
k]

instance Semigroup (Enumerate a) where
  <> :: Enumerate a -> Enumerate a -> Enumerate a
(<>)  = Enumerate a -> Enumerate a -> Enumerate a
forall a. Enumerate a -> Enumerate a -> Enumerate a
union

-- | The @'mappend'@ is (disjoint) @'union'@
instance Monoid (Enumerate a) where
  mempty :: Enumerate a
mempty      = Enumerate a
forall (f :: * -> *) a. Alternative f => f a
empty
  mconcat :: [Enumerate a] -> Enumerate a
mconcat     = [Enumerate a] -> Enumerate a
forall a. [Enumerate a] -> Enumerate a
econcat

-- | Optimal 'mconcat' on enumerations.
econcat :: [Enumerate a] -> Enumerate a
econcat :: [Enumerate a] -> Enumerate a
econcat []    = Enumerate a
forall a. Monoid a => a
mempty
econcat [Enumerate a
a]   = Enumerate a
a
econcat [Enumerate a
a,Enumerate a
b] = Enumerate a -> Enumerate a -> Enumerate a
forall a. Enumerate a -> Enumerate a -> Enumerate a
union Enumerate a
a Enumerate a
b
econcat [Enumerate a]
xs    = RevList (Finite a) -> Enumerate a
forall a. RevList (Finite a) -> Enumerate a
Enumerate
  ([Finite a] -> RevList (Finite a)
forall a. [a] -> RevList a
toRev ([Finite a] -> RevList (Finite a))
-> ([[Finite a]] -> [Finite a])
-> [[Finite a]]
-> RevList (Finite a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Finite a] -> Finite a) -> [[Finite a]] -> [Finite a]
forall a b. (a -> b) -> [a] -> [b]
map [Finite a] -> Finite a
forall a. Monoid a => [a] -> a
mconcat ([[Finite a]] -> [Finite a])
-> ([[Finite a]] -> [[Finite a]]) -> [[Finite a]] -> [Finite a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[Finite a]] -> [[Finite a]]
forall a. [[a]] -> [[a]]
transpose ([[Finite a]] -> RevList (Finite a))
-> [[Finite a]] -> RevList (Finite a)
forall a b. (a -> b) -> a -> b
$ (Enumerate a -> [Finite a]) -> [Enumerate a] -> [[Finite a]]
forall a b. (a -> b) -> [a] -> [b]
map Enumerate a -> [Finite a]
forall a. Enumerate a -> [Finite a]
parts [Enumerate a]
xs)


-- Product of two enumerations
cartesian :: Enumerate a -> Enumerate b -> Enumerate (a,b)
cartesian :: Enumerate a -> Enumerate b -> Enumerate (a, b)
cartesian (Enumerate RevList (Finite a)
xs1) (Enumerate RevList (Finite b)
xs2) = RevList (Finite (a, b)) -> Enumerate (a, b)
forall a. RevList (Finite a) -> Enumerate a
Enumerate (RevList (Finite a)
xs1 RevList (Finite a) -> RevList (Finite b) -> RevList (Finite (a, b))
forall a b.
RevList (Finite a) -> RevList (Finite b) -> RevList (Finite (a, b))
`prod` RevList (Finite b)
xs2)

prod :: RevList (Finite a) -> RevList (Finite b) -> RevList (Finite (a,b))
prod :: RevList (Finite a) -> RevList (Finite b) -> RevList (Finite (a, b))
prod (RevList [] [[Finite a]]
_)           RevList (Finite b)
_                 = RevList (Finite (a, b))
forall a. Monoid a => a
mempty
prod (RevList xs0 :: [Finite a]
xs0@(Finite a
_:[Finite a]
xst) [[Finite a]]
_)  (RevList [Finite b]
_ [[Finite b]]
rys0)  = [Finite (a, b)] -> RevList (Finite (a, b))
forall a. [a] -> RevList a
toRev([Finite (a, b)] -> RevList (Finite (a, b)))
-> [Finite (a, b)] -> RevList (Finite (a, b))
forall a b. (a -> b) -> a -> b
$ [[Finite b]] -> [Finite (a, b)]
forall b. [[Finite b]] -> [Finite (a, b)]
prod' [[Finite b]]
rys0 where

  -- We need to thread carefully here, making sure that guarded recursion is safe
  prod' :: [[Finite b]] -> [Finite (a, b)]
prod' []        = []
  prod' ([Finite b]
h:[[Finite b]]
t)  = [Finite b] -> [[Finite b]] -> [Finite (a, b)]
forall b. [Finite b] -> [[Finite b]] -> [Finite (a, b)]
go [Finite b]
h [[Finite b]]
t where
    go :: [Finite b] -> [[Finite b]] -> [Finite (a, b)]
go [Finite b]
ry [[Finite b]]
rys = [Finite a] -> [Finite b] -> Finite (a, b)
forall a b. [Finite a] -> [Finite b] -> Finite (a, b)
conv [Finite a]
xs0 [Finite b]
ry Finite (a, b) -> [Finite (a, b)] -> [Finite (a, b)]
forall a. a -> [a] -> [a]
: case [[Finite b]]
rys of
      ([Finite b]
ry':[[Finite b]]
rys')   -> [Finite b] -> [[Finite b]] -> [Finite (a, b)]
go [Finite b]
ry' [[Finite b]]
rys'
      []           -> [Finite b] -> [Finite a] -> [Finite (a, b)]
forall b a. [Finite b] -> [Finite a] -> [Finite (a, b)]
prod'' [Finite b]
ry [Finite a]
xst

  -- rys0 is exhausted, slide a window over xs0 until it is exhausted
  prod'' :: [Finite b] -> [Finite a] -> [Finite (a,b)]
  prod'' :: [Finite b] -> [Finite a] -> [Finite (a, b)]
prod'' [Finite b]
ry = [Finite a] -> [Finite (a, b)]
forall a. [Finite a] -> [Finite (a, b)]
go where
    go :: [Finite a] -> [Finite (a, b)]
go []         = []
    go xs :: [Finite a]
xs@(Finite a
_:[Finite a]
xs') = [Finite a] -> [Finite b] -> Finite (a, b)
forall a b. [Finite a] -> [Finite b] -> Finite (a, b)
conv [Finite a]
xs [Finite b]
ry Finite (a, b) -> [Finite (a, b)] -> [Finite (a, b)]
forall a. a -> [a] -> [a]
: [Finite a] -> [Finite (a, b)]
go [Finite a]
xs'

  conv :: [Finite a] -> [Finite b] -> Finite (a,b)
  conv :: [Finite a] -> [Finite b] -> Finite (a, b)
conv [Finite a]
xs [Finite b]
ys = Integer -> (Integer -> (a, b)) -> Finite (a, b)
forall a. Integer -> (Integer -> a) -> Finite a
Finite
    ([Integer] -> Integer
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ (Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(*) ((Finite a -> Integer) -> [Finite a] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Finite a -> Integer
forall a. Finite a -> Integer
fCard [Finite a]
xs) ((Finite b -> Integer) -> [Finite b] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Finite b -> Integer
forall a. Finite a -> Integer
fCard [Finite b]
ys ))
    ([Finite a] -> [Finite b] -> Integer -> (a, b)
forall a b. [Finite a] -> [Finite b] -> Integer -> (a, b)
prodSel [Finite a]
xs [Finite b]
ys)

  prodSel :: [Finite a] -> [Finite b] -> (Index -> (a,b))
  prodSel :: [Finite a] -> [Finite b] -> Integer -> (a, b)
prodSel (Finite a
f1:[Finite a]
f1s) (Finite b
f2:[Finite b]
f2s) = \Integer
i ->
    let mul :: Integer
mul = Finite a -> Integer
forall a. Finite a -> Integer
fCard Finite a
f1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Finite b -> Integer
forall a. Finite a -> Integer
fCard Finite b
f2
    in  if Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
mul
        then  let (Integer
q, Integer
r) = (Integer
i Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Finite b -> Integer
forall a. Finite a -> Integer
fCard Finite b
f2)
              in (Finite a -> Integer -> a
forall a. Finite a -> Integer -> a
fIndex Finite a
f1 Integer
q, Finite b -> Integer -> b
forall a. Finite a -> Integer -> a
fIndex Finite b
f2 Integer
r)
        else [Finite a] -> [Finite b] -> Integer -> (a, b)
forall a b. [Finite a] -> [Finite b] -> Integer -> (a, b)
prodSel [Finite a]
f1s [Finite b]
f2s (Integer
iInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
mul)
  prodSel [Finite a]
_ [Finite b]
_ = \Integer
_ -> [Char] -> (a, b)
forall a. HasCallStack => [Char] -> a
error [Char]
"index out of bounds"


union :: Enumerate a -> Enumerate a -> Enumerate a
union :: Enumerate a -> Enumerate a -> Enumerate a
union (Enumerate RevList (Finite a)
xs1) (Enumerate RevList (Finite a)
xs2) = RevList (Finite a) -> Enumerate a
forall a. RevList (Finite a) -> Enumerate a
Enumerate (RevList (Finite a)
xs1 RevList (Finite a) -> RevList (Finite a) -> RevList (Finite a)
forall a. Monoid a => a -> a -> a
`mappend` RevList (Finite a)
xs2)


-- | The definition of @pure@ for the applicative instance.
singleton :: a -> Enumerate a
singleton :: a -> Enumerate a
singleton a
a = RevList (Finite a) -> Enumerate a
forall a. RevList (Finite a) -> Enumerate a
Enumerate (Finite a -> RevList (Finite a)
forall a. a -> RevList a
revPure (Finite a -> RevList (Finite a)) -> Finite a -> RevList (Finite a)
forall a b. (a -> b) -> a -> b
$ a -> Finite a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a)



------------------------------------------------------------------
-- Reverse lists

-- | A data structure that contains a list and the reversals of all initial
-- segments of the list. Intuitively
--
-- @reversals xs !! n = reverse (take (n+1) (fromRev xs))@
--
-- Any operation on a @RevList@ typically discards the reversals and constructs
-- new reversals on demand.
data RevList a = RevList {RevList a -> [a]
fromRev :: [a], RevList a -> [[a]]
reversals :: [[a]]} deriving Int -> RevList a -> ShowS
[RevList a] -> ShowS
RevList a -> [Char]
(Int -> RevList a -> ShowS)
-> (RevList a -> [Char])
-> ([RevList a] -> ShowS)
-> Show (RevList a)
forall a. Show a => Int -> RevList a -> ShowS
forall a. Show a => [RevList a] -> ShowS
forall a. Show a => RevList a -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
showList :: [RevList a] -> ShowS
$cshowList :: forall a. Show a => [RevList a] -> ShowS
show :: RevList a -> [Char]
$cshow :: forall a. Show a => RevList a -> [Char]
showsPrec :: Int -> RevList a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> RevList a -> ShowS
Show

instance Functor RevList where
  fmap :: (a -> b) -> RevList a -> RevList b
fmap a -> b
f = [b] -> RevList b
forall a. [a] -> RevList a
toRev ([b] -> RevList b) -> (RevList a -> [b]) -> RevList a -> RevList b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f ([a] -> [b]) -> (RevList a -> [a]) -> RevList a -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RevList a -> [a]
forall a. RevList a -> [a]
fromRev

instance Semigroup a => Semigroup (RevList a) where
  <> :: RevList a -> RevList a -> RevList a
(<>) RevList a
as RevList a
bs  = [a] -> RevList a
forall a. [a] -> RevList a
toRev ([a] -> RevList a) -> [a] -> RevList a
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> [a]
forall a. Semigroup a => [a] -> [a] -> [a]
zipMon (RevList a -> [a]
forall a. RevList a -> [a]
fromRev RevList a
as) (RevList a -> [a]
forall a. RevList a -> [a]
fromRev RevList a
bs) where
    zipMon :: Semigroup a => [a] -> [a] -> [a]
    zipMon :: [a] -> [a] -> [a]
zipMon (a
x:[a]
xs) (a
y:[a]
ys) = a
x a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. Semigroup a => [a] -> [a] -> [a]
zipMon [a]
xs [a]
ys
    zipMon [a]
xs [a]
ys         = [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
ys

-- Maybe this should be append instead?
-- | Padded zip
instance Semigroup a => Monoid (RevList a) where
  mempty :: RevList a
mempty   = [a] -> RevList a
forall a. [a] -> RevList a
toRev[]
  mappend :: RevList a -> RevList a -> RevList a
mappend  = RevList a -> RevList a -> RevList a
forall a. Semigroup a => a -> a -> a
(<>)

-- | Constructs a "Reverse list" variant of a given list. In a sensible
-- Haskell implementation evaluating any inital segment of
-- @'reversals' (toRev xs)@ uses linear memory in the size of the segment.
toRev:: [a] -> RevList a
toRev :: [a] -> RevList a
toRev [a]
as = [a] -> [[a]] -> RevList a
forall a. [a] -> [[a]] -> RevList a
RevList [a]
as ([[a]] -> RevList a) -> [[a]] -> RevList a
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> [[a]]
forall a. [a] -> [a] -> [[a]]
go [] [a]
as where
  go :: [a] -> [a] -> [[a]]
go [a]
_ []       = []
  go [a]
rev (a
x:[a]
xs) = let rev' :: [a]
rev' = a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
rev in [a]
rev' [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [[a]]
go [a]
rev' [a]
xs

-- | Adds an  element to the head of a @RevList@. Constant memory iff the
-- the reversals of the resulting list are not evaluated (which is frequently
-- the case in @Feat@).
revCons :: a -> RevList a -> RevList a
revCons :: a -> RevList a -> RevList a
revCons a
a = [a] -> RevList a
forall a. [a] -> RevList a
toRev([a] -> RevList a) -> (RevList a -> [a]) -> RevList a -> RevList a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> (RevList a -> [a]) -> RevList a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RevList a -> [a]
forall a. RevList a -> [a]
fromRev

revPure :: a -> RevList a
revPure :: a -> RevList a
revPure a
a = [a] -> [[a]] -> RevList a
forall a. [a] -> [[a]] -> RevList a
RevList [a
a] [[a
a]]