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

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

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Enumeration.Invertible
-- Copyright   :  Brent Yorgey
-- Maintainer  :  byorgey@gmail.com
--
-- An /invertible enumeration/ is a bijection between a set of values
-- and the natural numbers (or a finite prefix thereof), represented
-- as a pair of inverse functions, one in each direction.  Hence they
-- support efficient indexing and can be constructed for very large
-- finite sets.  A few examples are shown below.
--
-- Compared to "Data.Enumeration", one can also build invertible
-- enumerations of functions (or other type formers with contravariant
-- arguments); however, invertible enumerations no longer make for
-- valid 'Functor', 'Applicative', or 'Alternative' instances.
--
-- This module exports many of the same names as "Data.Enumeration";
-- the expectation is that you will choose one or the other to import,
-- though of course it is possible to import both if you qualify the
-- imports.
--
-----------------------------------------------------------------------------

module Data.Enumeration.Invertible
  ( -- * Invertible enumerations

    IEnumeration

    -- ** Using enumerations

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

  , isFinite
  , enumerate

    -- ** Primitive enumerations

  , void
  , unit
  , singleton
  , finite
  , finiteList
  , boundedEnum

  , nat
  , int
  , cw
  , rat

  -- ** Enumeration combinators

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

  , maybeOf
  , eitherOf
  , listOf
  , finiteSubsetOf
  , finiteEnumerationOf
  , functionOf

  -- * Utilities

  , undiagonal
  ) where

import           Control.Applicative (Alternative (..))
import           Data.Bits           (shiftL, (.|.))
import           Data.List           (findIndex, foldl')
import           Data.Maybe          (fromJust)
import           Data.Ratio

import           Data.Enumeration    (Cardinality (..), Enumeration, Index)
import qualified Data.Enumeration    as E

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

-- $setup
-- >>> :set -XTypeApplications
-- >>> import Control.Arrow ((&&&))
-- >>> import Data.Maybe (fromMaybe, listToMaybe)
-- >>> :{
--   data Tree = L | B Tree Tree deriving Show
--   treesUpTo :: Int -> IEnumeration Tree
--   treesUpTo 0 = singleton L
--   treesUpTo n = mapE toTree fromTree (unit <+> (t' >< t'))
--     where
--       t' = treesUpTo (n-1)
--   trees :: IEnumeration Tree
--   trees = infinite $ mapE toTree fromTree (unit <+> (trees >< trees))
--   toTree :: Either () (Tree, Tree) -> Tree
--   toTree = either (const L) (uncurry B)
--   fromTree :: Tree -> Either () (Tree, Tree)
--   fromTree L = Left ()
--   fromTree (B l r) = Right (l,r)
-- :}

------------------------------------------------------------
-- Invertible enumerations
------------------------------------------------------------

-- | An invertible enumeration is a bijection between a set of
--   enumerated values and the natural numbers, or a finite prefix of
--   the natural numbers.  An invertible enumeration is represented as
--   a function from natural numbers to values, paired with an inverse
--   function that returns the natural number index of a given value.
--   Enumerations can thus easily be constructed for very large sets,
--   and support efficient indexing and random sampling.
--
--   Note that 'IEnumeration' cannot be made an instance of 'Functor',
--   'Applicative', or 'Alternative'.  However, it does support the
--   'functionOf' combinator which cannot be supported by
--   "Data.Enumeration".

data IEnumeration a = IEnumeration
  { forall a. IEnumeration a -> Enumeration a
baseEnum :: Enumeration a
    -- | Compute the index of a particular value in its enumeration.
    --   Note that the result of 'locate' is only valid when given a
    --   value which is actually in the range of the enumeration.
  , forall a. IEnumeration a -> a -> Integer
locate   :: a -> Index
  }

-- | Map a pair of inverse functions over an invertible enumeration of
--   @a@ values to turn it into an invertible enumeration of @b@
--   values.  Because invertible enumerations contain a /bijection/ to
--   the natural numbers, we really do need both directions of a
--   bijection between @a@ and @b@ in order to map.  This is why
--   'IEnumeration' cannot be an instance of 'Functor'.
mapE :: (a -> b) -> (b -> a) -> IEnumeration a -> IEnumeration b
mapE :: forall a b.
(a -> b) -> (b -> a) -> IEnumeration a -> IEnumeration b
mapE a -> b
f b -> a
g (IEnumeration Enumeration a
e a -> Integer
l) = Enumeration b -> (b -> Integer) -> IEnumeration b
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration (a -> b
f (a -> b) -> Enumeration a -> Enumeration b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Enumeration a
e) (a -> Integer
l (a -> Integer) -> (b -> a) -> b -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> a
g)

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

-- | Select the value at a particular index.  Precondition: the index
--   must be strictly less than the cardinality.
select :: IEnumeration a -> (Index -> a)
select :: forall a. IEnumeration a -> Integer -> a
select = Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
E.select (Enumeration a -> Integer -> a)
-> (IEnumeration a -> Enumeration a)
-> IEnumeration a
-> Integer
-> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IEnumeration a -> Enumeration a
forall a. IEnumeration a -> Enumeration a
baseEnum

-- | Get the cardinality of an enumeration.
card :: IEnumeration a -> Cardinality
card :: forall a. IEnumeration a -> Cardinality
card = Enumeration a -> Cardinality
forall a. Enumeration a -> Cardinality
E.card (Enumeration a -> Cardinality)
-> (IEnumeration a -> Enumeration a)
-> IEnumeration a
-> Cardinality
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IEnumeration a -> Enumeration a
forall a. IEnumeration a -> Enumeration a
baseEnum

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

-- | List the elements of an enumeration in order.  Inverse of
--   'finiteList'.
enumerate :: IEnumeration a -> [a]
enumerate :: forall a. IEnumeration a -> [a]
enumerate IEnumeration a
e = case IEnumeration a -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration a
e of
  Cardinality
Infinite -> (Integer -> a) -> [Integer] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (IEnumeration a -> Integer -> a
forall a. IEnumeration a -> Integer -> a
select IEnumeration a
e) [Integer
0 ..]
  Finite Integer
c -> (Integer -> a) -> [Integer] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (IEnumeration a -> Integer -> a
forall a. IEnumeration a -> Integer -> a
select IEnumeration a
e) [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 :: IEnumeration a
void :: forall a. IEnumeration a
void = Enumeration a -> (a -> Integer) -> IEnumeration a
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration Enumeration a
forall a. Enumeration a
forall (f :: * -> *) a. Alternative f => f a
empty ([Char] -> a -> Integer
forall a. HasCallStack => [Char] -> a
error [Char]
"locate void")

-- | The unit enumeration, with a single value of @()@ at index 0.
--
-- >>> card unit
-- Finite 1
--
-- >>> enumerate unit
-- [()]
--
-- >>> locate unit ()
-- 0
unit :: IEnumeration ()
unit :: IEnumeration ()
unit = Enumeration () -> (() -> Integer) -> IEnumeration ()
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration Enumeration ()
E.unit (Integer -> () -> Integer
forall a b. a -> b -> a
const Integer
0)

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

-- | 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)
-- []
--
-- >>> locate (finite 5) 2
-- 2
finite :: Integer -> IEnumeration Integer
finite :: Integer -> IEnumeration Integer
finite Integer
n = Enumeration Integer -> (Integer -> Integer) -> IEnumeration Integer
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration (Integer -> Enumeration Integer
E.finite Integer
n) Integer -> Integer
forall a. a -> a
id

-- | Construct an enumeration from the elements of a /finite/ list.
--   The elements of the list must all be distinct. 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
-- >>> locate (finiteList [2,3,8,1]) 8
-- 2
--
--   '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.
--
--   'finiteList' uses ('!!') and 'findIndex' internally (which both
--   take $O(n)$ time), 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 :: Eq a => [a] -> IEnumeration a
finiteList :: forall a. Eq a => [a] -> IEnumeration a
finiteList [a]
as = Enumeration a -> (a -> Integer) -> IEnumeration a
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration ([a] -> Enumeration a
forall a. [a] -> Enumeration a
E.finiteList [a]
as) a -> Integer
forall {c}. Num c => a -> c
locateFinite
  -- Note the use of !! and findIndex 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.

  where
    locateFinite :: a -> c
locateFinite a
a = Int -> c
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> c) -> (Maybe Int -> Int) -> Maybe Int -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Int -> c) -> Maybe Int -> c
forall a b. (a -> b) -> a -> b
$ (a -> Bool) -> [a] -> Maybe Int
forall a. (a -> Bool) -> [a] -> Maybe Int
findIndex (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==a
a) [a]
as

