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

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

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.CoEnumeration
-- Copyright   :  Brent Yorgey, Koji Miyazato
-- Maintainer  :  byorgey@gmail.com
-- 
-- A /coenumeration/ is a function from values to finite or countably infinite
-- sets, canonically represented by non-negative integers less than its cardinality.
-- 
-- Alternatively, a coenumeration can be thought of as a classification of values
-- into finite or countably infinite classes, with each class labelled with
-- integers.
-- 
-- This module provides many ways to construct @CoEnumeration@ values,
-- and most of them are implemented as inverses of enumerations made with
-- functions in "Data.Enumeration".
-- 
-- == Example
-- 
-- Through examples of this module, "Data.Enumeration" module is
-- referred by alias @E@.
-- 
-- > import qualified Data.Enumeration as E
-- 
-- >>> take 5 . drop 5 $ E.enumerate (E.listOf E.nat)
-- [[1,0],[2],[0,1],[1,0,0],[2,0]]
-- >>> locate (listOf nat) <$> [[1,0],[2],[0,1],[1,0,0],[2,0]]
-- [5,6,7,8,9]
--
-- >>> locate (listOf nat) [3,1,4,1,5,9,2]
-- 78651719792187121765701606023038613403478037124236785850350
-- >>> E.select (E.listOf E.nat) 78651719792187121765701606023038613403478037124236785850350
-- [3,1,4,1,5,9,2]
module Data.CoEnumeration
  ( -- * Coenumerations
    CoEnumeration(), card, locate, isFinite
  , unsafeMkCoEnumeration

    -- * Cardinality and Index
  , Index, Cardinality(..)

    -- * Primitive coenumerations
  , unit, lost
  , boundedEnum
  , nat
  , int
  , cw
  , rat

    -- * Coenumeration combinators
  , takeC, dropC, modulo, overlayC
  , infinite
  , (><), (<+>)
  , maybeOf, eitherOf, listOf, finiteSubsetOf
  , finiteFunctionOf

    -- * Utilities
  , unList, unSet
  ) where

import Data.Void
import Data.Bits
import Data.List (foldl')
import Data.Ratio

import Data.Functor.Contravariant
import Data.Functor.Contravariant.Divisible(lost, Divisible(..), Decidable(..))

import Data.Enumeration (Index, Cardinality(..))
import Data.Enumeration.Invertible (undiagonal)


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

-- $setup
-- >>> :set -XTypeApplications
-- >>> import qualified Data.Enumeration as E

-- | A /coenumeration/ is a function from values to finite or countably infinite
-- sets, canonically represented by non-negative integers less than its cardinality.
-- 
-- Alternatively, a coenumeration can be thought of as a classification of values
-- into finite or countably infinite classes, with each class labelled with
-- integers.
-- 
-- 'CoEnumeration' is an instance of the following type classes:
--
-- * 'Contravariant' (you can change the type of base values contravariantly)
-- * 'Divisible' (representing Cartesian product of finite number of coenumerations)
--
--     * Binary cartesian product ('><')
--     * Coenumeration onto singleton set as an unit ('unit')
--
-- * 'Decidable' (representing disjoint union of finite number of coenumerations)
--
--     * Binary disjoint union ('<+>')
--     * Coenumeration of uninhabited type 'Void'. It's not exported directly,
--       but only through a method of the class
--       
--         'lose' @:: Decidable f => (a -> Void) -> f Void@
--       
--         or
--       
--         'lost' @:: Decidable f => f Void@.
data CoEnumeration a = CoEnumeration
  { -- | Get the cardinality of a coenumeration.
    --   Under \"classification\" interpretation,
    --   it is the cardinality of the set of classes.
    forall a. CoEnumeration a -> Cardinality
card :: Cardinality

    -- | Compute the index of a particular value.
  , forall a. CoEnumeration a -> a -> Index
locate :: a -> Index
  }

-- | Returns if the the cardinality of coenumeration is finite.
isFinite :: CoEnumeration a -> Bool
isFinite :: forall a. CoEnumeration a -> Bool
isFinite = (Cardinality
Infinite Cardinality -> Cardinality -> Bool
forall a. Eq a => a -> a -> Bool
/=) (Cardinality -> Bool)
-> (CoEnumeration a -> Cardinality) -> CoEnumeration a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card

-- | Constructs a coenumeration.
--
--   To construct valid coenumeration by @unsafeMkCoEnumeration n f@,
--   for all @x :: a@, it must satisfy @(Finite (f x) < n)@.
--   
--   This functions does not (and never could) check the validity
--   of its arguments.
unsafeMkCoEnumeration :: Cardinality -> (a -> Index) -> CoEnumeration a
unsafeMkCoEnumeration :: forall a. Cardinality -> (a -> Index) -> CoEnumeration a
unsafeMkCoEnumeration = Cardinality -> (a -> Index) -> CoEnumeration a
forall a. Cardinality -> (a -> Index) -> CoEnumeration a
CoEnumeration

instance Contravariant CoEnumeration where
  contramap :: forall a' a. (a' -> a) -> CoEnumeration a -> CoEnumeration a'
