{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Constrained.SumList where

import Data.Either (partitionEithers)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE
import Data.Semigroup (sconcat)
import System.Random (Random (..))
import Test.QuickCheck (Gen, choose, shuffle, vectorOf)

-- =======================================================
-- Helper functions for genSizedList

data Solution t = Yes (NonEmpty [t]) | No [String]
  deriving (Solution t -> Solution t -> Bool
forall t. Eq t => Solution t -> Solution t -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Solution t -> Solution t -> Bool
$c/= :: forall t. Eq t => Solution t -> Solution t -> Bool
== :: Solution t -> Solution t -> Bool
$c== :: forall t. Eq t => Solution t -> Solution t -> Bool
Eq)

instance Show t => Show (Solution t) where
  show :: Solution t -> String
show (No [String]
xs) = String
"No" forall a. [a] -> [a] -> [a]
++ String
"\n" forall a. [a] -> [a] -> [a]
++ [String] -> String
unlines [String]
xs
  show (Yes NonEmpty [t]
xs) = String
"Yes " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show NonEmpty [t]
xs

-- | The basic idea is to concat all the Yes's and skip over the No's.
--   The one wrinkle is if everything is No, then in that case return an arbitrary one of the No's.
--   This can be done in linear time in the length of the list. Call that length n.
--   Check for all No. This takes time proportional to n. If it is true return one of them.
--   If it is not true. Concat all the Yes, and skip all the No.
--   We find the first No (if it exist), and all the Yes by partitioning the list
--   This is similar in spirit to Constrained.GenT.catGEs, but doesn't require a
--   a Monad to escape on the first No.
concatSolution :: Show t => t -> String -> t -> Int -> [Solution t] -> Solution t
concatSolution :: forall t.
Show t =>
t -> String -> t -> Int -> [Solution t] -> Solution t
concatSolution t
smallest String
pName t
total Int
count [Solution t]
sols =
  case forall a b. [Either a b] -> ([a], [b])
partitionEithers (forall a b. (a -> b) -> [a] -> [b]
map (\case Yes NonEmpty [t]
x -> forall a b. a -> Either a b
Left NonEmpty [t]
x; No [String]
x -> forall a b. b -> Either a b
Right [String]
x) [Solution t]
sols) of
    ([], [String]
n : [[String]]
_) -> forall t. [String] -> Solution t
No [String]
n -- All No, arbitrarily return the first.
    (NonEmpty [t]
y : [NonEmpty [t]]
ys, [[String]]
_) -> forall t. NonEmpty [t] -> Solution t
Yes forall a b. (a -> b) -> a -> b
$ forall a. Semigroup a => NonEmpty a -> a
sconcat (NonEmpty [t]
y forall a. a -> [a] -> NonEmpty a
:| [NonEmpty [t]]
ys) -- At least one Yes, and all No's skipped ('ys')
    ([], []) ->
      forall t. [String] -> Solution t
No -- The list is empty
        [ String
"The sample in pickAll was empty"
        , String
"smallest = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
smallest
        , String
"pred = " forall a. [a] -> [a] -> [a]
++ String
pName
        , String
"total = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
total
        , String
"count = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
count
        ]

newtype Cost = Cost Int deriving (Cost -> Cost -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Cost -> Cost -> Bool
$c/= :: Cost -> Cost -> Bool
== :: Cost -> Cost -> Bool
$c== :: Cost -> Cost -> Bool
Eq, Int -> Cost -> ShowS
[Cost] -> ShowS
Cost -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Cost] -> ShowS
$cshowList :: [Cost] -> ShowS
show :: Cost -> String
$cshow :: Cost -> String
showsPrec :: Int -> Cost -> ShowS
$cshowsPrec :: Int -> Cost -> ShowS
Show, Integer -> Cost
Cost -> Cost
Cost -> Cost -> Cost
forall a.
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> a)
-> (a -> a)
-> (Integer -> a)
-> Num a
fromInteger :: Integer -> Cost
$cfromInteger :: Integer -> Cost
signum :: Cost -> Cost
$csignum :: Cost -> Cost
abs :: Cost -> Cost
$cabs :: Cost -> Cost
negate :: Cost -> Cost
$cnegate :: Cost -> Cost
* :: Cost -> Cost -> Cost
$c* :: Cost -> Cost -> Cost
- :: Cost -> Cost -> Cost
$c- :: Cost -> Cost -> Cost
+ :: Cost -> Cost -> Cost
$c+ :: Cost -> Cost -> Cost
Num, Eq Cost
Cost -> Cost -> Bool
Cost -> Cost -> Ordering
Cost -> Cost -> Cost
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Cost -> Cost -> Cost
$cmin :: Cost -> Cost -> Cost
max :: Cost -> Cost -> Cost
$cmax :: Cost -> Cost -> Cost
>= :: Cost -> Cost -> Bool
$c>= :: Cost -> Cost -> Bool
> :: Cost -> Cost -> Bool
$c> :: Cost -> Cost -> Bool
<= :: Cost -> Cost -> Bool
$c<= :: Cost -> Cost -> Bool
< :: Cost -> Cost -> Bool
$c< :: Cost -> Cost -> Bool
compare :: Cost -> Cost -> Ordering
$ccompare :: Cost -> Cost -> Ordering
Ord)