-- | Enumerate all the values of a bounded 'Enum' instance.
--
-- >>> enumerate (boundedEnum @Bool)
-- [False,True]
--
-- >>> select (boundedEnum @Char) 97
-- 'a'
-- >>> locate (boundedEnum @Char) 'Z'
-- 90
--
-- >>> card (boundedEnum @Int)
-- Finite 18446744073709551616
-- >>> select (boundedEnum @Int) 0
-- -9223372036854775808
boundedEnum :: forall a. (Enum a, Bounded a) => IEnumeration a
boundedEnum :: forall a. (Enum a, Bounded a) => IEnumeration a
boundedEnum = Enumeration a -> (a -> Integer) -> IEnumeration a
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration Enumeration a
forall a. (Enum a, Bounded a) => Enumeration a
E.boundedEnum (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
subtract Integer
lo (Integer -> Integer) -> (a -> Integer) -> a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> (a -> Int) -> a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int
forall a. Enum a => a -> Int
fromEnum)
  where
    lo :: 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))

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

-- | All integers in the order @0, 1, -1, 2, -2, 3, -3, ...@.
int :: IEnumeration Integer
int :: IEnumeration Integer
int = Enumeration Integer -> (Integer -> Integer) -> IEnumeration Integer
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration Enumeration Integer
E.int Integer -> Integer
forall {a}. (Ord a, Num a) => a -> a
locateInt
  where
    locateInt :: a -> a
locateInt a
z
      | a
z a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0    = a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Num a => a -> a
abs a
z
      | Bool