contramap a' -> a
f CoEnumeration a
e = CoEnumeration a
e{ locate = locate e . f }

-- | Associativity of 'divide' is maintained only when
--   arguments are finite.
instance Divisible CoEnumeration where
  divide :: forall a b c.
(a -> (b, c))
-> CoEnumeration b -> CoEnumeration c -> CoEnumeration a
divide a -> (b, c)
f CoEnumeration b
a CoEnumeration c
b = (a -> (b, c)) -> CoEnumeration (b, c) -> CoEnumeration a
forall a' a. (a' -> a) -> CoEnumeration a -> CoEnumeration a'
forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap a -> (b, c)
f (CoEnumeration (b, c) -> CoEnumeration a)
-> CoEnumeration (b, c) -> CoEnumeration a
forall a b. (a -> b) -> a -> b
$ CoEnumeration b
a CoEnumeration b -> CoEnumeration c -> CoEnumeration (b, c)
forall a b.
CoEnumeration a -> CoEnumeration b -> CoEnumeration (a, b)
>< CoEnumeration c
b
  conquer :: forall a. CoEnumeration a
conquer = CoEnumeration a
forall a. CoEnumeration a
unit

-- | Associativity of 'choose' is maintained only when
--   arguments are finite.
instance Decidable CoEnumeration where
  choose :: forall a b c.
(a -> Either b c)
-> CoEnumeration b -> CoEnumeration c -> CoEnumeration a
choose a -> Either b c
f CoEnumeration b
a CoEnumeration c
b = (a -> Either b c) -> CoEnumeration (Either b c) -> CoEnumeration a
forall a' a. (a' -> a) -> CoEnumeration a -> CoEnumeration a'
forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap a -> Either b c
f (CoEnumeration (Either b c) -> CoEnumeration a)
-> CoEnumeration (Either b c) -> CoEnumeration a
forall a b. (a -> b) -> a -> b
$ CoEnumeration b
a CoEnumeration b -> CoEnumeration c -> CoEnumeration (Either b c)
forall a b.
CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
<+> CoEnumeration c
b
  lose :: forall a. (a -> Void) -> CoEnumeration a
lose a -> Void
f = (a -> Void) -> CoEnumeration Void -> CoEnumeration a
forall a' a. (a' -> a) -> CoEnumeration a -> CoEnumeration a'
forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap a -> Void
f CoEnumeration Void
void

-- | Coenumeration to the singleton set.
--
-- >>> card unit
-- Finite 1
-- >>> locate unit True
-- 0
-- >>> locate unit (3 :: Int)
-- 0
-- >>> locate unit (cos :: Float -> Float)
-- 0
unit :: CoEnumeration a
unit :: forall a. CoEnumeration a
unit = CoEnumeration{ card :: Cardinality
card = Cardinality
1, locate :: a -> Index
locate = Index -> a -> Index
forall a b. a -> b -> a
const Index
0 }

-- | Coenumeration of an uninhabited type 'Void'.
--
-- >>> card void
-- Finite 0
-- 
-- Note that when a coenumeration of a type @t@ has cardinality 0,
-- it should indicate /No/ value of @t@ can be created without
-- using bottoms like @undefined@.
void :: CoEnumeration Void
void :: CoEnumeration Void
void = CoEnumeration{ card :: Cardinality
card = Cardinality
0, locate :: Void -> Index
locate = Index -> Void -> Index
forall a b. a -> b -> a
const ([Char] -> Index
forall a. HasCallStack => [Char] -> a
error [Char]
"locate void") }

-- | An inverse of forward 'E.boundedEnum'
boundedEnum :: forall a. (Enum a, Bounded a) => CoEnumeration a
boundedEnum :: forall a. (Enum a, Bounded a) => CoEnumeration a
boundedEnum = CoEnumeration{ card :: Cardinality
card = Cardinality
size, locate :: a -> Index
locate = a -> Index
loc }
  where loc :: a -> Index