firstYesG ::
  Monad m => Solution t -> (x -> Cost -> m (Cost, Solution t)) -> [x] -> Cost -> m (Cost, Solution t)
firstYesG :: forall (m :: * -> *) t x.
Monad m =>
Solution t
-> (x -> Cost -> m (Cost, Solution t))
-> [x]
-> Cost
-> m (Cost, Solution t)
firstYesG Solution t
nullSolution x -> Cost -> m (Cost, Solution t)
f [x]
xs Cost
c = [x] -> Cost -> m (Cost, Solution t)
go [x]
xs Cost
c
  where
    go :: [x] -> Cost -> m (Cost, Solution t)
go [] Cost
cost = forall (f :: * -> *) a. Applicative f => a -> f a
pure (Cost
cost, Solution t
nullSolution)
    go [x
x] Cost
cost = x -> Cost -> m (Cost, Solution t)
f x
x (Cost
cost forall a. Num a => a -> a -> a
+ Cost
1)
    go (x
x : [x]
more) Cost
cost = do
      (Cost, Solution t)
ans <- x -> Cost -> m (Cost, Solution t)
f x
x (Cost
cost forall a. Num a => a -> a -> a
+ Cost
1)
      case (Cost, Solution t)
ans of
        (Cost
cost1, No [String]
_) -> [x] -> Cost -> m (Cost, Solution t)
go [x]
more Cost
cost1
        (Cost, Solution t)
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Cost, Solution t)
ans

noChoices :: Show t => Cost -> String -> t -> t -> t -> Int -> [(t, t)] -> Solution t
noChoices :: forall t.
Show t =>
Cost -> String -> t -> t -> t -> Int -> [(t, t)] -> Solution t
noChoices Cost
cost String
p t
smallest t
largest t
total Int
count [(t, t)]
samp =
  forall t. [String] -> Solution t
No
    [ String
"No legal choice can be found, where"
    , String
"  predicate = " forall a. [a] -> [a] -> [a]
++ String
p
    , String
"  smallest = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
smallest
    , String
"  largest = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
largest
    , String
"  total = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
total
    , String
"  count = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
count
    , String
"  cost = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Cost
cost
    , String
"Small sample of what was explored"
    , forall a. Show a => a -> String
show [(t, t)]
samp
    ]

-- =====================================================