otherwise = a
2a -> a -> a
forall a. Num a => a -> a -> a
*a
z a -> a -> a
forall a. Num a => a -> a -> a
- a
1

-- | 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]
-- >>> locate cw (3 % 2)
-- 4
-- >>> locate cw (23 % 99)
-- 3183
cw :: IEnumeration Rational
cw :: IEnumeration Rational
cw = Enumeration Rational
-> (Rational -> Integer) -> IEnumeration Rational
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration Enumeration Rational
E.cw (Integer -> Integer
forall a. Enum a => a -> a
pred (Integer -> Integer)
-> (Rational -> Integer) -> Rational -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rational -> Integer
forall {b} {a}. (Num b, Num a, Ord b) => Ratio b -> a
locateCW)
  where
    locateCW :: Ratio b -> a
locateCW Ratio b
r = (b, b) -> a
forall {b} {a}. (Num b, Num a, Ord b) => (b, b) -> a
go (Ratio b -> b
forall a. Ratio a -> a
numerator Ratio b
r, Ratio b -> b
forall a. Ratio a -> a
denominator Ratio b
r)
    go :: (b, b) -> a
go (b
1,b
1) = a
1
    go (b
a,b
b)
      | b
a b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b
b     = a
2 a -> a -> a
forall a. Num a => a -> a -> a
* (b, b) -> a
go (b
a, b
b b -> b -> b
forall a. Num a => a -> a -> a
- b
a)
      | Bool
otherwise = a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
2 a -> a -> a
forall a. Num a => a -> a -> a
* (b, b) -> a
go (b
a b -> 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]
-- >>> locate rat (-45 % 61)
-- 2540

rat :: IEnumeration Rational
rat :: IEnumeration Rational
rat = (Either () (Either Rational Rational) -> Rational)
-> (Rational -> Either () (Either Rational Rational))
-> IEnumeration (Either () (Either Rational Rational))
-> IEnumeration Rational
forall a b.
(a -> b) -> (b -> a) -> IEnumeration a -> IEnumeration b
mapE
  ((() -> Rational)
-> (Either Rational Rational -> Rational)
-> Either () (Either Rational Rational)
-> Rational
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Rational -> () -> Rational
forall a b. a -> b -> a
const Rational
0) ((Rational -> Rational)
-> (Rational -> Rational) -> Either Rational Rational -> Rational
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Rational -> Rational
forall a. a -> a
id Rational -> Rational
forall a. Num a => a -> a
negate))
  Rational -> Either () (Either Rational Rational)
forall {b}. (Num b, Ord b) => b -> Either () (Either b b)
unrat
  (IEnumeration ()
unit IEnumeration ()
-> IEnumeration (Either Rational Rational)
-> IEnumeration (Either () (Either Rational Rational))
forall a b.
IEnumeration a -> IEnumeration b -> IEnumeration (Either a b)
<+> (IEnumeration Rational
cw IEnumeration Rational
-> IEnumeration Rational -> IEnumeration (Either Rational Rational)
forall a b.
IEnumeration a -> IEnumeration b -> IEnumeration (Either a b)
<+> IEnumeration Rational
cw))
  where
    unrat :: b -> Either () (Either b b)
unrat b
0 = () -> Either () (Either b b)
forall a b. a -> Either a b
Left ()
    unrat b
r
      | b
r b -> b -> Bool
forall a. Ord a => a -> a -> Bool
> b
0     = Either b b -> Either () (Either b b)
forall a b. b -> Either a b
Right (b -> Either b b
forall a b. a -> Either a b
Left b
r)
      | Bool
otherwise = Either b b -> Either () (Either b b)
forall a b. b -> Either a b
Right (b -> Either b b
forall a b. b -> Either a b
Right (-b
r))

-- | 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 -> IEnumeration a -> IEnumeration a
takeE :: forall a. Integer -> IEnumeration a -> IEnumeration a
takeE Integer
k (IEnumeration Enumeration a
e a -> Integer
l) = Enumeration a -> (a -> Integer) -> IEnumeration a
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration (Integer -> Enumeration a -> Enumeration a
forall a. Integer -> Enumeration a -> Enumeration a
E.takeE Integer
k Enumeration a
e) a -> Integer
l

-- | 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 -> IEnumeration a -> IEnumeration a
dropE :: forall a. Integer -> IEnumeration a -> IEnumeration a
dropE Integer
k (IEnumeration Enumeration a
e a -> Integer
l) = Enumeration a -> (a -> Integer) -> IEnumeration a
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration (Integer -> Enumeration a -> Enumeration a
forall a. Integer -> Enumeration a -> Enumeration a
E.dropE Integer
k Enumeration a
e) (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
subtract (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max Integer
0 Integer
k) (Integer -> Integer) -> (a -> Integer) -> a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Integer
l)