loc = Int -> Index
forall a. Integral a => a -> Index
toInteger (Int -> Index) -> (a -> Int) -> a -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
lo (Int -> Int) -> (a -> Int) -> a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int
forall a. Enum a => a -> Int
fromEnum
        lo :: Int
lo = a -> Int
forall a. Enum a => a -> Int
fromEnum (forall a. Bounded a => a
minBound @a)
        hi :: Int
hi = a -> Int
forall a. Enum a => a -> Int
fromEnum (forall a. Bounded a => a
maxBound @a)
        size :: Cardinality
size = Index -> Cardinality
Finite (Index -> Cardinality) -> Index -> Cardinality
forall a b. (a -> b) -> a -> b
$ Index
1 Index -> Index -> Index
forall a. Num a => a -> a -> a
+ Int -> Index
forall a. Integral a => a -> Index
toInteger Int
hi Index -> Index -> Index
forall a. Num a => a -> a -> a
- Int -> Index
forall a. Integral a => a -> Index
toInteger Int
lo

-- | 'nat' is an inverse of forward enumeration 'E.nat'.
--  
-- For a negative integer @x@ which 'E.nat' doesn't enumerate,
-- @locate nat x@ returns the same index to the absolute value of @x@,
-- i.e. @locate nat x == locate nat (abs x)@.
-- 
-- >>> locate nat <$> [-3 .. 3]
-- [3,2,1,0,1,2,3]
nat :: CoEnumeration Integer
nat :: CoEnumeration Index
nat = CoEnumeration{ card :: Cardinality
card = Cardinality
Infinite, locate :: Index -> Index
locate = Index -> Index
forall a. Num a => a -> a
abs }

-- | 'int' is the inverse of forward enumeration 'E.int', which is
--   all integers ordered as the sequence @0, 1, -1, 2, -2, ...@
-- 
-- >>> locate int <$> [1, 2, 3, 4, 5]
-- [1,3,5,7,9]
-- >>> locate int <$> [0, -1 .. -5]
-- [0,2,4,6,8,10]
int :: CoEnumeration Integer
int :: CoEnumeration Index
int = CoEnumeration{ card :: Cardinality
card = Cardinality
Infinite, locate :: Index -> Index
locate = Index -> Index
integerToNat }
  where
    integerToNat :: Integer -> Integer
    integerToNat :: Index -> Index
integerToNat Index
n
      | Index
n Index -> Index -> Bool
forall a. Ord a => a -> a -> Bool
<= Index
0    = Index
2 Index -> Index -> Index
forall a. Num a => a -> a -> a
* Index -> Index
forall a. Num a => a -> a
negate Index
n
      | Bool
otherwise = Index
2 Index -> Index -> Index
forall a. Num a => a -> a -> a
* Index
n Index -> Index -> Index
forall a. Num a => a -> a -> a
- Index
1

-- | 'cw' is an inverse of forward enumeration 'E.cw'.
--
-- Because 'E.cw' only enumerates positive 'Rational' values,
-- @locate cw x@ for zero or less rational number @x@ could be arbitrary.
-- 
-- This implementation chose @locate cw x = 33448095@ for all @x <= 0@, which is
-- the index of @765 % 4321@, rather than returning @undefined@.
-- 
-- >>> locate cw <$> [1 % 1, 1 % 2, 2 % 1, 1 % 3, 3 % 2]
-- [0,1,2,3,4]
-- >>> locate cw (765 % 4321)
-- 33448095
-- >>> locate cw 0
-- 33448095
cw :: CoEnumeration Rational
cw :: CoEnumeration Rational
cw = CoEnumeration{ card :: Cardinality
card = Cardinality
Infinite, locate :: Rational -> Index
locate = Rational -> Index
forall {a} {a}. (Ord a, Num a, Num a) => Ratio a -> a
locateCW }
  where
    locateCW :: Ratio a -> a
locateCW Ratio a
x = case Ratio a -> a
forall a. Ratio a -> a
numerator Ratio a
x of
      a
n | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0     -> a -> a -> [a] -> a
forall {a} {a}. (Num a, Num a, Ord a) => a -> a -> [a] -> a
go a
n (Ratio a -> a
forall a. Ratio a -> a
denominator Ratio a
x) [] a -> a -> a
forall a. Num a => a -> a -> a
- a
1
        | Bool