-- | Given 'count', return a list if pairs, that add to 'count'
--   splitsOf 6 --> [(1,5),(2,4),(3,3)].
--   Note we don't return reflections like (5,1) and (4,2),
--   as they have the same information as (1,5) and (2,4).
splitsOf :: Integral b => b -> [(b, b)]
splitsOf :: forall b. Integral b => b -> [(b, b)]
splitsOf b
count = [(b
i, b
j) | b
i <- [b
1 .. forall a. Integral a => a -> a -> a
div b
count b
2], let j :: b
j = b
count forall a. Num a => a -> a -> a
- b
i]
{-# SPECIALIZE splitsOf :: Int -> [(Int, Int)] #-}

-- | Given a Path, find a representative solution, 'ans', for that path, such that
--   1) (length ans) == 'count',
--   2) (sum ans) == 'total'
--   3) (all p ans) is True
--   What is a path?
--   Suppose i==5, then we recursively explore every way to split 5 into
--   split pairs that add to 5. I.e. (1,4) (2,3), then we split each of those.
--   Here is a picture of the graph of all paths for i==5. A path goes from the root '5'
--   to one of the leaves. Note all leaves are count == '1 (where the solution is '[total]').
--   To solve for 5, we could solve either of the sub problems rooted at 5: [1,4] or [2,3].
--   In 'pickAll' we will try to solve both, but in pick1, we only attempt 1 of those sub problems.
--   5
--   |
--   [1,4]
--   |  |
--   |  [1,3]
--   |  |  |
--   |  |  [1,2]
--   |  |     |
--   |  |     [1,1]
--   |  |
--   |  [2,2]
--   |   | |
--   |   | [1,1]
--   |   |
--   |   [1,1]
--   |
--   [2,3]
--    | |
--    | [1,2]
--    |    |
--    |    [1,1]
--    [1,1]
--  In 'pickAll' will explore a path for every split of 'count'
--  so if it returns (No _), we can be somewhat confidant that no solution exists.
--  Note that count of 1 and 2, are base cases.
--  When 'count' is greater than 1, we need to sample from [smallest..total],
--  so 'smallest' better be less that or equal to 'total'
pickAll ::
  forall t.
  (Show t, Integral t, Random t) =>
  t -> t -> (String, t -> Bool) -> t -> Int -> Cost -> Gen (Cost, Solution t)
pickAll :: forall t.
(Show t, Integral t, Random t) =>
t
-> t
-> (String, t -> Bool)
-> t
-> Int
-> Cost
-> Gen (Cost, Solution t)
pickAll t
smallest t
largest (String
pName, t -> Bool
_) t
total Int
count Cost
cost
  | Cost
cost forall a. Ord a => a -> a -> Bool
> Cost
1000 =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
        ( Cost
cost
        , forall t. [String] -> Solution t
No
            [ String
"pickAll exceeds cost limit " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Cost
cost
            , String
"  predicate = " forall a. [a] -> [a] -> [a]
++ String
pName
            , String
"  smallest = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
smallest
            , String
"  largest = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
largest
            , String
"  total = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
total
            , String
"  count = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
count
            ]
        )
pickAll t
smallest t
largest (String
pName, t -> Bool
p) t
total Int
1 Cost
cost =
  if t -> Bool
p t
total
    then forall (f :: * -> *) a. Applicative f => a -> f a
pure (Cost
cost, forall t. NonEmpty [t] -> Solution t
Yes forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure [t
total])
    else forall (f :: * -> *) a. Applicative f => a -> f a
pure (Cost
cost, forall t.
Show t =>
Cost -> String -> t -> t -> t -> Int -> [(t, t)] -> Solution t
noChoices Cost
cost String
pName t
smallest t
largest t
total Int
1 [(t
total, t
0)])
pickAll t
smallest t
largest (String
pName, t -> Bool
_) t
total Int
count Cost
cost
  | t
smallest forall a. Ord a => a -> a -> Bool
> t
largest =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
        ( Cost
cost
        , forall t. [String] -> Solution t
No
            [ String
"The feasible range to pickAll [" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
smallest forall a. [a] -> [a] -> [a]
++ String
" .. " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall a. Integral a => a -> a -> a
div t
total t
2) forall a. [a] -> [a] -> [a]
++ String
"] was empty"
            , String
"  predicate = " forall a. [a] -> [a] -> [a]
++ String
pName
            , String
"  smallest = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
smallest
            , String
"  largest = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
largest
            , String
"  total = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
total
            , String
"  count = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
count
            , String
"  cost = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Cost
cost
            ]
        )
