{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}

-- SPDX-License-Identifier: BSD-3-Clause

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Enumeration
-- Copyright   :  Brent Yorgey
-- Maintainer  :  byorgey@gmail.com
--
-- An /enumeration/ is a finite or countably infinite sequence of
-- values, that is, enumerations are isomorphic to lists.  However,
-- enumerations are represented a functions from index to value, so
-- they support efficient indexing and can be constructed for very
-- large finite sets.  A few examples are shown below.
--
-- >>> enumerate . takeE 15 $ listOf nat
-- [[],[0],[0,0],[1],[0,0,0],[1,0],[2],[0,1],[1,0,0],[2,0],[3],[0,0,0,0],[1,1],[2,0,0],[3,0]]
-- >>> select (listOf nat) 986235087203970702008108646
-- [11987363624969,1854392,1613,15,0,2,0]
--
-- @
-- data Tree = L | B Tree Tree deriving Show
--
-- treesUpTo :: Int -> Enumeration Tree
-- treesUpTo 0 = 'singleton' L
-- treesUpTo n = 'singleton' L '<|>' B '<$>' t' '<*>' t'
--   where t' = treesUpTo (n-1)
--
-- trees :: Enumeration Tree
-- trees = 'infinite' $ 'singleton' L '<|>' B '<$>' trees '<*>' trees
-- @
--
-- >>> card (treesUpTo 1)
-- Finite 2
-- >>> card (treesUpTo 10)
-- Finite 14378219780015246281818710879551167697596193767663736497089725524386087657390556152293078723153293423353330879856663164406809615688082297859526620035327291442156498380795040822304677
-- >>> select (treesUpTo 5) 12345
-- B (B L (B (B (B L L) L) (B L L))) (B (B (B L L) L) (B L L))
--
-- >>> card trees
-- Infinite
-- >>> select trees 12345
-- B (B (B (B L (B L L)) L) (B L (B (B L L) L))) (B (B L (B L L)) (B (B L L) (B L (B L L))))
--
-- For /invertible/ enumerations, /i.e./ bijections between some set
-- of values and natural numbers (or finite prefix thereof), see
-- "Data.Enumeration.Invertible".

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

module Data.Enumeration
  ( -- * Enumerations

    Enumeration
  , mkEnumeration

    -- ** Using enumerations

  , Cardinality(..), card
  , Index, select

  , isFinite
  , enumerate

    -- ** Primitive enumerations

  , unit
  , singleton
  , always
  , finite
  , finiteList
  , boundedEnum

  , nat
  , int
  , cw
  , rat

  -- ** Enumeration combinators

  , takeE
  , dropE
  , infinite
  , zipE, zipWithE
  , (<+>)
  , (><)
  , interleave

  , maybeOf
  , eitherOf
  , listOf
  , finiteSubsetOf
  , finiteEnumerationOf

    -- * Utilities

  , diagonal

  ) where

import           Control.Applicative

import           Data.Bits              ((.&.))
import           Data.Ratio
import           Data.Tuple             (swap)