otherwise -> a
33448095 {- Magic number, see the haddock above -}
    
    go :: a -> a -> [a] -> a
go a
1 a
1 [a]
acc = (a -> a -> a) -> a -> [a] -> a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\a
i a
b -> a
2 a -> a -> a
forall a. Num a => a -> a -> a
* a
i a -> a -> a
forall a. Num a => a -> a -> a
+ a
b) a
1 [a]
acc
    go a
a a
b [a]
acc
      | a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
b = a -> a -> [a] -> a
go (a
a a -> a -> a
forall a. Num a => a -> a -> a
- a
b) a
b (a
1 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc)
      | a
a a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
b = a -> a -> [a] -> a
go a
a (a
b a -> a -> a
forall a. Num a => a -> a -> a
- a
a) (a
0 a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc)
      | Bool
otherwise = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"cw: locateCW: Never reach here!"

-- | 'rat' is the inverse of forward enumeration 'E.rat'.
--
-- >>> let xs = E.enumerate . E.takeE 6 $ E.rat
-- >>> (xs, locate rat <$> xs)
-- ([0 % 1,1 % 1,(-1) % 1,1 % 2,(-1) % 2,2 % 1],[0,1,2,3,4,5])
-- >>> locate rat (E.select E.rat 1000)
-- 1000
rat :: CoEnumeration Rational
rat :: CoEnumeration Rational
rat = (Rational -> Maybe (Either Rational Rational))
-> CoEnumeration (Maybe (Either Rational Rational))
-> CoEnumeration Rational
forall a' a. (a' -> a) -> CoEnumeration a -> CoEnumeration a'
forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap Rational -> Maybe (Either Rational Rational)
caseBySign (CoEnumeration (Maybe (Either Rational Rational))
 -> CoEnumeration Rational)
-> CoEnumeration (Maybe (Either Rational Rational))
-> CoEnumeration Rational
forall a b. (a -> b) -> a -> b
$ CoEnumeration (Either Rational Rational)
-> CoEnumeration (Maybe (Either Rational Rational))
forall a. CoEnumeration a -> CoEnumeration (Maybe a)
maybeOf (CoEnumeration Rational
cw CoEnumeration Rational
-> CoEnumeration Rational
-> CoEnumeration (Either Rational Rational)
forall a b.
CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
<+> CoEnumeration Rational
cw)
  where
    caseBySign :: Rational -> Maybe (Either Rational Rational)
    caseBySign :: Rational -> Maybe (Either Rational Rational)
caseBySign Rational
x = case Rational -> Rational -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Rational
x Rational
0 of
      Ordering
LT -> Either Rational Rational -> Maybe (Either Rational Rational)
forall a. a -> Maybe a
Just (Rational -> Either Rational Rational
forall a b. b -> Either a b
Right (Rational -> Rational
forall a. Num a => a -> a
negate Rational
x))
      Ordering
EQ -> Maybe (Either Rational Rational)
forall a. Maybe a
Nothing
      Ordering
GT -> Either Rational Rational -> Maybe (Either Rational Rational)
forall a. a -> Maybe a
Just (Rational -> Either Rational Rational
forall a b. a -> Either a b
Left Rational
x)

-- | Sets the cardinality of given coenumeration to 'Infinite'
infinite :: CoEnumeration a -> CoEnumeration a
infinite :: forall a. CoEnumeration a -> CoEnumeration a
infinite CoEnumeration a
e = CoEnumeration a
e{ card = Infinite }

