packages feed

simple-enumeration-0.3: src/Data/CoEnumeration.hs

{-# 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.
    card :: Cardinality

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

-- | Returns if the the cardinality of coenumeration is finite.
isFinite :: CoEnumeration a -> Bool
isFinite = (Infinite /=) . 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 = CoEnumeration

instance Contravariant CoEnumeration where
  contramap f e = e{ locate = locate e . f }

-- | Associativity of 'divide' is maintained only when
--   arguments are finite.
instance Divisible CoEnumeration where
  divide f a b = contramap f $ a >< b
  conquer = unit

-- | Associativity of 'choose' is maintained only when
--   arguments are finite.
instance Decidable CoEnumeration where
  choose f a b = contramap f $ a <+> b
  lose f = contramap f 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 = CoEnumeration{ card = 1, locate = const 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{ card = 0, locate = const (error "locate void") }

-- | An inverse of forward 'E.boundedEnum'
boundedEnum :: forall a. (Enum a, Bounded a) => CoEnumeration a
boundedEnum = CoEnumeration{ card = size, locate = loc }
  where loc = toInteger . subtract lo . fromEnum
        lo = fromEnum (minBound @a)
        hi = fromEnum (maxBound @a)
        size = Finite $ 1 + toInteger hi - toInteger 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{ card = Infinite, locate = 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{ card = Infinite, locate = integerToNat }
  where
    integerToNat :: Integer -> Integer
    integerToNat n
      | n <= 0    = 2 * negate n
      | otherwise = 2 * n - 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{ card = Infinite, locate = locateCW }
  where
    locateCW x = case numerator x of
      n | n > 0     -> go n (denominator x) [] - 1
        | otherwise -> 33448095 {- Magic number, see the haddock above -}
    
    go 1 1 acc = foldl' (\i b -> 2 * i + b) 1 acc
    go a b acc
      | a > b = go (a - b) b (1 : acc)
      | a < b = go a (b - a) (0 : acc)
      | otherwise = error "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 = contramap caseBySign $ maybeOf (cw <+> cw)
  where
    caseBySign :: Rational -> Maybe (Either Rational Rational)
    caseBySign x = case compare x 0 of
      LT -> Just (Right (negate x))
      EQ -> Nothing
      GT -> Just (Left x)

-- | Sets the cardinality of given coenumeration to 'Infinite'
infinite :: CoEnumeration a -> CoEnumeration a
infinite e = 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)
e1 >< e2 = CoEnumeration{ card = n1 * n2, locate = locatePair }
  where
    n1 = card e1
    n2 = card e2
    locatePair = case (n1, n2) of
      (_,          Finite n2') -> \(a,b) -> locate e1 a * n2' + locate e2 b
      (Finite n1', Infinite)   -> \(a,b) -> locate e1 a + locate e2 b * n1'
      (Infinite,   Infinite)   -> \(a,b) -> undiagonal (locate e1 a, locate e2 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)
e1 <+> e2 = CoEnumeration{ card = n1 + n2, locate = locateEither }
  where
    n1 = card e1
    n2 = card e2
    locateEither = case (n1, n2) of
      (Finite n1', _)          -> either (locate e1) ((n1' +) . locate e2)
      (Infinite,   Finite n2') -> either ((n2' +) . locate e1) (locate e2)
      (Infinite,   Infinite)   -> either ((*2) . locate e1) (succ . (*2) . locate e2)

-- |
--
-- >>> locate (dropC 3 nat) <$> [0..5]
-- [0,0,0,0,1,2]
dropC :: Integer -> CoEnumeration a -> CoEnumeration a
dropC k e
  | k == 0      = e
  | card e == 0 = e
  | card e <= Finite k = error "Impossible empty coenumeration"
  | otherwise = CoEnumeration{ card = size, locate = loc }
  where
    size = card e - Finite k
    loc = max 0 . subtract k . locate e

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

checkEmpty :: CoEnumeration a -> CoEnumeration a
checkEmpty e
  | card e == 0 = e
  | otherwise   = error "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 n
  | n <= 0    = error $ "modulo: invalid argument " ++ show n
  | otherwise = CoEnumeration{ card = Finite n, locate = (`mod` 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 e1 e2 = CoEnumeration{
    card = max (card e1) (card e2)
  , locate = either (locate e1) (locate e2)
  }

-- | The inverse of forward 'E.maybeOf'
maybeOf :: CoEnumeration a -> CoEnumeration (Maybe a)
maybeOf e = contramap (maybe (Left ()) Right) $ unit <+> e

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

-- | 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 e = CoEnumeration{ card = size, locate = locateList }
  where
    n = card e
    size | n == 0    = 1
         | otherwise = Infinite
    locateList = unList n . fmap (locate e)

unList :: Cardinality -> [Index] -> Index
unList (Finite k) = foldl' (\n a -> 1 + (a + n * k)) 0 . reverse
unList Infinite   = foldl' (\n a -> 1 + undiagonal (a, n)) 0 . 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 e = CoEnumeration{ card = size, locate = unSet . fmap (locate e) }
  where
    size = case card e of
      Finite k -> Finite (2 ^ k)
      Infinite -> Infinite

unSet :: [Index] -> Index
unSet = foldl' (\n i -> n .|. bit (fromInteger i)) 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 0 _ = unit
finiteFunctionOf n a = CoEnumeration{ card = size, locate = locateEnum }
  where
    size = case card a of
      Finite k -> Finite (k^n)
      Infinite -> Infinite
    
    step = case card a of
      Finite k -> \r d -> k * r + d
      Infinite -> curry undiagonal

    locateEnum f =
      let go i !acc
            | i == n    = acc
            | otherwise = go (i+1) (step acc (locate a (f i)))
      in go 0 0