pickAll t
smallest t
largest (String
pName, t -> Bool
p) t
total Int
2 Cost
cost = do
  -- for small things, enumerate all possibilities
  -- for large things, use a fair sample.
  [(t, t)]
choices <- forall t.
(Random t, Integral t) =>
t -> t -> t -> t -> Int -> Gen [(t, t)]
smallSample t
smallest t
largest t
total t
1000 Int
100
  case forall a. (a -> Bool) -> [a] -> [a]
filter (\(t
x, t
y) -> t -> Bool
p t
x Bool -> Bool -> Bool
&& t -> Bool
p t
y) [(t, t)]
choices of
    [] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (Cost
cost forall a. Num a => a -> a -> a
+ Cost
1, forall t.
Show t =>
Cost -> String -> t -> t -> t -> Int -> [(t, t)] -> Solution t
noChoices Cost
cost String
pName t
smallest t
largest t
total Int
2 (forall a. Int -> [a] -> [a]
take Int
10 [(t, t)]
choices))
    [(t, t)]
zs -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (Cost
cost forall a. Num a => a -> a -> a
+ Cost
1, forall t. NonEmpty [t] -> Solution t
Yes forall a b. (a -> b) -> a -> b
$ forall a. [a] -> NonEmpty a
NE.fromList (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(t
x, t
y) -> [t
x, t
y]) [(t, t)]
zs))
pickAll t
smallest t
largest (String
pName, t -> Bool
p) t
total Int
count Cost
cost = do
  -- Compute a representative sample of the choices between smallest and total.
  -- E.g. when smallest = -2, and total = 5, the complete set of values is:
  -- [(-2,7),(-1,6),(0,5),(1,4),(2,3),(3,2),(4,1),(5,0)]  Note they all add to 5
  -- We could explore the whole set of values, but that can be millions of choices.
  -- so we choose to explore a representative subset. See the function 'fairSample', for details.
  -- Remember this is just 1 step on one path. So if this step fails, there are many more
  -- paths to explore. In fact there are usually many many solutions. We need to find just 1.
  [(t, t)]
choices <- forall t.
(Random t, Integral t) =>
t -> t -> t -> t -> Int -> Gen [(t, t)]
smallSample t
smallest t
largest t
total t
1000 Int
20
  -- The choice of splits is crucial. If total >> count, we want the larger splits first
  -- if count >> total , we want smaller splits first
  [(Int, Int)]
splits <-
    if Int
count forall a. Ord a => a -> a -> Bool
>= Int
20
      then forall a. [a] -> Gen [a]
shuffle forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take Int
10 (forall b. Integral b => b -> [(b, b)]
splitsOf Int
count)
      else
        if t
total forall a. Ord a => a -> a -> Bool
> forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
count
          then forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. [a] -> [a]
reverse (forall b. Integral b => b -> [(b, b)]
splitsOf Int
count))
          else forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall b. Integral b => b -> [(b, b)]
splitsOf Int
count)

  forall (m :: * -> *) t x.
Monad m =>
Solution t
-> (x -> Cost -> m (Cost, Solution t))
-> [x]
-> Cost
-> m (Cost, Solution t)
firstYesG
    (forall t. [String] -> Solution t
No [String
"No split has a solution", String
"cost = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Cost
cost])
    (forall t.
(Random t, Show t, Integral t) =>
t
-> t
-> (String, t -> Bool)
-> t
-> [(t, t)]
-> (Int, Int)
-> Cost
-> Gen (Cost, Solution t)
doSplit t
smallest t
largest (String
pName, t -> Bool
p) t
total [(t, t)]
choices)
    [(Int, Int)]