-- | Cartesian product of coenumeration, made to be an inverse of
--   cartesian product of enumeration '(E.><)'.
--   
-- >>> let a  = E.finite 3 E.>< (E.boundedEnum @Bool)
-- >>> let a' = modulo 3     >< boundedEnum @Bool
-- >>> (E.enumerate a, locate a' <$> E.enumerate a)
-- ([(0,False),(0,True),(1,False),(1,True),(2,False),(2,True)],[0,1,2,3,4,5])
--
-- This operation is not associative if and only if one of arguments
-- is not finite.
(><) :: CoEnumeration a -> CoEnumeration b -> CoEnumeration (a,b)
CoEnumeration a
e1 >< :: forall a b.
CoEnumeration a -> CoEnumeration b -> CoEnumeration (a, b)
>< CoEnumeration b
e2 = CoEnumeration{ card :: Cardinality
card = Cardinality
n1 Cardinality -> Cardinality -> Cardinality
forall a. Num a => a -> a -> a
* Cardinality
n2, locate :: (a, b) -> Index
locate = (a, b) -> Index
locatePair }
  where
    n1 :: Cardinality
n1 = CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e1
    n2 :: Cardinality
n2 = CoEnumeration b -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration b
e2
    locatePair :: (a, b) -> Index
locatePair = case (Cardinality
n1, Cardinality
n2) of
      (Cardinality
_,          Finite Index
n2') -> \(a
a,b
b) -> CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e1 a
a Index -> Index -> Index
forall a. Num a => a -> a -> a
* Index
n2' Index -> Index -> Index
forall a. Num a => a -> a -> a
+ CoEnumeration b -> b -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration b
e2 b
b
      (Finite Index
n1', Cardinality
Infinite)   -> \(a
a,b
b) -> CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e1 a
a Index -> Index -> Index
forall a. Num a => a -> a -> a
+ CoEnumeration b -> b -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration b
e2 b
b Index -> Index -> Index
forall a. Num a => a -> a -> a
* Index
n1'
      (Cardinality
Infinite,   Cardinality
Infinite)   -> \(a
a,b
b) -> (Index, Index) -> Index
undiagonal (CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e1 a
a, CoEnumeration b -> b -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration b
e2 b
b)

-- | Sum, or disjoint union, of two coenumerations.
--
--   It corresponds to disjoint union of enumerations 'E.eitherOf'.
--   
--   Its type can't be @CoEnumeration a -> CoEnumeration a -> CoEnumeration a@,
--   like 'Data.Enumeration.Enumeration' which is covariant functor,
--   because @CoEnumeration@ is 'Contravariant' functor.
--   
-- >>> let a  = E.finite 3 `E.eitherOf` (E.boundedEnum @Bool)
-- >>> let a' = modulo 3    <+>          boundedEnum @Bool
-- >>> (E.enumerate a, locate a' <$> E.enumerate a)
-- ([Left 0,Left 1,Left 2,Right False,Right True],[0,1,2,3,4])
--
-- This operation is not associative if and only if one of arguments
-- is not finite.
(<+>) :: CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
CoEnumeration a
e1 <+> :: forall a b.
CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
<+> CoEnumeration b
e2 = CoEnumeration{ card :: Cardinality
card = Cardinality
n1 Cardinality -> Cardinality -> Cardinality
forall a. Num a => a -> a -> a
+ Cardinality
n2, locate :: Either a b -> Index
locate = Either a b -> Index
locateEither }
  where
    n1 :: Cardinality
n1 = CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e1
    n2 :: Cardinality
n2 = CoEnumeration b -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration b
e2
    locateEither :: Either a b -> Index
locateEither = case (Cardinality
n1, Cardinality
n2) of
      (Finite Index
n1', Cardinality
_)          -> (a -> Index) -> (b -> Index) -> Either a b -> Index
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e1) ((Index
n1' Index -> Index -> Index
forall a. Num a => a -> a -> a
+) (Index -> Index) -> (b -> Index) -> b -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoEnumeration b -> b -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration b
e2)
      (Cardinality
Infinite,   Finite Index
n2') -> (a -> Index) -> (b -> Index) -> Either a b -> Index
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ((Index
n2' Index -> Index -> Index
forall a. Num a => a -> a -> a
+) (Index -> Index) -> (a -> Index) -> a -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e1) (CoEnumeration b -> b -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration b
e2)
      (Cardinality
Infinite,   Cardinality
Infinite)   -> (a -> Index) -> (b -> Index) -> Either a b -> Index
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ((Index -> Index -> Index
forall a. Num a => a -> a -> a
*Index
2) (Index -> Index) -> (a -> Index) -> a -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e1) (Index -> Index
forall a. Enum a => a -> a
succ (Index -> Index) -> (b -> Index) -> b -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Index -> Index -> Index
forall a. Num a => a -> a -> a
*Index
2) (Index -> Index) -> (b -> Index) -> b -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoEnumeration b -> b -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration b
e2)

-- |
--
-- >>> locate (dropC 3 nat) <$> [0..5]
-- [0,0,0,0,1,2]
dropC :: Integer -> CoEnumeration a -> CoEnumeration a
dropC :: forall a. Index -> CoEnumeration a -> CoEnumeration a
dropC Index
k CoEnumeration a
e
  | Index
k Index -> Index -> Bool
forall a. Eq a => a -> a -> Bool
== Index
0      = CoEnumeration a
e
  | CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e Cardinality -> Cardinality -> Bool
forall a. Eq a => a -> a -> Bool
== Cardinality
0 = CoEnumeration a
e
  | CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e Cardinality -> Cardinality -> Bool
forall a. Ord a => a -> a -> Bool
<= Index -> Cardinality
Finite Index
k = [Char] -> CoEnumeration a
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible empty coenumeration"
  | Bool
otherwise = CoEnumeration{ card :: Cardinality
card = Cardinality
size, locate :: a -> Index
locate = a -> Index
loc }
  where
    size :: Cardinality
size = CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e Cardinality -> Cardinality -> Cardinality
forall a. Num a => a -> a -> a
- Index -> Cardinality
Finite Index
k
    loc :: a -> Index
loc = Index -> Index -> Index
forall a. Ord a => a -> a -> a
max Index
0 (Index -> Index) -> (a -> Index) -> a -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index -> Index -> Index
forall a. Num a => a -> a -> a
subtract Index
k (Index -> Index) -> (a -> Index) -> a -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e

-- |
-- >>> locate (takeC 3 nat) <$> [0..5]
-- [0,1,2,2,2,2]
takeC :: Integer -> CoEnumeration a -> CoEnumeration a
takeC :: forall a. Index -> CoEnumeration a -> CoEnumeration a
takeC Index
k
  | Index
k Index -> Index -> Bool
forall a. Ord a => a -> a -> Bool
<= Index
0 = CoEnumeration a -> CoEnumeration a
forall a. CoEnumeration a -> CoEnumeration a
checkEmpty
  | Bool
otherwise = CoEnumeration a -> CoEnumeration a
forall a. CoEnumeration a -> CoEnumeration a
aux
  where
    aux :: CoEnumeration a -> CoEnumeration a
aux CoEnumeration a
e =
      let size :: Cardinality
size = Cardinality -> Cardinality -> Cardinality
forall a. Ord a => a -> a -> a
min (Index -> Cardinality
Finite Index
k) (CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e)
          loc :: a -> Index
loc = Index -> Index -> Index
forall a. Ord a => a -> a -> a
min (Index
kIndex -> Index -> Index
forall a. Num a => a -> a -> a
-Index
1) (Index -> Index) -> (a -> Index) -> a -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e
      in CoEnumeration{ card :: Cardinality
card = Cardinality
size, locate :: a -> Index
locate = a -> Index
loc }

checkEmpty :: CoEnumeration a -> CoEnumeration a
checkEmpty :: forall a. CoEnumeration a -> CoEnumeration a
checkEmpty CoEnumeration a
e
  | CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e Cardinality -> Cardinality -> Bool
forall a. Eq a => a -> a -> Bool
== Cardinality
0 = CoEnumeration a
e
  | Bool
otherwise   = [Char] -> CoEnumeration a
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible empty coenumeration"

-- |
-- >>> locate (modulo 3) <$> [0..7]
-- [0,1,2,0,1,2,0,1]
-- >>> locate (modulo 3) (-4)
-- 2
modulo :: Integer -> CoEnumeration Integer
modulo :: Index -> CoEnumeration Index
modulo Index
n
  | Index
n Index -> Index -> Bool
forall a. Ord a => a -> a -> Bool
<= Index
0    = [Char] -> CoEnumeration Index
forall a. HasCallStack => [Char] -> a
error ([Char] -> CoEnumeration Index) -> [Char] -> CoEnumeration Index
forall a b. (a -> b) -> a -> b
$ [Char]
"modulo: invalid argument " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Index -> [Char]
forall a. Show a => a -> [Char]
show Index
n
  | Bool
otherwise = CoEnumeration{ card :: Cardinality
card = Index -> Cardinality
Finite Index
n, locate :: Index -> Index
locate = (Index -> Index -> Index
forall a. Integral a => a -> a -> a
`mod` Index
n) }

-- | @overlayC a b@ combines two coenumerations in parallel, sharing
--   indices of two coenumerations.
--
--   The resulting coenumeration has cardinality of the larger of the
--   two arguments.
overlayC :: CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
overlayC :: forall a b.
CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
overlayC CoEnumeration a
e1 CoEnumeration b
e2 = CoEnumeration{
    card :: Cardinality
card = Cardinality -> Cardinality -> Cardinality
forall a. Ord a => a -> a -> a
max (CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e1) (CoEnumeration b -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration b
e2)
  , locate :: Either a b -> Index
locate = (a -> Index) -> (b -> Index) -> Either a b -> Index
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e1) (CoEnumeration b -> b -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration b
e2)
  }