-- | 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
--
-- toTree :: Either () (Tree, Tree) -> Tree
-- toTree = either (const L) (uncurry B)
--
-- fromTree :: Tree -> Either () (Tree, Tree)
-- fromTree L       = Left ()
-- fromTree (B l r) = Right (l,r)
--
-- treesBad :: IEnumeration Tree
-- treesBad = mapE toTree fromTree (unit '<+>' (treesBad '><' treesBad))
--
-- trees :: IEnumeration Tree
-- trees = infinite $ mapE toTree fromTree (unit '<+>' (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))
-- >>> select trees 123
-- B (B L (B L L)) (B (B L (B L L)) (B L (B L L)))
-- >>> locate trees (B (B L (B L L)) (B (B L (B L L)) (B L (B L L))))
-- 123

infinite :: IEnumeration a -> IEnumeration a
infinite :: forall a. IEnumeration a -> IEnumeration a
infinite (IEnumeration Enumeration a
e a -> Integer
l) = Enumeration a -> (a -> Integer) -> IEnumeration a
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration (Enumeration a -> Enumeration a
forall a. Enumeration a -> Enumeration a
E.infinite Enumeration a
e) a -> Integer
l

-- | 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
--   'Data.Enumeration.diagonal's.
--
--   Note that the type of this function is slightly different than
--   its counterpart in "Data.Enumeration": each enumerated value in
--   the output is tagged with an index indicating which input
--   enumeration it came from.  This is required to make the result
--   invertible, and is analogous to the way the output values of
--   '<+>' are tagged with 'Left' or 'Right'; in fact, 'interleave'
--   can be thought of as an iterated version of '<+>', but with a
--   more efficient implementation.

interleave :: IEnumeration (IEnumeration a) -> IEnumeration (Index, a)
interleave :: forall a.
IEnumeration (IEnumeration a) -> IEnumeration (Integer, a)
interleave IEnumeration (IEnumeration a)
e = IEnumeration
  { baseEnum :: Enumeration (Integer, a)
baseEnum = Cardinality
-> (Integer -> (Integer, a)) -> Enumeration (Integer, a)
forall a. Cardinality -> (Integer -> a) -> Enumeration a
E.mkEnumeration Cardinality
Infinite ((Integer -> (Integer, a)) -> Enumeration (Integer, a))
-> (Integer -> (Integer, a)) -> Enumeration (Integer, a)
forall a b. (a -> b) -> a -> b
$ \Integer
k ->
      let (Integer
i,Integer
j) = case IEnumeration (IEnumeration a) -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration (IEnumeration 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)
E.diagonal Integer
k
      in  (Integer
j, IEnumeration a -> Integer -> a
forall a. IEnumeration a -> Integer -> a
select (IEnumeration (IEnumeration a) -> Integer -> IEnumeration a
forall a. IEnumeration a -> Integer -> a
select IEnumeration (IEnumeration a)
e Integer
j) Integer
i)
  , locate :: (Integer, a) -> Integer
locate   = \(Integer
j, a
a) ->
      let i :: Integer
i = IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate (IEnumeration (IEnumeration a) -> Integer -> IEnumeration a
forall a. IEnumeration a -> Integer -> a
select IEnumeration (IEnumeration a)
e Integer
j) a
a
      in  case IEnumeration (IEnumeration a) -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration (IEnumeration a)
e of
            Finite Integer
n -> Integer
iInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
j
            Cardinality
Infinite -> (Integer, Integer) -> Integer
undiagonal (Integer
i,Integer
j)
  }

-- | 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.
--
--   Note that defining @zipWithE@ as in "Data.Enumeration" is not
--   possible since there would be no way to invert it in general.
--   However, one can use 'zipE' in combination with 'mapE' to achieve
--   a similar result.
--
-- >>> enumerate $ zipE nat (boundedEnum @Bool)
-- [(0,False),(1,True)]
--
-- >>> headD x = fromMaybe x . listToMaybe
-- >>> cs = mapE (uncurry replicate) (length &&& headD ' ') (zipE (finiteList [1..10]) (dropE 35 (boundedEnum @Char)))
-- >>> enumerate cs
-- ["#","$$","%%%","&&&&","'''''","((((((",")))))))","********","+++++++++",",,,,,,,,,,"]
-- >>> locate cs "********"
-- 7

zipE :: IEnumeration a -> IEnumeration b -> IEnumeration (a,b)
zipE :: forall a b. IEnumeration a -> IEnumeration b -> IEnumeration (a, b)
zipE IEnumeration a
ea IEnumeration b
eb = IEnumeration
  { baseEnum :: Enumeration (a, b)
baseEnum = Enumeration a -> Enumeration b -> Enumeration (a, b)
forall a b. Enumeration a -> Enumeration b -> Enumeration (a, b)
E.zipE (IEnumeration a -> Enumeration a
forall a. IEnumeration a -> Enumeration a
baseEnum IEnumeration a
ea) (IEnumeration b -> Enumeration b
forall a. IEnumeration a -> Enumeration a
baseEnum IEnumeration b
eb)
  , locate :: (a, b) -> Integer
locate   = IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
ea (a -> Integer) -> ((a, b) -> a) -> (a, b) -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> a
forall a b. (a, b) -> a
fst
  }