splits
    Cost
cost

-- TODO run some tests to see if this is a better solution than firstYesG
-- concatSolution smallest pName total count
--  <$> mapM  (doSplit smallest total (pName, p) choices (pickAll (depth +1) smallest)) splits

-- {-# SPECIALIZE pickAll::Int -> (String, Int -> Bool) -> Int -> Int -> Cost -> Gen (Cost, Solution Int) #-}

doSplit ::
  (Random t, Show t, Integral t) =>
  t ->
  t ->
  (String, t -> Bool) ->
  t ->
  [(t, t)] ->
  -- (t -> (String, t -> Bool) -> t -> Int -> Cost -> Gen (Cost, Solution t)) ->
  (Int, Int) ->
  Cost ->
  Gen (Cost, Solution t)
doSplit :: forall t.
(Random t, Show t, Integral t) =>
t
-> t
-> (String, t -> Bool)
-> t
-> [(t, t)]
-> (Int, Int)
-> Cost
-> Gen (Cost, Solution t)
doSplit t
smallest t
largest (String
pName, t -> Bool
p) t
total [(t, t)]
sample (Int
i, Int
j) Cost
c = [(t, t)] -> Cost -> Gen (Cost, Solution t)
go [(t, t)]
sample Cost
c
  where
    -- The 'sample' is a list of pairs (x,y), where we know (x+y) == total.
    -- We will search for the first good solution in the given sample
    -- to build a representative value for this path, with split (i,j).
    go :: [(t, t)] -> Cost -> Gen (Cost, Solution t)
go ((t
x, t
y) : [(t, t)]
more) Cost
cost0 = do
      -- Note (i+j) = current length of the ans we are looking for
      --      (x+y) = total
      -- pick 'ans1' such that (sum ans1 == x) and (length ans1 == i)
      (Cost
cost1, Solution t
ans1) <- forall t.
(Show t, Integral t, Random t) =>
t
-> t
-> (String, t -> Bool)
-> t
-> Int
-> Cost
-> Gen (Cost, Solution t)
pickAll t
smallest t
largest (String
pName, t -> Bool
p) t
x Int
i Cost
cost0
      -- pick 'ans2' such that (sum ans2 == y) and (length ans2 == j)
      (Cost
cost2, Solution t
ans2) <- forall t.
(Show t, Integral t, Random t) =>
t
-> t
-> (String, t -> Bool)
-> t
-> Int
-> Cost
-> Gen (Cost, Solution t)
pickAll t
smallest t
largest (String
pName, t -> Bool
p) t
y Int
j Cost
cost1
      case (Solution t
ans1, Solution t
ans2) of
        (Yes NonEmpty [t]
ys, Yes NonEmpty [t]
zs) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (Cost
cost2, forall t. NonEmpty [t] -> Solution t
Yes (forall a. [a] -> NonEmpty a
NE.fromList [[t]
a forall a. Semigroup a => a -> a -> a
<> [t]
b | [t]
a <- forall a. NonEmpty a -> [a]
NE.toList NonEmpty [t]
ys, [t]
b <- forall a. NonEmpty a -> [a]
NE.toList NonEmpty [t]
zs]))
        (Solution t, Solution t)
_ -> [(t, t)] -> Cost -> Gen (Cost, Solution t)
go [(t, t)]
more Cost
cost2
    go [] Cost
cost =
      case [(t, t)]
sample of
        [] ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
            ( Cost
cost
            , forall t. [String] -> Solution t
No
                [ String
"The sample passed to doSplit [" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
smallest forall a. [a] -> [a] -> [a]
++ String
" .. " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall a. Integral a => a -> a -> a
div t
total t
2) forall a. [a] -> [a] -> [a]
++ String
"] was empty"
                , String
"  predicate = " forall a. [a] -> [a] -> [a]
++ String
pName
                , String
"  smallest = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
smallest
                , String
"  total " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
total
                , String
"  count = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (Int
i forall a. Num a => a -> a -> a
+ Int
j)
                , String
"  split of count = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (Int
i, Int
j)
                ]
            )
        ((t
left, t
right) : [(t, t)]
_) ->
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
            ( Cost
cost
            , forall t. [String] -> Solution t
No
                [ String
"All choices in (genSizedList " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (Int
i forall a. Num a => a -> a -> a
+ Int
j) forall a. [a] -> [a] -> [a]
++ String
" 'p' " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
total forall a. [a] -> [a] -> [a]
++ String
") have failed."
                , String
"Here is 1 example failure."
                , String
"  smallest = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
smallest
                , String
"  total " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
total forall a. [a] -> [a] -> [a]
++ String
" = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
left forall a. [a] -> [a] -> [a]
++ String
" + " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
right
                , String
"  count = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (Int
i forall a. Num a => a -> a -> a
+ Int
j) forall a. [a] -> [a] -> [a]
++ String
", split of count = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (Int
i, Int
j)
                , String