-- | The inverse of forward 'E.maybeOf'
maybeOf :: CoEnumeration a -> CoEnumeration (Maybe a)
maybeOf :: forall a. CoEnumeration a -> CoEnumeration (Maybe a)
maybeOf CoEnumeration a
e = (Maybe a -> Either () a)
-> CoEnumeration (Either () a) -> CoEnumeration (Maybe a)
forall a' a. (a' -> a) -> CoEnumeration a -> CoEnumeration a'
forall (f :: * -> *) a' a.
Contravariant f =>
(a' -> a) -> f a -> f a'
contramap (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) (CoEnumeration (Either () a) -> CoEnumeration (Maybe a))
-> CoEnumeration (Either () a) -> CoEnumeration (Maybe a)
forall a b. (a -> b) -> a -> b
$ CoEnumeration ()
forall a. CoEnumeration a
unit CoEnumeration () -> CoEnumeration a -> CoEnumeration (Either () a)
forall a b.
CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
<+> CoEnumeration a
e

-- | Synonym of '(<+>)'
eitherOf :: CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
eitherOf :: forall a b.
CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
eitherOf = CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
forall a b.
CoEnumeration a -> CoEnumeration b -> CoEnumeration (Either a b)
(<+>)

-- | Coenumerate all possible finite lists using given coenumeration.
--
--   If a coenumeration @a@ is the inverse of enumeration @b@,
--   'listOf' @a@ is the inverse of forward enumeration 'E.listOf' @b@.
-- 
-- >>> E.enumerate . E.takeE 6 $ E.listOf E.nat
-- [[],[0],[0,0],[1],[0,0,0],[1,0]]
-- >>> locate (listOf nat) <$> [[],[0],[0,0],[1],[0,0,0],[1,0]]
-- [0,1,2,3,4,5]
-- >>> E.select (E.listOf E.nat) 1000000
-- [1008,26,0]
-- >>> locate (listOf nat) [1008,26,0]
-- 1000000
listOf :: CoEnumeration a -> CoEnumeration [a]
listOf :: forall a. CoEnumeration a -> CoEnumeration [a]
listOf CoEnumeration a
e = CoEnumeration{ card :: Cardinality
card = Cardinality
size, locate :: [a] -> Index
locate = [a] -> Index
locateList }
  where
    n :: Cardinality
n = CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e
    size :: Cardinality
size | Cardinality
n Cardinality -> Cardinality -> Bool
forall a. Eq a => a -> a -> Bool
== Cardinality
0    = Cardinality
1
         | Bool
otherwise = Cardinality
Infinite
    locateList :: [a] -> Index
locateList = Cardinality -> [Index] -> Index
unList Cardinality
n ([Index] -> Index) -> ([a] -> [Index]) -> [a] -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Index) -> [a] -> [Index]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e)