-- | 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 this has a different type than the version in
--   "Data.Enumeration".  Here we require the output to carry an
--   explicit 'Either' tag to make it invertible.
--
-- >>> enumerate . takeE 5 $ singleton 17 <+> nat
-- [Left 17,Right 0,Right 1,Right 2,Right 3]
--
-- >>> enumerate . takeE 5 $ nat <+> singleton 17
-- [Right 17,Left 0,Left 1,Left 2,Left 3]
--
-- >>> enumerate . takeE 5 $ nat <+> nat
-- [Left 0,Right 0,Left 1,Right 1,Left 2]
--
-- >>> locate (nat <+> nat) (Right 35)
-- 71

(<+>) :: IEnumeration a -> IEnumeration b -> IEnumeration (Either a b)
IEnumeration a
a <+> :: forall a b.
IEnumeration a -> IEnumeration b -> IEnumeration (Either a b)
<+> IEnumeration b
b = Enumeration (Either a b)
-> (Either a b -> Integer) -> IEnumeration (Either a b)
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration (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
<$> IEnumeration a -> Enumeration a
forall a. IEnumeration a -> Enumeration a
baseEnum IEnumeration 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
<$> IEnumeration b -> Enumeration b
forall a. IEnumeration a -> Enumeration a
baseEnum IEnumeration b
b) (IEnumeration a -> IEnumeration b -> Either a b -> Integer
forall a b.
IEnumeration a -> IEnumeration b -> Either a b -> Integer
locateEither IEnumeration a
a IEnumeration b
b)
  where
    locateEither :: IEnumeration a -> IEnumeration b -> (Either a b -> Index)
    locateEither :: forall a b.
IEnumeration a -> IEnumeration b -> Either a b -> Integer
locateEither IEnumeration a
a IEnumeration b
b = case (IEnumeration a -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration a
a, IEnumeration b -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration b
b) of
      (Finite Integer
k1, Cardinality
_) -> (a -> Integer) -> (b -> Integer) -> Either a b -> Integer
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
a) ((Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
k1) (Integer -> Integer) -> (b -> Integer) -> b -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IEnumeration b -> b -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration b
b)
      (Cardinality
_, Finite Integer
k2) -> (a -> Integer) -> (b -> Integer) -> Either a b -> Integer
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ((Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
k2) (Integer -> Integer) -> (a -> Integer) -> a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
a) (IEnumeration b -> b -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration b
b)
      (Cardinality, Cardinality)
_              -> (a -> Integer) -> (b -> Integer) -> Either a b -> Integer
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ((Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
2) (Integer -> Integer) -> (a -> Integer) -> a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
a) (Integer -> Integer
forall a. Enum a => a -> a
succ (Integer -> Integer) -> (b -> Integer) -> b -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
2) (Integer -> Integer) -> (b -> Integer) -> b -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IEnumeration b -> b -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration b
b)


-- | The other half of the isomorphism between \(\mathbb{N}\) and
--   \(\mathbb{N} \times \mathbb{N}\) which enumerates by diagonals:
--   turn a pair of natural numbers giving a position in the 2D grid
--   into the number in the cell, according to this numbering scheme:
--
--   @
--   0 1 3 6 ...
--   2 4 7
--   5 8
--   9
--   @
undiagonal :: (Integer, Integer) -> Integer
undiagonal :: (Integer, Integer) -> Integer
undiagonal (Integer
r,Integer
c) = (Integer
rInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
c) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
rInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
cInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
r

-- | 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 'Data.Enumeration.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)]
--
-- >>> locate (nat >< nat) (1,1)
-- 4
-- >>> locate (nat >< nat) (36,45)
-- 3357
--
--   Like ('<+>'), this operation is also not associative (not even up
--   to reassociating tuples).
(><) :: IEnumeration a -> IEnumeration b -> IEnumeration (a,b)
IEnumeration a
a >< :: forall a b. IEnumeration a -> IEnumeration b -> IEnumeration (a, b)
>< IEnumeration b
b = Enumeration (a, b) -> ((a, b) -> Integer) -> IEnumeration (a, b)
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration (IEnumeration a -> Enumeration a
forall a. IEnumeration a -> Enumeration a
baseEnum IEnumeration a
a Enumeration a -> Enumeration b -> Enumeration (a, b)
forall a b. Enumeration a -> Enumeration b -> Enumeration (a, b)
E.>< IEnumeration b -> Enumeration b
forall a. IEnumeration a -> Enumeration a
baseEnum IEnumeration b
b) (IEnumeration a -> IEnumeration b -> (a, b) -> Integer
forall a b. IEnumeration a -> IEnumeration b -> (a, b) -> Integer
locatePair IEnumeration a
a IEnumeration b
b)
  where
    locatePair :: IEnumeration a -> IEnumeration b -> ((a,b) -> Index)
    locatePair :: forall a b. IEnumeration a -> IEnumeration b -> (a, b) -> Integer
locatePair IEnumeration a
a IEnumeration b
b = case (IEnumeration a -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration a
a, IEnumeration b -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration b
b) of
      (Cardinality
_, Finite Integer
k2) -> \(a
x,b
y) -> Integer
k2 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
a a
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ IEnumeration b -> b -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration b
b b
y
      (Finite Integer
k1, Cardinality
_) -> \(a
x,b
y) -> Integer
k1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* IEnumeration b -> b -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration b
b b
y Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
a a
x
      (Cardinality, Cardinality)