"We are trying to solve sub-problems like:"
                , String
"  split " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
left forall a. [a] -> [a] -> [a]
++ String
" into " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
i forall a. [a] -> [a] -> [a]
++ String
" parts, where all parts meet 'p'"
                , String
"  split " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show t
right forall a. [a] -> [a] -> [a]
++ String
" into " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
j forall a. [a] -> [a] -> [a]
++ String
" parts, where all parts meet 'p'"
                , String
"predicate 'p' = " forall a. [a] -> [a] -> [a]
++ String
pName
                , String
"prefix of the sample"
                , [String] -> String
unlines (forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> String
show (forall a. Int -> [a] -> [a]
take Int
10 [(t, t)]
sample))
                ]
            )
{-# INLINE doSplit #-}

-- | If the sample is small enough, then enumerate all of it, otherwise take a fair sample.
smallSample :: (Random t, Integral t) => t -> t -> t -> t -> Int -> Gen [(t, t)]
smallSample :: forall t.
(Random t, Integral t) =>
t -> t -> t -> t -> Int -> Gen [(t, t)]
smallSample t
smallest t
largest t
total t
bound Int
size
  | t
largest forall a. Num a => a -> a -> a
- t
smallest forall a. Ord a => a -> a -> Bool
<= t
bound = do
      forall a. [a] -> Gen [a]
shuffle forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
takeWhile (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. Ord a => a -> a -> Bool
(<=)) [(t
x, t
total forall a. Num a => a -> a -> a
- t
x) | t
x <- [t
smallest .. t
total]]
  | Bool
otherwise = do
      [t]
choices <- forall a.
(Random a, Integral a) =>
a -> a -> Int -> Int -> Bool -> Gen [a]
fair t
smallest t
largest Int
size Int
5 Bool
True
      forall a. [a] -> Gen [a]
shuffle [(t
x, t
total forall a. Num a => a -> a -> a
- t
x) | t
x <- [t]
choices]
{-# INLINE smallSample #-}

-- | Generates a fair sample of numbers between 'smallest' and 'largest'.
--   makes sure there are numbers of all sizes. Controls both the size of the sample
--   and the precision (how many powers of 10 are covered)
--   Here is how we generate one sample when we call (fair (-3455) (10234) 12 3 True)
--   raw = [(-9999,-1000),(-999,-100),(-99,-10),(-9,-1),(0,9),(10,99),(100,999),(1000,9999),(10000,99999)]
--   ranges = [(-3455,-1000),(-999,-100),(-99,-10),(-9,-1),(0,9),(10,99),(100,999),(1000,9999),(10000,10234)]
--   count = 4
--   largePrecision = [(10000,10234),(1000,9999),(100,999)]
--   smallPrecision = [(-3455,-1000),(-999,-100),(-99,-10)]
--   answer generated = [10128,10104,10027,10048,4911,7821,5585,2157,448,630,802,889]
--   isLarge==True   means be biased towards the large end of the range,
--   isLArge==False  means be biased towards the small end of the range,
fair :: (Random a, Integral a) => a -> a -> Int -> Int -> Bool -> Gen [a]
fair :: forall a.
(Random a, Integral a) =>
a -> a -> Int -> Int -> Bool -> Gen [a]
fair a
smallest a
largest Int
size Int
precision Bool
isLarge =
  forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {a}. Random a => (a, a) -> Gen [a]
oneRange (if Bool
isLarge then [(a, a)]
largePrecision else [(a, a)]
smallPrecision)
  where
    raw :: [(a, a)]
raw = forall a b. (a -> b) -> [a] -> [b]
map forall a. Integral a => a -> (a, a)
logRange [forall t. Integral t => t -> t
logish a
smallest .. forall t. Integral t => t -> t
logish a
largest]
    fixEnds :: (a, a) -> (a, a)
fixEnds (a
x, a
y) = (forall a. Ord a => a -> a -> a
max a
smallest a
x, forall a. Ord a => a -> a -> a
min a
largest a
y)
    ranges :: [(a, a)]
ranges = forall a b. (a -> b) -> [a] -> [b]
map (a, a) -> (a, a)
fixEnds [(a, a)]
raw
    count :: Int
count = forall a. Integral a => a -> a -> a
div Int
size Int
precision
    largePrecision :: [(a, a)]
largePrecision = forall a. Int -> [a] -> [a]
take Int
precision (forall a. [a] -> [a]
reverse [(a, a)]
ranges)
    smallPrecision :: [(a, a)]
smallPrecision = forall a. Int -> [a] -> [a]
take Int
precision [(a, a)]
ranges
    oneRange :: (a, a) -> Gen [a]
oneRange (a
x, a
y) = forall a. Int -> Gen a -> Gen [a]
vectorOf Int
count (forall a. Random a => (a, a) -> Gen a
choose (a
x, a
y))

logRange :: Integral a => a -> (a, a)
logRange :: forall a. Integral a => a -> (a, a)
logRange a
1 = (a
10, a
99)
logRange (-1) = (-a
9, -a
1)
logRange a
n = case forall a. Ord a => a -> a -> Ordering
compare a
n a
0 of
  Ordering
EQ -> (a
0, a
9)
  Ordering
LT -> (forall a. Num a => a -> a
negate (forall a. Integral a => a -> a -> a
div a
b a
10), forall a. Num a => a -> a
negate (forall a. Integral a => a -> a -> a
div a
a a
10))
  Ordering
GT -> (a
10 forall a b. (Num a, Integral b) => a -> b -> a
^ a
n, a
10 forall a b. (Num a, Integral b) => a -> b -> a
^ (a
n forall a. Num a => a -> a -> a
+ a
1) forall a. Num a => a -> a -> a
- a
1)
  where
    (a
a, a
b) = forall a. Integral a => a -> (a, a)
logRange (forall a. Num a => a -> a
negate a
n)

-- | like (logBase10 n), except negative answers mean negative numbers, rather than fractions less than 1.
logish :: Integral t => t -> t
logish :: forall t. Integral t => t -> t
logish t
n
  | t
0 forall a. Ord a => a -> a -> Bool
<= t
n Bool -> Bool -> Bool
&& t
n forall a. Ord a => a -> a -> Bool
<= t
9 = t
0
  | t
n forall a. Ord a => a -> a -> Bool
> t
9 = t
1 forall a. Num a => a -> a -> a
+ forall t. Integral t => t -> t
logish (t
n forall a. Integral a => a -> a -> a
`div` t
10)
  | (-t
9) forall a. Ord a => a -> a -> Bool
<= t
n Bool -> Bool -> Bool
&& t
n forall a. Ord a => a -> a -> Bool
<= (-t
1) = -t
1
  | Bool
True = forall a. Num a => a -> a
negate (t
1 forall a. Num a => a -> a -> a
+ forall t. Integral t => t -> t
logish (forall a. Num a => a -> a
negate t
n))

-- =====================================================================