unList :: Cardinality -> [Index] -> Index
unList :: Cardinality -> [Index] -> Index
unList (Finite Index
k) = (Index -> Index -> Index) -> Index -> [Index] -> Index
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Index
n Index
a -> Index
1 Index -> Index -> Index
forall a. Num a => a -> a -> a
+ (Index
a Index -> Index -> Index
forall a. Num a => a -> a -> a
+ Index
n Index -> Index -> Index
forall a. Num a => a -> a -> a
* Index
k)) Index
0 ([Index] -> Index) -> ([Index] -> [Index]) -> [Index] -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Index] -> [Index]
forall a. [a] -> [a]
reverse
unList Cardinality
Infinite   = (Index -> Index -> Index) -> Index -> [Index] -> Index
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Index
n Index
a -> Index
1 Index -> Index -> Index
forall a. Num a => a -> a -> a
+ (Index, Index) -> Index
undiagonal (Index
a, Index
n)) Index
0 ([Index] -> Index) -> ([Index] -> [Index]) -> [Index] -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Index] -> [Index]
forall a. [a] -> [a]
reverse

-- | An inverse of 'E.finiteSubsetOf'.
--
--   Given a coenumeration of @a@, make a coenumeration of finite sets of
--   @a@, by ignoring order and repetition from @[a]@.
-- 
-- >>> as = take 11 . E.enumerate $ E.finiteSubsetOf E.nat
-- >>> (as, locate (finiteSubsetOf nat) <$> as)
-- ([[],[0],[1],[0,1],[2],[0,2],[1,2],[0,1,2],[3],[0,3],[1,3]],[0,1,2,3,4,5,6,7,8,9,10])
finiteSubsetOf :: CoEnumeration a -> CoEnumeration [a]
finiteSubsetOf :: forall a. CoEnumeration a -> CoEnumeration [a]
finiteSubsetOf CoEnumeration a
e = CoEnumeration{ card :: Cardinality
card = Cardinality
size, locate :: [a] -> Index
locate = [Index] -> Index
unSet ([Index] -> Index) -> ([a] -> [Index]) -> [a] -> Index
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Index) -> [a] -> [Index]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
e) }
  where
    size :: Cardinality