_              -> \(a
x,b
y) -> (Integer, Integer) -> Integer
undiagonal (IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
a a
x, IEnumeration b -> b -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration b
b b
y)

------------------------------------------------------------
-- 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]
-- >>> locate (maybeOf (maybeOf (finiteList [1,2,3]))) (Just (Just 2))
-- 3
maybeOf :: IEnumeration a -> IEnumeration (Maybe a)
maybeOf :: forall a. IEnumeration a -> IEnumeration (Maybe a)
maybeOf IEnumeration a
a = (Either () a -> Maybe a)
-> (Maybe a -> Either () a)
-> IEnumeration (Either () a)
-> IEnumeration (Maybe a)
forall a b.
(a -> b) -> (b -> a) -> IEnumeration a -> IEnumeration b
mapE ((() -> Maybe a) -> (a -> Maybe a) -> Either () a -> Maybe a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe a -> () -> Maybe a
forall a b. a -> b -> a
const Maybe a
forall a. Maybe a
Nothing) a -> Maybe a
forall a. a -> Maybe a
Just) (Either () a -> (a -> Either () a) -> Maybe a -> Either () a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> Either () a
forall a b. a -> Either a b
Left ()) a -> Either () a
forall a b. b -> Either a b
Right) (IEnumeration ()
unit IEnumeration () -> IEnumeration a -> IEnumeration (Either () a)
forall a b.
IEnumeration a -> IEnumeration b -> IEnumeration (Either a b)
<+> IEnumeration a
a)

-- | Enumerae all possible values of type @Either a b@ with inner values
--   taken from the given enumerations.
--
--   Note that for invertible enumerations, 'eitherOf' is simply a
--   synonym for '<+>'.
--
-- >>> enumerate . takeE 6 $ eitherOf nat nat
-- [Left 0,Right 0,Left 1,Right 1,Left 2,Right 2]
eitherOf :: IEnumeration a -> IEnumeration b -> IEnumeration (Either a b)
eitherOf :: forall a b.
IEnumeration a -> IEnumeration b -> IEnumeration (Either a b)
eitherOf = IEnumeration a -> IEnumeration b -> IEnumeration (Either a b)
forall a b.
IEnumeration a -> IEnumeration b -> IEnumeration (Either a 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]]
-- >>> locate (listOf nat) [3,4,20,5,19]
-- 666270815854068922513792635440014
listOf :: IEnumeration a -> IEnumeration [a]
listOf :: forall a. IEnumeration a -> IEnumeration [a]
listOf IEnumeration a
a = case IEnumeration a -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration a
a of
  Finite Integer
0 -> [a] -> IEnumeration [a]
forall a. a -> IEnumeration a
singleton []
  Cardinality
_        -> IEnumeration [a]
listOfA
    where
      listOfA :: IEnumeration [a]
listOfA = IEnumeration [a] -> IEnumeration [a]
forall a. IEnumeration a -> IEnumeration a
infinite (IEnumeration [a] -> IEnumeration [a])
-> IEnumeration [a] -> IEnumeration [a]
forall a b. (a -> b) -> a -> b
$
        (Either () (a, [a]) -> [a])
-> ([a] -> Either () (a, [a]))
-> IEnumeration (Either () (a, [a]))
-> IEnumeration [a]
forall a b.
(a -> b) -> (b -> a) -> IEnumeration a -> IEnumeration b
mapE ((() -> [a]) -> ((a, [a]) -> [a]) -> Either () (a, [a]) -> [a]
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ([a] -> () -> [a]
forall a b. a -> b -> a
const []) ((a -> [a] -> [a]) -> (a, [a]) -> [a]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (:))) [a] -> Either () (a, [a])
forall {a}. [a] -> Either () (a, [a])
uncons (IEnumeration ()
unit IEnumeration ()
-> IEnumeration (a, [a]) -> IEnumeration (Either () (a, [a]))
forall a b.
IEnumeration a -> IEnumeration b -> IEnumeration (Either a b)
<+> (IEnumeration a
a IEnumeration a -> IEnumeration [a] -> IEnumeration (a, [a])
forall a b. IEnumeration a -> IEnumeration b -> IEnumeration (a, b)
>< IEnumeration [a]
listOfA))
      uncons :: [a] -> Either () (a, [a])
uncons []     = () -> Either () (a, [a])
forall a b. a -> Either a b
Left ()
      uncons (a
a:[a]
as) = (a, [a]) -> Either () (a, [a])
forall a b. b -> Either a b
Right (a
a, [a]
as)