import           GHC.Base               (Int (I#))
import           GHC.Integer.Logarithms (integerLog2#)

------------------------------------------------------------
-- Setup for doctest examples
------------------------------------------------------------

-- $setup
-- >>> :set -XTypeApplications
-- >>> :{
--   data Tree = L | B Tree Tree deriving Show
--   treesUpTo :: Int -> Enumeration Tree
--   treesUpTo 0 = singleton L
--   treesUpTo n = singleton L <|> B <$> t' <*> t'
--     where t' = treesUpTo (n-1)
--   trees :: Enumeration Tree
--   trees = infinite $ singleton L <|> B <$> trees <*> trees
-- :}

------------------------------------------------------------
-- Enumerations
------------------------------------------------------------

-- | The cardinality of a countable set: either a specific finite
--   natural number, or countably infinite.
data Cardinality = Finite !Integer | Infinite
  deriving (Int -> Cardinality -> ShowS
[Cardinality] -> ShowS
Cardinality -> String
(Int -> Cardinality -> ShowS)
-> (Cardinality -> String)
-> ([Cardinality] -> ShowS)
-> Show Cardinality
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Cardinality -> ShowS
showsPrec :: Int -> Cardinality -> ShowS
$cshow :: Cardinality -> String
show :: Cardinality -> String
$cshowList :: [Cardinality] -> ShowS
showList :: [Cardinality] -> ShowS
Show, Cardinality -> Cardinality -> Bool
(Cardinality -> Cardinality -> Bool)
-> (Cardinality -> Cardinality -> Bool) -> Eq Cardinality
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Cardinality -> Cardinality -> Bool
== :: Cardinality -> Cardinality -> Bool
$c/= :: Cardinality -> Cardinality -> Bool
/= :: Cardinality -> Cardinality -> Bool
Eq, Eq Cardinality
Eq Cardinality =>
(Cardinality -> Cardinality -> Ordering)
-> (Cardinality -> Cardinality -> Bool)
-> (Cardinality -> Cardinality -> Bool)
-> (Cardinality -> Cardinality -> Bool)
-> (Cardinality -> Cardinality -> Bool)
-> (Cardinality -> Cardinality -> Cardinality)
-> (Cardinality -> Cardinality -> Cardinality)
-> Ord Cardinality
Cardinality -> Cardinality -> Bool
Cardinality -> Cardinality -> Ordering
Cardinality -> Cardinality -> Cardinality
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
$ccompare :: Cardinality -> Cardinality -> Ordering
compare :: Cardinality -> Cardinality -> Ordering
$c< :: Cardinality -> Cardinality -> Bool
< :: Cardinality -> Cardinality -> Bool
$c<= :: Cardinality -> Cardinality -> Bool
<= :: Cardinality -> Cardinality -> Bool
$c> :: Cardinality -> Cardinality -> Bool
> :: Cardinality -> Cardinality -> Bool
$c>= :: Cardinality -> Cardinality -> Bool
>= :: Cardinality -> Cardinality -> Bool
$cmax :: Cardinality -> Cardinality -> Cardinality
max :: Cardinality -> Cardinality -> Cardinality
$cmin :: Cardinality -> Cardinality -> Cardinality
min :: Cardinality -> Cardinality -> Cardinality
Ord)

-- | @Cardinality@ has a @Num@ instance for convenience, so we can use
--   numeric literals as finite cardinalities, and add, subtract, and
--   multiply cardinalities.  Note that:
--
--   * subtraction is saturating (/i.e./ 3 - 5 = 0)
--
--   * infinity - infinity is treated as zero
--
--   * zero is treated as a "very strong" annihilator for multiplication:
--     even infinity * zero = zero.
instance Num Cardinality where
  fromInteger :: Integer -> Cardinality
fromInteger = Integer -> Cardinality
Finite

  Cardinality
Infinite + :: Cardinality -> Cardinality -> Cardinality
+ Cardinality
_        = Cardinality
Infinite
  Cardinality
_        + Cardinality
Infinite = Cardinality
Infinite
  Finite Integer
a + Finite Integer
b = Integer -> Cardinality
Finite (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
b)

  Finite Integer
0 * :: Cardinality -> Cardinality -> Cardinality
* Cardinality
_        = Integer -> Cardinality
Finite Integer
0
  Cardinality
_        * Finite Integer
0 = Integer -> Cardinality
Finite Integer
0
  Cardinality
Infinite * Cardinality
_        = Cardinality
Infinite
  Cardinality
_        * Cardinality
Infinite = Cardinality
Infinite
  Finite Integer
a * Finite Integer
b = Integer -> Cardinality
Finite (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b)

  Finite Integer
a - :: Cardinality -> Cardinality -> Cardinality
- Finite Integer
b = Integer -> Cardinality
Finite (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max Integer
0 (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
b))
  Cardinality
_        - Cardinality
Infinite = Integer -> Cardinality
Finite Integer
0
  Cardinality
_        - Cardinality
_        = Cardinality
Infinite

  negate :: Cardinality -> Cardinality
negate = String -> Cardinality -> Cardinality
forall a. HasCallStack => String -> a
error String
"Can't negate Cardinality"
  signum :: Cardinality -> Cardinality
signum = String -> Cardinality -> Cardinality
forall a. HasCallStack => String -> a
error String
"No signum for Cardinality"
  abs :: Cardinality -> Cardinality
abs    = String -> Cardinality -> Cardinality
forall a. HasCallStack => String -> a
error String
"No abs for Cardinality"

-- | An index into an enumeration.
type Index = Integer

-- | An enumeration of a finite or countably infinite set of
--   values. An enumeration is represented as a function from the natural numbers
--   (for infinite enumerations) or a finite prefix of the natural numbers (for finite ones)
--   to values.  Enumerations can thus easily be constructed for very large sets, and
--   support efficient indexing and random sampling.
--
--   'Enumeration' is an instance of the following type classes:
--
--   * 'Functor' (you can map a function over every element of an enumeration)
--   * 'Applicative' (representing Cartesian product of enumerations; see ('><'))
--   * 'Alternative' (representing disjoint union of enumerations; see ('<+>'))
--
--   'Enumeration' is /not/ a 'Monad', since there is no way to
--   implement 'Control.Monad.join' that works for any combination of
--   finite and infinite enumerations (but see 'interleave').
data Enumeration a = Enumeration
  { -- | Get the cardinality of an enumeration.
    forall a. Enumeration a -> Cardinality
card   :: Cardinality

    -- | Select the value at a particular index of an enumeration.
    --   Precondition: the index must be strictly less than the
    --   cardinality.  For infinite sets, every possible value must
    --   occur at some finite index.
  , forall a. Enumeration a -> Integer -> a
select :: Index -> a
  }
  deriving (forall a b. (a -> b) -> Enumeration a -> Enumeration b)
-> (forall a b. a -> Enumeration b -> Enumeration a)
-> Functor Enumeration
forall a b. a -> Enumeration b -> Enumeration a
forall a b. (a -> b) -> Enumeration a -> Enumeration b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Enumeration a -> Enumeration b
fmap :: forall a b. (a -> b) -> Enumeration a -> Enumeration b
$c<$ :: forall a b. a -> Enumeration b -> Enumeration a
<$ :: forall a b. a -> Enumeration b -> Enumeration a
Functor

-- | Create an enumeration primitively out of a cardinality and an
--   index function.
mkEnumeration :: Cardinality -> (Index -> a) -> Enumeration a
mkEnumeration :: forall a. Cardinality -> (Integer -> a) -> Enumeration a
mkEnumeration = Cardinality -> (Integer -> a) -> Enumeration a
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration

-- | The @Applicative@ instance for @Enumeration@ works similarly to
--   the instance for lists: @pure = singleton@, and @f '<*>' x@ takes
--   the Cartesian product of @f@ and @x@ (see ('><')) and applies
--   each paired function and argument.
instance Applicative Enumeration where
  pure :: forall a. a -> Enumeration a
pure    = a -> Enumeration a
forall a. a -> Enumeration a
singleton
  Enumeration (a -> b)
f <*> :: forall a b. Enumeration (a -> b) -> Enumeration a -> Enumeration b
<*> Enumeration a
x = ((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
($) ((a -> b, a) -> b) -> Enumeration (a -> b, a) -> Enumeration b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Enumeration (a -> b)
f Enumeration (a -> b) -> Enumeration a -> Enumeration (a -> b, a)
forall a b. Enumeration a -> Enumeration b -> Enumeration (a, b)
>< Enumeration a
x)

-- | The @Alternative@ instance for @Enumeration@ represents the sum
--   monoidal structure on enumerations: @empty@ is the empty
--   enumeration, and @('<|>') = ('<+>')@ is disjoint union.
instance Alternative Enumeration where
  empty :: forall a. Enumeration a
empty = Enumeration a
forall a. Enumeration a
void
  <|> :: forall a. Enumeration a -> Enumeration a -> Enumeration a
(<|>) = Enumeration a -> Enumeration a -> Enumeration a
forall a. Enumeration a -> Enumeration a -> Enumeration a
(<+>)

------------------------------------------------------------
-- Using enumerations
------------------------------------------------------------

-- | Test whether an enumeration is finite.
--
-- >>> isFinite (finiteList [1,2,3])
-- True
--
-- >>> isFinite nat
-- False
isFinite :: Enumeration a -> Bool
isFinite :: forall a. Enumeration a -> Bool
isFinite (Enumeration (Finite Integer
_) Integer -> a
_) = Bool
True
isFinite Enumeration a
_                          = Bool
False

-- | List the elements of an enumeration in order.  Inverse of
--   'finiteList'.
enumerate :: Enumeration a -> [a]
enumerate :: forall a. Enumeration a -> [a]
enumerate Enumeration a
e = (Integer -> a) -> [Integer] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
e) ([Integer] -> [a]) -> [Integer] -> [a]
forall a b. (a -> b) -> a -> b
$
  case Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e of
    Cardinality
Infinite -> [Integer
0 ..]
    Finite Integer
c -> [Integer
0 .. Integer
cInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1]

------------------------------------------------------------
-- Constructing Enumerations
------------------------------------------------------------

-- | The empty enumeration, with cardinality zero and no elements.
--
-- >>> card void
-- Finite 0
--
-- >>> enumerate void
-- []
void :: Enumeration a
void :: forall a. Enumeration a
void = Cardinality -> (Integer -> a) -> Enumeration a
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration Cardinality
0 (String -> Integer -> a
forall a. HasCallStack => String -> a
error String
"select void")

-- | The unit enumeration, with a single value of @()@.
--
-- >>> card unit
-- Finite 1
--
-- >>> enumerate unit
-- [()]
unit :: Enumeration ()
unit :: Enumeration ()
unit = Enumeration
  { card :: Cardinality
card   = Cardinality
1
  , select :: Integer -> ()
select = () -> Integer -> ()
forall a b. a -> b -> a
const ()
  }

-- | An enumeration of a single given element.
--
-- >>> card (singleton 17)
-- Finite 1
--
-- >>> enumerate (singleton 17)
-- [17]
singleton :: a -> Enumeration a
singleton :: forall a. a -> Enumeration a
singleton a
a = Cardinality -> (Integer -> a) -> Enumeration a
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration Cardinality
1 (a -> Integer -> a
forall a b. a -> b -> a
const a
a)

-- | A constant infinite enumeration.
--
-- >>> card (always 17)
-- Infinite
--
-- >>> enumerate . takeE 10 $ always 17
-- [17,17,17,17,17,17,17,17,17,17]
always :: a -> Enumeration a
always :: forall a. a -> Enumeration a
always a
a = Cardinality -> (Integer -> a) -> Enumeration a
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration Cardinality
Infinite (a -> Integer -> a
forall a b. a -> b -> a
const a
a)

-- | A finite prefix of the natural numbers.
--
-- >>> card (finite 5)
-- Finite 5
-- >>> card (finite 1234567890987654321)
-- Finite 1234567890987654321
--
-- >>> enumerate (finite 5)
-- [0,1,2,3,4]
-- >>> enumerate (finite 0)
-- []
finite :: Integer -> Enumeration Integer
finite :: Integer -> Enumeration Integer
finite Integer
n = Cardinality -> (Integer -> Integer) -> Enumeration Integer
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration (Integer -> Cardinality
Finite Integer
n) Integer -> Integer
forall a. a -> a
id

-- | Construct an enumeration from the elements of a /finite/ list.  To
--   turn an enumeration back into a list, use 'enumerate'.
--
-- >>> enumerate (finiteList [2,3,8,1])
-- [2,3,8,1]
-- >>> select (finiteList [2,3,8,1]) 2
-- 8
--
--   'finiteList' does not work on infinite lists: inspecting the
--   cardinality of the resulting enumeration (something many of the
--   enumeration combinators need to do) will hang trying to compute
--   the length of the infinite list.  To make an infinite enumeration,
--   use something like @f '<$>' 'nat'@ where @f@ is a function to
--   compute the value at any given index.
--
--   'finiteList' uses ('!!') internally, so you probably want to
--   avoid using it on long lists.  It would be possible to make a
--   version with better indexing performance by allocating a vector
--   internally, but I am too lazy to do it.  If you have a good use
--   case let me know (better yet, submit a pull request).
finiteList :: [a] -> Enumeration a
finiteList :: forall a. [a] -> Enumeration a
finiteList [a]
as = Cardinality -> (Integer -> a) -> Enumeration a
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration (Integer -> Cardinality
Finite (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
as)) (\Integer
k -> [a]
as [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
k)
  -- Note the use of !! is not very efficient, but for small lists it
  -- probably still beats the overhead of allocating a vector.  Most
  -- likely this will only ever be used with very small lists anyway.
  -- If it becomes a problem we could add another combinator that
  -- behaves just like finiteList but allocates a Vector internally.

-- | Enumerate all the values of a bounded 'Enum' instance.
--
-- >>> enumerate (boundedEnum @Bool)
-- [False,True]
--
-- >>> select (boundedEnum @Char) 97
-- 'a'
--
-- >>> card (boundedEnum @Int)
-- Finite 18446744073709551616
-- >>> select (boundedEnum @Int) 0
-- -9223372036854775808
boundedEnum :: forall a. (Enum a, Bounded a) => Enumeration a
boundedEnum :: forall a. (Enum a, Bounded a) => Enumeration a
boundedEnum = Enumeration
  { card :: Cardinality
card = Integer -> Cardinality
Finite (Integer
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
lo Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
  , select :: Integer -> a
select = Int -> a
forall a. Enum a => Int -> a
toEnum (Int -> a) -> (Integer -> Int) -> Integer -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> (Integer -> Integer) -> Integer -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
lo)
  }
  where
    lo, hi :: Index
    lo :: Integer
lo = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a -> Int
forall a. Enum a => a -> Int
fromEnum (forall a. Bounded a => a
minBound @a))
    hi :: Integer
hi = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a -> Int
forall a. Enum a => a -> Int
fromEnum (forall a. Bounded a => a
maxBound @a))

-- | The natural numbers, @0, 1, 2, ...@.
--
-- >>> enumerate . takeE 10 $ nat
-- [0,1,2,3,4,5,6,7,8,9]
nat :: Enumeration Integer
nat :: Enumeration Integer
nat = Cardinality -> (Integer -> Integer) -> Enumeration Integer
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration Cardinality
Infinite Integer -> Integer
forall a. a -> a
id

-- | All integers in the order @0, 1, -1, 2, -2, 3, -3, ...@.
int :: Enumeration Integer
int :: Enumeration Integer
int = Integer -> Integer
forall a. Num a => a -> a
negate (Integer -> Integer) -> Enumeration Integer -> Enumeration Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Enumeration Integer
nat Enumeration Integer -> Enumeration Integer -> Enumeration Integer
forall a. Enumeration a -> Enumeration a -> Enumeration a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Integer -> Enumeration Integer -> Enumeration Integer
forall a. Integer -> Enumeration a -> Enumeration a
dropE Integer
1 Enumeration Integer
nat

-- | The positive rational numbers, enumerated according to the
--   [Calkin-Wilf sequence](http://www.cs.ox.ac.uk/publications/publication1664-abstract.html).
--
-- >>> enumerate . takeE 10 $ cw
-- [1 % 1,1 % 2,2 % 1,1 % 3,3 % 2,2 % 3,3 % 1,1 % 4,4 % 3,3 % 5]
cw :: Enumeration Rational
cw :: Enumeration Rational
cw = Enumeration { card :: Cardinality
card = Cardinality
Infinite, select :: Integer -> Rational
select = (Integer -> Integer -> Rational) -> (Integer, Integer) -> Rational
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Integer -> Integer -> Rational
forall a. Integral a => a -> a -> Ratio a
(%) ((Integer, Integer) -> Rational)
-> (Integer -> (Integer, Integer)) -> Integer -> Rational
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> (Integer, Integer)
forall {t} {b}. (Num b, Integral t) => t -> (b, b)
go (Integer -> (Integer, Integer))
-> (Integer -> Integer) -> Integer -> (Integer, Integer)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer
forall a. Enum a => a -> a
succ }
  where
    go :: t -> (b, b)
go t
1 = (b
1,b
1)
    go t
n
      | t -> Bool
forall a. Integral a => a -> Bool
even t
n    = (b, b) -> (b, b)
forall {b}. Num b => (b, b) -> (b, b)
left (t -> (b, b)
go (t
n t -> t -> t
forall a. Integral a => a -> a -> a
`div` t
2))
      | Bool
otherwise = (b, b) -> (b, b)
forall {b}. Num b => (b, b) -> (b, b)
right (t -> (b, b)
go (t
n t -> t -> t
forall a. Integral a => a -> a -> a
`div` t
2))
    left :: (b, b) -> (b, b)
left  (!b
a, !b
b) = (b
a, b
ab -> b -> b
forall a. Num a => a -> a -> a
+b
b)
    right :: (b, b) -> (b, b)
right (!b
a, !b
b) = (b
ab -> b -> b
forall a. Num a => a -> a -> a
+b
b, b
b)

-- | An enumeration of all rational numbers: 0 first, then each
--   rational in the Calkin-Wilf sequence followed by its negative.
--
-- >>> enumerate . takeE 10 $ rat
-- [0 % 1,1 % 1,(-1) % 1,1 % 2,(-1) % 2,2 % 1,(-2) % 1,1 % 3,(-1) % 3,3 % 2]
rat :: Enumeration Rational
rat :: Enumeration Rational
rat = Rational -> Enumeration Rational
forall a. a -> Enumeration a
singleton Rational
0 Enumeration Rational
-> Enumeration Rational -> Enumeration Rational
forall a. Enumeration a -> Enumeration a -> Enumeration a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (Enumeration Rational
cw Enumeration Rational
-> Enumeration Rational -> Enumeration Rational
forall a. Enumeration a -> Enumeration a -> Enumeration a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Rational -> Rational
forall a. Num a => a -> a
negate (Rational -> Rational)
-> Enumeration Rational -> Enumeration Rational
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Enumeration Rational
cw)

-- | Take a finite prefix from the beginning of an enumeration.  @takeE
--   k e@ always yields the empty enumeration for \(k \leq 0\), and
--   results in @e@ whenever @k@ is greater than or equal to the
--   cardinality of the enumeration.  Otherwise @takeE k e@ has
--   cardinality @k@ and matches @e@ from @0@ to @k-1@.
--
-- >>> enumerate $ takeE 3 (boundedEnum @Int)
-- [-9223372036854775808,-9223372036854775807,-9223372036854775806]
--
-- >>> enumerate $ takeE 2 (finiteList [1..5])
-- [1,2]
--
-- >>> enumerate $ takeE 0 (finiteList [1..5])
-- []
--
-- >>> enumerate $ takeE 7 (finiteList [1..5])
-- [1,2,3,4,5]
takeE :: Integer -> Enumeration a -> Enumeration a
takeE :: forall a. Integer -> Enumeration a -> Enumeration a
takeE Integer
k Enumeration a
e
  | Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0             = Enumeration a
forall a. Enumeration a
void
  | Integer -> Cardinality
Finite Integer
k Cardinality -> Cardinality -> Bool
forall a. Ord a => a -> a -> Bool
>= Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e = Enumeration a
e
  | Bool
otherwise = Cardinality -> (Integer -> a) -> Enumeration a
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration (Integer -> Cardinality
Finite Integer
k) (Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
e)

-- | Drop some elements from the beginning of an enumeration.  @dropE k
--   e@ yields @e@ unchanged if \(k \leq 0\), and results in the empty
--   enumeration whenever @k@ is greater than or equal to the
--   cardinality of @e@.
--
-- >>> enumerate $ dropE 2 (finiteList [1..5])
-- [3,4,5]
--
-- >>> enumerate $ dropE 0 (finiteList [1..5])
-- [1,2,3,4,5]
--
-- >>> enumerate $ dropE 7 (finiteList [1..5])
-- []
dropE :: Integer -> Enumeration a -> Enumeration a
dropE :: forall a. Integer -> Enumeration a -> Enumeration a
dropE Integer
k Enumeration a
e
  | Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0             = Enumeration a
e
  | Integer -> Cardinality
Finite Integer
k Cardinality -> Cardinality -> Bool
forall a. Ord a => a -> a -> Bool
>= Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e = Enumeration a
forall a. Enumeration a
void
  | Bool
otherwise          = Enumeration
      { card :: Cardinality
card = Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e Cardinality -> Cardinality -> Cardinality
forall a. Num a => a -> a -> a
- Integer -> Cardinality
Finite Integer
k, select :: Integer -> a
select = Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
e (Integer -> a) -> (Integer -> Integer) -> Integer -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
k) }

-- | Explicitly mark an enumeration as having an infinite cardinality,
--   ignoring the previous cardinality. It is sometimes necessary to
--   use this as a "hint" when constructing a recursive enumeration
--   whose cardinality would otherwise consist of an infinite sum of
--   finite cardinalities.
--
--   For example, consider the following definitions:
--
-- @
-- data Tree = L | B Tree Tree deriving Show
--
-- treesBad :: Enumeration Tree
-- treesBad = singleton L '<|>' B '<$>' treesBad '<*>' treesBad
--
-- trees :: Enumeration Tree
-- trees = infinite $ singleton L '<|>' B '<$>' trees '<*>' trees
-- @
--
--   Trying to use @treesBad@ at all will simply hang, since trying to
--   compute its cardinality leads to infinite recursion.
--
-- @
-- \>>>\ select treesBad 5
-- ^CInterrupted.
-- @
--
--   However, using 'infinite', as in the definition of @trees@,
--   provides the needed laziness:
--
-- >>> card trees
-- Infinite
-- >>> enumerate . takeE 3 $ trees
-- [L,B L L,B L (B L L)]
-- >>> select trees 87239862967296
-- B (B (B (B (B L L) (B (B (B L L) L) L)) (B L (B L (B L L)))) (B (B (B L (B L (B L L))) (B (B L L) (B L L))) (B (B L (B L (B L L))) L))) (B (B L (B (B (B L (B L L)) (B L L)) L)) (B (B (B L (B L L)) L) L))
infinite :: Enumeration a -> Enumeration a
infinite :: forall a. Enumeration a -> Enumeration a
infinite (Enumeration Cardinality
_ Integer -> a
s) = Cardinality -> (Integer -> a) -> Enumeration a
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration Cardinality
Infinite Integer -> a
s

-- | Fairly interleave a set of /infinite/ enumerations.
--
--   For a finite set of infinite enumerations, a round-robin
--   interleaving is used. That is, if we think of an enumeration of
--   enumerations as a 2D matrix read off row-by-row, this corresponds
--   to taking the transpose of a matrix with finitely many infinite
--   rows, turning it into one with infinitely many finite rows.  For
--   an infinite set of infinite enumerations, /i.e./ an infinite 2D
--   matrix, the resulting enumeration reads off the matrix by
--   'diagonal's.
--
-- >>> enumerate . takeE 15 $ interleave (finiteList [nat, negate <$> nat, (*10) <$> nat])
-- [0,0,0,1,-1,10,2,-2,20,3,-3,30,4,-4,40]
--
-- >>> enumerate . takeE 15 $ interleave (always nat)
-- [0,0,1,0,1,2,0,1,2,3,0,1,2,3,4]
--
--   This function is similar to 'Control.Monad.join' in a
--   hypothetical 'Monad' instance for 'Enumeration', but it only
--   works when the inner enumerations are all infinite.
--
--   To interleave a finite enumeration of enumerations, some of which
--   may be finite, you can use @'Data.Foldable.asum' . 'enumerate'@.
--   If you want to interleave an infinite enumeration of finite
--   enumerations, you are out of luck.
interleave :: Enumeration (Enumeration a) -> Enumeration a
interleave :: forall a. Enumeration (Enumeration a) -> Enumeration a
interleave Enumeration (Enumeration a)
e = Enumeration
  { card :: Cardinality
card   = Cardinality
Infinite
  , select :: Integer -> a
select = \Integer
k ->
      let (Integer
i,Integer
j) = case Enumeration (Enumeration a) -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration (Enumeration a)
e of
            Finite Integer
n -> Integer
k Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`divMod` Integer
n
            Cardinality
Infinite -> Integer -> (Integer, Integer)
diagonal Integer
k
      in  Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select (Enumeration (Enumeration a) -> Integer -> Enumeration a
forall a. Enumeration a -> Integer -> a
select Enumeration (Enumeration a)
e Integer
j) Integer
i
  }

-- | Zip two enumerations in parallel, producing the pair of
--   elements at each index.  The resulting enumeration is truncated
--   to the cardinality of the smaller of the two arguments.
--
-- >>> enumerate $ zipE nat (boundedEnum @Bool)
-- [(0,False),(1,True)]
zipE :: Enumeration a -> Enumeration b -> Enumeration (a,b)
zipE :: forall a b. Enumeration a -> Enumeration b -> Enumeration (a, b)
zipE = (a -> b -> (a, b))
-> Enumeration a -> Enumeration b -> Enumeration (a, b)
forall a b c.
(a -> b -> c) -> Enumeration a -> Enumeration b -> Enumeration c
zipWithE (,)

-- | Zip two enumerations in parallel, applying the given function to
--   the pair of elements at each index to produce a new element.  The
--   resulting enumeration is truncated to the cardinality of the
--   smaller of the two arguments.
--
-- >>> enumerate $ zipWithE replicate (finiteList [1..10]) (dropE 35 (boundedEnum @Char))
-- ["#","$$","%%%","&&&&","'''''","((((((",")))))))","********","+++++++++",",,,,,,,,,,"]

zipWithE :: (a -> b -> c) -> Enumeration a -> Enumeration b -> Enumeration c
zipWithE :: forall a b c.
(a -> b -> c) -> Enumeration a -> Enumeration b -> Enumeration c
zipWithE a -> b -> c
f Enumeration a
e1 Enumeration b
e2 =
  Cardinality -> (Integer -> c) -> Enumeration c
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration (Cardinality -> Cardinality -> Cardinality
forall a. Ord a => a -> a -> a
min (Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e1) (Enumeration b -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration b
e2)) (\Integer
k -> a -> b -> c
f (Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
e1 Integer
k) (Enumeration b -> Integer -> b
forall a. Enumeration a -> Integer -> a
select Enumeration b
e2 Integer
k))

-- | Sum, /i.e./ disjoint union, of two enumerations.  If both are
--   finite, all the values of the first will be enumerated before the
--   values of the second.  If only one is finite, the values from the
--   finite enumeration will be listed first.  If both are infinite, a
--   fair (alternating) interleaving is used, so that every value ends
--   up at a finite index in the result.
--
--   Note that the ('<+>') operator is a synonym for ('<|>') from the
--   'Alternative' instance for 'Enumeration', which should be used in
--   preference to ('<+>').  ('<+>') is provided as a separate
--   standalone operator to make it easier to document.
--
-- >>> enumerate . takeE 10 $ singleton 17 <|> nat
-- [17,0,1,2,3,4,5,6,7,8]
--
-- >>> enumerate . takeE 10 $ nat <|> singleton 17
-- [17,0,1,2,3,4,5,6,7,8]
--
-- >>> enumerate . takeE 10 $ nat <|> (negate <$> nat)
-- [0,0,1,-1,2,-2,3,-3,4,-4]
--
--   Note that this is not associative in a strict sense.  In
--   particular, it may fail to be associative when mixing finite and
--   infinite enumerations:
--
-- >>> enumerate . takeE 10 $ nat <|> (singleton 17 <|> nat)
-- [0,17,1,0,2,1,3,2,4,3]
--
-- >>> enumerate . takeE 10 $ (nat <|> singleton 17) <|> nat
-- [17,0,0,1,1,2,2,3,3,4]
--
-- However, it is associative in several weaker senses:
--
--   * If all the enumerations are finite
--   * If all the enumerations are infinite
--   * If enumerations are considered equivalent up to reordering
--     (they are not, but considering them so may be acceptable in
--     some applications).
(<+>) :: Enumeration a -> Enumeration a -> Enumeration a
Enumeration a
e1 <+> :: forall a. Enumeration a -> Enumeration a -> Enumeration a
<+> Enumeration a
e2 = case (Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e1, Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e2) of

  -- optimize for void <+> e2.
  (Finite Integer
0, Cardinality
_)  -> Enumeration a
e2

  -- Note we don't want to add a case for e1 <+> void right away since
  -- that would require forcing the cardinality of e2, and we'd rather
  -- let the following case work lazily in the cardinality of e2.

  -- First enumeration is finite: just put it first
  (Finite Integer
k1, Cardinality
_) -> Enumeration
    { card :: Cardinality
card   = Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e1 Cardinality -> Cardinality -> Cardinality
forall a. Num a => a -> a -> a
+ Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e2
    , select :: Integer -> a
select = \Integer
k -> if Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
k1 then Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
e1 Integer
k else Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
e2 (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
k1)
    }

  -- First is infinite but second is finite: put all the second values first
  (Cardinality
_, Finite Integer
_) -> Enumeration a
e2 Enumeration a -> Enumeration a -> Enumeration a
forall a. Enumeration a -> Enumeration a -> Enumeration a
<+> Enumeration a
e1

  -- Both are infinite: use a fair (alternating) interleaving
  (Cardinality, Cardinality)
_ -> Enumeration (Enumeration a) -> Enumeration a
forall a. Enumeration (Enumeration a) -> Enumeration a
interleave (Cardinality
-> (Integer -> Enumeration a) -> Enumeration (Enumeration a)
forall a. Cardinality -> (Integer -> a) -> Enumeration a
Enumeration Cardinality
2 (\case {Integer
0 -> Enumeration a
e1; Integer
1 -> Enumeration a
e2}))

-- | One half of the isomorphism between \(\mathbb{N}\) and
--   \(\mathbb{N} \times \mathbb{N}\) which enumerates by diagonals:
--   turn a particular natural number index into its position in the
--   2D grid.  That is, given this numbering of a 2D grid:
--
--   @
--   0 1 3 6 ...
--   2 4 7
--   5 8
--   9
--   @
--
--   'diagonal' maps \(0 \mapsto (0,0), 1 \mapsto (0,1), 2 \mapsto (1,0) \dots\)
diagonal :: Integer -> (Integer, Integer)
diagonal :: Integer -> (Integer, Integer)
diagonal Integer
k = (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
t, Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
t))
  where
    d :: Integer
d = (Integer -> Integer
integerSqrt (Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
8Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
k) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2
    t :: Integer
t = Integer
dInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*(Integer
dInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2

-- | Cartesian product of enumerations. If both are finite, uses a
--   simple lexicographic ordering.  If only one is finite, the
--   resulting enumeration is still in lexicographic order, with the
--   infinite enumeration as the most significant component.  For two
--   infinite enumerations, uses a fair 'diagonal' interleaving.
--
-- >>> enumerate $ finiteList [1..3] >< finiteList "abcd"
-- [(1,'a'),(1,'b'),(1,'c'),(1,'d'),(2,'a'),(2,'b'),(2,'c'),(2,'d'),(3,'a'),(3,'b'),(3,'c'),(3,'d')]
--
-- >>> enumerate . takeE 10 $ finiteList "abc" >< nat
-- [('a',0),('b',0),('c',0),('a',1),('b',1),('c',1),('a',2),('b',2),('c',2),('a',3)]
--
-- >>> enumerate . takeE 10 $ nat >< finiteList "abc"
-- [(0,'a'),(0,'b'),(0,'c'),(1,'a'),(1,'b'),(1,'c'),(2,'a'),(2,'b'),(2,'c'),(3,'a')]
--
-- >>> enumerate . takeE 10 $ nat >< nat
-- [(0,0),(0,1),(1,0),(0,2),(1,1),(2,0),(0,3),(1,2),(2,1),(3,0)]
--
--   Like ('<+>'), this operation is also not associative (not even up
--   to reassociating tuples).
(><) :: Enumeration a -> Enumeration b -> Enumeration (a,b)
Enumeration a
e1 >< :: forall a b. Enumeration a -> Enumeration b -> Enumeration (a, b)
>< Enumeration b
e2 = case (Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e1, Enumeration b -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration b
e2) of

  -- The second enumeration is finite: use lexicographic ordering with
  -- the first as the most significant component
  (Cardinality
_, Finite Integer
k2) -> Enumeration
    { card :: Cardinality
card   = Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
e1 Cardinality -> Cardinality -> Cardinality
forall a. Num a => a -> a -> a
* Enumeration b -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration b
e2
    , select :: Integer -> (a, b)
select = \Integer
k -> let (Integer
i,Integer
j) = Integer
k Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`divMod` Integer
k2 in (Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
e1 Integer
i, Enumeration b -> Integer -> b
forall a. Enumeration a -> Integer -> a
select Enumeration b
e2 Integer
j)
    }

  -- The first is finite but the second is infinite: lexicographic
  -- with the second as most significant.
  (Finite Integer
_, Cardinality
_) -> (b, a) -> (a, b)
forall a b. (a, b) -> (b, a)
swap ((b, a) -> (a, b)) -> Enumeration (b, a) -> Enumeration (a, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Enumeration b
e2 Enumeration b -> Enumeration a -> Enumeration (b, a)
forall a b. Enumeration a -> Enumeration b -> Enumeration (a, b)
>< Enumeration a
e1)

  -- Both are infinite: enumerate by diagonals
  (Cardinality, Cardinality)
_ -> Enumeration
    { card :: Cardinality
card = Cardinality
Infinite
    , select :: Integer -> (a, b)
select = \Integer
k -> let (Integer
i,Integer
j) = Integer -> (Integer, Integer)
diagonal Integer
k in (Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
e1 Integer
i, Enumeration b -> Integer -> b
forall a. Enumeration a -> Integer -> a
select Enumeration b
e2 Integer
j)
    }

------------------------------------------------------------
-- Building standard data types
------------------------------------------------------------

-- | Enumerate all possible values of type `Maybe a`, where the values
--   of type `a` are taken from the given enumeration.
--
-- >>> enumerate $ maybeOf (finiteList [1,2,3])
-- [Nothing,Just 1,Just 2,Just 3]
maybeOf :: Enumeration a -> Enumeration (Maybe a)
maybeOf :: forall a. Enumeration a -> Enumeration (Maybe a)
maybeOf Enumeration a
a = Maybe a -> Enumeration (Maybe a)
forall a. a -> Enumeration a
singleton Maybe a
forall a. Maybe a
Nothing Enumeration (Maybe a)
-> Enumeration (Maybe a) -> Enumeration (Maybe a)
forall a. Enumeration a -> Enumeration a -> Enumeration a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> Enumeration a -> Enumeration (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Enumeration a
a

-- | Enumerae all possible values of type @Either a b@ with inner values
--   taken from the given enumerations.
--
-- >>> enumerate . takeE 6 $ eitherOf nat nat
-- [Left 0,Right 0,Left 1,Right 1,Left 2,Right 2]
eitherOf :: Enumeration a -> Enumeration b -> Enumeration (Either a b)
eitherOf :: forall a b.
Enumeration a -> Enumeration b -> Enumeration (Either a b)
eitherOf Enumeration a
a Enumeration b
b = a -> Either a b
forall a b. a -> Either a b
Left (a -> Either a b) -> Enumeration a -> Enumeration (Either a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Enumeration a
a Enumeration (Either a b)
-> Enumeration (Either a b) -> Enumeration (Either a b)
forall a. Enumeration a -> Enumeration a -> Enumeration a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> b -> Either a b
forall a b. b -> Either a b
Right (b -> Either a b) -> Enumeration b -> Enumeration (Either a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Enumeration b
b

-- | Enumerate all possible finite lists containing values from the given enumeration.
--
-- >>> enumerate . takeE 15 $ listOf nat
-- [[],[0],[0,0],[1],[0,0,0],[1,0],[2],[0,1],[1,0,0],[2,0],[3],[0,0,0,0],[1,1],[2,0,0],[3,0]]
-- >>> enumerate $ listOf empty :: [[Data.Void.Void]]
-- [[]]
listOf :: Enumeration a -> Enumeration [a]
listOf :: forall a. Enumeration a -> Enumeration [a]
listOf Enumeration a
a = case Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
a of
  Finite Integer
0 -> [a] -> Enumeration [a]
forall a. a -> Enumeration a
singleton []
  Cardinality
_        -> Enumeration [a]
listOfA
    where
      listOfA :: Enumeration [a]
listOfA = Enumeration [a] -> Enumeration [a]
forall a. Enumeration a -> Enumeration a
infinite (Enumeration [a] -> Enumeration [a])
-> Enumeration [a] -> Enumeration [a]
forall a b. (a -> b) -> a -> b
$ [a] -> Enumeration [a]
forall a. a -> Enumeration a
singleton [] Enumeration [a] -> Enumeration [a] -> Enumeration [a]
forall a. Enumeration a -> Enumeration a -> Enumeration a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (:) (a -> [a] -> [a]) -> Enumeration a -> Enumeration ([a] -> [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Enumeration a
a Enumeration ([a] -> [a]) -> Enumeration [a] -> Enumeration [a]
forall a b. Enumeration (a -> b) -> Enumeration a -> Enumeration b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Enumeration [a]
listOfA

-- | Enumerate all possible finite subsets of values from the given enumeration.
--
-- >>> enumerate $ finiteSubsetOf (finite 3)
-- [[],[0],[1],[0,1],[2],[0,2],[1,2],[0,1,2]]
finiteSubsetOf :: Enumeration a -> Enumeration [a]
finiteSubsetOf :: forall a. Enumeration a -> Enumeration [a]
finiteSubsetOf Enumeration a
as = Integer -> [a]
pick (Integer -> [a]) -> Enumeration Integer -> Enumeration [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Enumeration Integer
bitstrings
  where
    bitstrings :: Enumeration Integer
bitstrings = case Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
as of
      Cardinality
Infinite -> Enumeration Integer
nat
      Finite Integer
k -> Integer -> Enumeration Integer
finite (Integer
2Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
k)

    pick :: Integer -> [a]
pick Integer
0 = []
    pick Integer
n = Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
as (Integer -> Integer
integerLog2 Integer
l) a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Integer -> [a]
pick (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
l)
      where
        l :: Integer
l = Integer -> Integer
lsb Integer
n

    lsb :: Integer -> Integer
    lsb :: Integer -> Integer
lsb Integer
n = Integer
n Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. (-Integer
n)

    integerLog2 :: Integer -> Integer
    integerLog2 :: Integer -> Integer
integerLog2 Integer
n = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int# -> Int
I# (Integer -> Int#
integerLog2# Integer
n))

-- | @finiteEnumerationOf n a@ creates an enumeration of all sequences
--   of exactly n items taken from the enumeration @a@.
finiteEnumerationOf :: Int -> Enumeration a -> Enumeration (Enumeration a)
finiteEnumerationOf :: forall a. Int -> Enumeration a -> Enumeration (Enumeration a)
finiteEnumerationOf Int
0 Enumeration a
_ = Enumeration a -> Enumeration (Enumeration a)
forall a. a -> Enumeration a
singleton Enumeration a
forall a. Enumeration a
forall (f :: * -> *) a. Alternative f => f a
empty
finiteEnumerationOf Int
n Enumeration a
a = case Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
card Enumeration a
a of
  Finite Integer
k -> Integer -> Integer -> Enumeration a
selectEnum Integer
k (Integer -> Enumeration a)
-> Enumeration Integer -> Enumeration (Enumeration a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Integer -> Enumeration Integer
finite (Integer
kInteger -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Int
n)
  Cardinality
Infinite -> (Enumeration a
 -> Enumeration (Enumeration a) -> Enumeration (Enumeration a))
-> Enumeration (Enumeration a)
-> [Enumeration a]
-> Enumeration (Enumeration a)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Enumeration a
-> Enumeration (Enumeration a) -> Enumeration (Enumeration a)
forall a.
Enumeration a
-> Enumeration (Enumeration a) -> Enumeration (Enumeration a)
cons (Enumeration a -> Enumeration (Enumeration a)
forall a. a -> Enumeration a
singleton Enumeration a
forall a. Enumeration a
forall (f :: * -> *) a. Alternative f => f a
empty) (Int -> Enumeration a -> [Enumeration a]
forall a. Int -> a -> [a]
replicate Int
n Enumeration a
a)

  where
    selectEnum :: Integer -> Integer -> Enumeration a
selectEnum Integer
k = (Integer -> a) -> Enumeration Integer -> Enumeration a
forall a b. (a -> b) -> Enumeration a -> Enumeration b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
select Enumeration a
a) (Enumeration Integer -> Enumeration a)
-> (Integer -> Enumeration Integer) -> Integer -> Enumeration a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> Enumeration Integer
forall a. [a] -> Enumeration a
finiteList ([Integer] -> Enumeration Integer)
-> (Integer -> [Integer]) -> Integer -> Enumeration Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> [Integer]
forall a. [a] -> [a]
reverse ([Integer] -> [Integer])
-> (Integer -> [Integer]) -> Integer -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
n ([Integer] -> [Integer])
-> (Integer -> [Integer]) -> Integer -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer -> [Integer]
forall {t}. Integral t => t -> t -> [t]
toBase Integer
k

    toBase :: t -> t -> [t]
toBase t
_ t
0 = t -> [t]
forall a. a -> [a]
repeat t
0
    toBase t
k t
n = t
n t -> t -> t
forall a. Integral a => a -> a -> a
`mod` t
k t -> [t] -> [t]
forall a. a -> [a] -> [a]
: t -> t -> [t]
toBase t
k (t
n t -> t -> t
forall a. Integral a => a -> a -> a
`div` t
k)

    cons :: Enumeration a -> Enumeration (Enumeration a) -> Enumeration (Enumeration a)
    cons :: forall a.
Enumeration a
-> Enumeration (Enumeration a) -> Enumeration (Enumeration a)
cons Enumeration a
a Enumeration (Enumeration a)
as = Enumeration a -> Enumeration a -> Enumeration a
forall a. Enumeration a -> Enumeration a -> Enumeration a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(<|>) (Enumeration a -> Enumeration a -> Enumeration a)
-> Enumeration (Enumeration a)
-> Enumeration (Enumeration a -> Enumeration a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> Enumeration a
forall a. a -> Enumeration a
singleton (a -> Enumeration a)
-> Enumeration a -> Enumeration (Enumeration a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Enumeration a
a) Enumeration (Enumeration a -> Enumeration a)
-> Enumeration (Enumeration a) -> Enumeration (Enumeration a)
forall a b. Enumeration (a -> b) -> Enumeration a -> Enumeration b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Enumeration (Enumeration a)
as

-- https://mail.haskell.org/pipermail/haskell-cafe/2008-February/039465.html
-- imLog :: Integer->Integer->Integer
-- > >   imLog b x
-- > >     = if x < b then
-- > >         0
-- > >       else
-- > >         let
-- > >           l = 2 * imLog (b*b) x
-- > >           doDiv x l = if x < b then l else doDiv (x`div`b) (l+1)
-- > >         in
-- > >           doDiv (x`div`(b^l)) l

-- Note: more efficient integerSqrt in arithmoi
-- (Math.NumberTheory.Powers.Squares), but it's a rather heavyweight
-- dependency to pull in just for this.

-- Implementation of `integerSqrt` taken from the Haskell wiki:
-- https://wiki.haskell.org/Generic_number_type#squareRoot

-- | Find the square root (rounded down) of a positive integer.
--
-- >>> integerSqrt 0
-- 0
-- >>> integerSqrt 1
-- 1
-- >>> integerSqrt 3
-- 1
-- >>> integerSqrt 4
-- 2
-- >>> integerSqrt 38
-- 6
-- >>> integerSqrt 763686362402795580983595318628819602756
-- 27634875834763498734

integerSqrt :: Integer -> Integer
integerSqrt :: Integer -> Integer
integerSqrt Integer
0 = Integer
0
integerSqrt Integer
1 = Integer
1
integerSqrt Integer
n =
  let twopows :: [Integer]
twopows = (Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (Integer -> Int -> Integer
forall a. Num a => a -> Int -> a
^!Int
2) Integer
2
      (Integer
lowerRoot, Integer
lowerN) =
        [(Integer, Integer)] -> (Integer, Integer)
forall a. HasCallStack => [a] -> a
last ([(Integer, Integer)] -> (Integer, Integer))
-> [(Integer, Integer)] -> (Integer, Integer)
forall a b. (a -> b) -> a -> b
$ ((Integer, Integer) -> Bool)
-> [(Integer, Integer)] -> [(Integer, Integer)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Integer
nInteger -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>=) (Integer -> Bool)
-> ((Integer, Integer) -> Integer) -> (Integer, Integer) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Integer) -> Integer
forall a b. (a, b) -> b
snd) ([(Integer, Integer)] -> [(Integer, Integer)])
-> [(Integer, Integer)] -> [(Integer, Integer)]
forall a b. (a -> b) -> a -> b
$ [Integer] -> [Integer] -> [(Integer, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Integer
1Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
:[Integer]
twopows) [Integer]
twopows
      newtonStep :: Integer -> Integer
newtonStep Integer
x = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div Integer
n Integer
x) Integer
2
      isRoot :: Integer -> Bool
isRoot Integer
r = Integer
rInteger -> Int -> Integer
forall a. Num a => a -> Int -> a
^!Int
2 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
n Bool -> Bool -> Bool
&& Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< (Integer
rInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1)Integer -> Int -> Integer
forall a. Num a => a -> Int -> a
^!Int
2
      initGuess :: Integer
initGuess = Integer -> Integer
integerSqrt (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div Integer
n Integer
lowerN ) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
lowerRoot
  in  (Integer -> Bool) -> (Integer -> Integer) -> Integer -> Integer
forall a. (a -> Bool) -> (a -> a) -> a -> a
iterUntil Integer -> Bool
isRoot Integer -> Integer
newtonStep Integer
initGuess

iterUntil :: (a -> Bool) -> (a -> a) -> a -> a
iterUntil :: forall a. (a -> Bool) -> (a -> a) -> a -> a
iterUntil a -> Bool
p a -> a
f a
a
  | a -> Bool
p a
a = a
a
  | Bool
otherwise = (a -> Bool) -> (a -> a) -> a -> a
forall a. (a -> Bool) -> (a -> a) -> a -> a
iterUntil a -> Bool
p a -> a
f (a -> a
f a
a)

(^!) :: Num a => a -> Int -> a
^! :: forall a. Num a => a -> Int -> a
(^!) a
x Int
n = a
xa -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^Int
n