size = case CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
e of
      Finite Index
k -> Index -> Cardinality
Finite (Index
2 Index -> Index -> Index
forall a b. (Num a, Integral b) => a -> b -> a
^ Index
k)
      Cardinality
Infinite -> Cardinality
Infinite

unSet :: [Index] -> Index
unSet :: [Index] -> Index
unSet = (Index -> Index -> Index) -> Index -> [Index] -> Index
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Index
n Index
i -> Index
n Index -> Index -> Index
forall a. Bits a => a -> a -> a
.|. Int -> Index
forall a. Bits a => Int -> a
bit (Index -> Int
forall a. Num a => Index -> a
fromInteger Index
i)) Index
0

-- | An inverse of 'E.finiteEnumerationOf'.
--   
--   Given a coenumeration of @a@, make a coenumeration of function from
--   finite sets to @a@.
--   
--   Ideally, its type should be the following dependent type
--   
--   > {n :: Integer} -> CoEnumeration a -> CoEnumeration ({k :: Integer | k < n} -> a)
--
-- >>> let as = E.finiteEnumerationOf 3 (E.takeE 2 E.nat)
-- >>> map E.enumerate $ E.enumerate as
-- [[0,0,0],[0,0,1],[0,1,0],[0,1,1],[1,0,0],[1,0,1],[1,1,0],[1,1,1]]
-- >>> let inv_as = finiteFunctionOf 3 (takeC 2 nat)
-- >>> locate inv_as (E.select (E.finiteList [0,1,1]))
-- 3
finiteFunctionOf :: Integer -> CoEnumeration a -> CoEnumeration (Integer -> a)
finiteFunctionOf :: forall a. Index -> CoEnumeration a -> CoEnumeration (Index -> a)
finiteFunctionOf Index
0 CoEnumeration a
_ = CoEnumeration (Index -> a)
forall a. CoEnumeration a
unit
finiteFunctionOf Index
n CoEnumeration a
a = CoEnumeration{ card :: Cardinality
card = Cardinality
size, locate :: (Index -> a) -> Index
locate = (Index -> a) -> Index
locateEnum }
  where
    size :: Cardinality
size = case CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
a of
      Finite Index
k -> Index -> Cardinality
Finite (Index
kIndex -> Index -> Index
forall a b. (Num a, Integral b) => a -> b -> a
^Index
n)
      Cardinality
Infinite -> Cardinality
Infinite
    
    step :: Index -> Index -> Index
step = case CoEnumeration a -> Cardinality
forall a. CoEnumeration a -> Cardinality
card CoEnumeration a
a of
      Finite Index
k -> \Index
r Index
d -> Index
k Index -> Index -> Index
forall a. Num a => a -> a -> a
* Index
r Index -> Index -> Index
forall a. Num a => a -> a -> a
+ Index
d
      Cardinality
Infinite -> ((Index, Index) -> Index) -> Index -> Index -> Index
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (Index, Index) -> Index
undiagonal

    locateEnum :: (Index -> a) -> Index
locateEnum Index -> a
f =
      let go :: Index -> Index -> Index
go Index
i !Index
acc
            | Index
i Index -> Index -> Bool
forall a. Eq a => a -> a -> Bool
== Index
n    = Index
acc
            | Bool
otherwise = Index -> Index -> Index
go (Index
iIndex -> Index -> Index
forall a. Num a => a -> a -> a
+Index
1) (Index -> Index -> Index
step Index
acc (CoEnumeration a -> a -> Index
forall a. CoEnumeration a -> a -> Index
locate CoEnumeration a
a (Index -> a
f Index
i)))
      in Index -> Index -> Index
go Index
0 Index
0