-- | Enumerate all possible finite subsets of values from the given
--   enumeration.  The elements in each list will always occur in
--   increasing order of their index in the given enumeration.
--
-- >>> enumerate $ finiteSubsetOf (finite 3)
-- [[],[0],[1],[0,1],[2],[0,2],[1,2],[0,1,2]]
--
-- >>> locate (finiteSubsetOf nat) [2,3,6,8]
-- 332
-- >>> 332 == 2^8 + 2^6 + 2^3 + 2^2
-- True
finiteSubsetOf :: IEnumeration a -> IEnumeration [a]
finiteSubsetOf :: forall a. IEnumeration a -> IEnumeration [a]
finiteSubsetOf IEnumeration a
a = Enumeration [a] -> ([a] -> Integer) -> IEnumeration [a]
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration (Enumeration a -> Enumeration [a]
forall a. Enumeration a -> Enumeration [a]
E.finiteSubsetOf (IEnumeration a -> Enumeration a
forall a. IEnumeration a -> Enumeration a
baseEnum IEnumeration a
a)) [a] -> Integer
unpick
  where
    unpick :: [a] -> Integer
unpick = (Integer -> Integer -> Integer) -> Integer -> [Integer] -> Integer
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
(.|.) Integer
0 ([Integer] -> Integer) -> ([a] -> [Integer]) -> [a] -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Integer) -> [a] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map ((Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL`) (Int -> Integer) -> (a -> Int) -> a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> (a -> Integer) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
a)

-- | @finiteEnumerationOf n a@ creates an enumeration of all sequences
--   of exactly n items taken from the enumeration @a@.
--
-- >>> map E.enumerate . enumerate $ finiteEnumerationOf 2 (finite 3)
-- [[0,0],[0,1],[0,2],[1,0],[1,1],[1,2],[2,0],[2,1],[2,2]]
--
-- >>> map E.enumerate . take 10 . enumerate $ finiteEnumerationOf 3 nat
-- [[0,0,0],[0,0,1],[1,0,0],[0,1,0],[1,0,1],[2,0,0],[0,0,2],[1,1,0],[2,0,1],[3,0,0]]
finiteEnumerationOf :: Int -> IEnumeration a -> IEnumeration (Enumeration a)
finiteEnumerationOf :: forall a. Int -> IEnumeration a -> IEnumeration (Enumeration a)
finiteEnumerationOf Int
0 IEnumeration a
_ = Enumeration a -> IEnumeration (Enumeration a)
forall a. a -> IEnumeration a
singleton Enumeration a
forall a. Enumeration a
forall (f :: * -> *) a. Alternative f => f a
empty
finiteEnumerationOf Int
n IEnumeration a
a = case IEnumeration a -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration a
a of
  Finite Integer
k -> Enumeration (Enumeration a)
-> (Enumeration a -> Integer) -> IEnumeration (Enumeration a)
forall a. Enumeration a -> (a -> Integer) -> IEnumeration a
IEnumeration (Int -> Enumeration a -> Enumeration (Enumeration a)
forall a. Int -> Enumeration a -> Enumeration (Enumeration a)
E.finiteEnumerationOf Int
n (IEnumeration a -> Enumeration a
forall a. IEnumeration a -> Enumeration a
baseEnum IEnumeration a
a)) (Integer -> Enumeration a -> Integer
locateEnum Integer
k)
  Cardinality
Infinite -> (IEnumeration a
 -> IEnumeration (Enumeration a) -> IEnumeration (Enumeration a))
-> IEnumeration (Enumeration a)
-> [IEnumeration a]
-> IEnumeration (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 IEnumeration a
-> IEnumeration (Enumeration a) -> IEnumeration (Enumeration a)
forall a.
IEnumeration a
-> IEnumeration (Enumeration a) -> IEnumeration (Enumeration a)
prod (Enumeration a -> IEnumeration (Enumeration a)
forall a. a -> IEnumeration a
singleton Enumeration a
forall a. Enumeration a
forall (f :: * -> *) a. Alternative f => f a
empty) (Int -> IEnumeration a -> [IEnumeration a]
forall a. Int -> a -> [a]
replicate Int
n IEnumeration a
a)

  where
    locateEnum :: Integer -> Enumeration a -> Integer
locateEnum Integer
k = Integer -> [Integer] -> Integer
forall {t :: * -> *} {b}. (Foldable t, Num b) => b -> t b -> b
fromBase Integer
k ([Integer] -> Integer)
-> (Enumeration a -> [Integer]) -> Enumeration a -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> [Integer]
forall a. [a] -> [a]
reverse ([Integer] -> [Integer])
-> (Enumeration a -> [Integer]) -> Enumeration a -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Enumeration Integer -> [Integer]
forall a. Enumeration a -> [a]
E.enumerate (Enumeration Integer -> [Integer])
-> (Enumeration a -> Enumeration Integer)
-> Enumeration a
-> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Integer) -> Enumeration a -> Enumeration Integer
forall a b. (a -> b) -> Enumeration a -> Enumeration b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
a)

    fromBase :: b -> t b -> b
fromBase b
k = (b -> b -> b) -> b -> t b -> b
forall a b. (a -> b -> b) -> b -> t a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\b
d b
r -> b
d b -> b -> b
forall a. Num a => a -> a -> a
+ b
kb -> b -> b
forall a. Num a => a -> a -> a
*b
r) b
0

    prod :: IEnumeration a -> IEnumeration (Enumeration a) -> IEnumeration (Enumeration a)
    prod :: forall a.
IEnumeration a
-> IEnumeration (Enumeration a) -> IEnumeration (Enumeration a)
prod IEnumeration a
a IEnumeration (Enumeration a)
as = ((a, Enumeration a) -> Enumeration a)
-> (Enumeration a -> (a, Enumeration a))
-> IEnumeration (a, Enumeration a)
-> IEnumeration (Enumeration a)
forall a b.
(a -> b) -> (b -> a) -> IEnumeration a -> IEnumeration b
mapE (\(a
a,Enumeration a
e) -> a -> Enumeration a
forall a. a -> Enumeration a
E.singleton a
a 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
e) (\Enumeration a
e -> (Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
E.select Enumeration a
e Integer
0, Integer -> Enumeration a -> Enumeration a
forall a. Integer -> Enumeration a -> Enumeration a
E.dropE Integer
1 Enumeration a
e))
                  (IEnumeration a
a IEnumeration a
-> IEnumeration (Enumeration a) -> IEnumeration (a, Enumeration a)
forall a b. IEnumeration a -> IEnumeration b -> IEnumeration (a, b)
>< IEnumeration (Enumeration a)
as)

-- | @functionOf a b@ creates an enumeration of all functions taking
--   values from the enumeration @a@ and returning values from the
--   enumeration @b@.  As a precondition, @a@ must be finite;
--   otherwise @functionOf@ throws an error. There are two exceptions:
--   first, if @b@ has cardinality 1, we get an enumeration of exactly
--   one function which constantly returns the one element of @b@,
--   even if @a@ is infinite.  Second, if @b@ has cardinality 0, we
--   get a singleton enumeration if @a@ also has cardinality 0, and an
--   empty enumeration otherwise (even if @a@ is infinite).
--
-- >>> bbs = functionOf (boundedEnum @Bool) (boundedEnum @Bool)
-- >>> card bbs
-- Finite 4
-- >>> map (select bbs 2) [False, True]
-- [True,False]
-- >>> locate bbs not
-- 2
--
-- >>> locate (functionOf bbs (boundedEnum @Bool)) (\f -> f True)
-- 5
--
-- >>> n2u = functionOf nat unit
-- >>> card n2u
-- Finite 1
-- >>> (select n2u 0) 57
-- ()
--
-- >>> n2o = functionOf nat void
-- >>> card n2o
-- Finite 0
-- >>> o2o = functionOf void void
-- >>> card o2o
-- Finite 1
functionOf :: IEnumeration a -> IEnumeration b -> IEnumeration (a -> b)
functionOf :: forall a b.
IEnumeration a -> IEnumeration b -> IEnumeration (a -> b)
functionOf IEnumeration a
as IEnumeration b
bs = case IEnumeration b -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration b
bs of
  Finite Integer
1 -> (a -> b) -> IEnumeration (a -> b)
forall a. a -> IEnumeration a
singleton (\a
_ -> IEnumeration b -> Integer -> b
forall a. IEnumeration a -> Integer -> a
select IEnumeration b
bs Integer
0)   -- 1^x = 1
  Finite Integer
0 -> case IEnumeration a -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration a
as of                 -- 0^0 = 1, 0^x = 0
    Finite Integer
0 -> (a -> b) -> IEnumeration (a -> b)
forall a. a -> IEnumeration a
singleton (\a
_ -> [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"called function with empty domain")
    Cardinality
_        -> IEnumeration (a -> b)
forall a. IEnumeration a
void
  Cardinality
_        -> case IEnumeration a -> Cardinality
forall a. IEnumeration a -> Cardinality
card IEnumeration a
as of
    Cardinality
Infinite -> [Char] -> IEnumeration (a -> b)
forall a. HasCallStack => [Char] -> a
error [Char]
"functionOf with infinite domain"
    Finite Integer
n -> (Enumeration b -> a -> b)
-> ((a -> b) -> Enumeration b)
-> IEnumeration (Enumeration b)
-> IEnumeration (a -> b)
forall a b.
(a -> b) -> (b -> a) -> IEnumeration a -> IEnumeration b
mapE Enumeration b -> a -> b
forall {a}. Enumeration a -> a -> a
toFunc (a -> b) -> Enumeration b
forall {b}. (a -> b) -> Enumeration b
fromFunc (Int -> IEnumeration b -> IEnumeration (Enumeration b)
forall a. Int -> IEnumeration a -> IEnumeration (Enumeration a)
finiteEnumerationOf (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n) IEnumeration b
bs)

  where
    toFunc :: Enumeration a -> a -> a
toFunc Enumeration a
bTuple a
a = Enumeration a -> Integer -> a
forall a. Enumeration a -> Integer -> a
E.select Enumeration a
bTuple (IEnumeration a -> a -> Integer
forall a. IEnumeration a -> a -> Integer
locate IEnumeration a
as a
a)
    fromFunc :: (a -> b) -> Enumeration b
fromFunc a -> b
f = a -> b
f (a -> b) -> Enumeration a -> Enumeration b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IEnumeration a -> Enumeration a
forall a. IEnumeration a -> Enumeration a
baseEnum IEnumeration a
as