-- | Since the arithmetic operations of composition might get
-- computationally expensive on very large codes, a similar interface is
-- provided here which produces and consumes bits in chunks whenever
-- spaces are about to grow beyond a certain size given in bytes, at the
-- cost of at most one bit per chunk.
--
-- While the Haskell language standard defines `Integer` as having no
-- upper bound, GHC most commonly uses the GNU Multiple Precision
-- Arithmetic Library (GMP) as a backend for it, which incurs a limit of
-- 16GiB (or a little over 17GB) on the size of `Integer` values.
module Codec.Arithmetic.Variety.Bounded
  ( encode
  , codeLen
  , decode
  ) where

import Data.Bits (Bits(bit))
import Data.Bifunctor (Bifunctor(first))

import qualified Codec.Arithmetic.Variety as Var
import Codec.Arithmetic.Variety.BitVec (BitVec)

err :: String -> a
err :: forall a. String -> a
err = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> (String -> String) -> String -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"Variety.Bounded: " String -> String -> String
forall a. [a] -> [a] -> [a]
++)

-- | Given a base getter function and precision, chunk the values into
-- lists that do not exceed the given precision or are singleton lists,
-- annotated with the total base for each chunk.
groupWithinPrec :: (a -> Integer) -> Int -> [a] -> [(Integer,[a])]
groupWithinPrec :: forall a. (a -> Integer) -> Int -> [a] -> [(Integer, [a])]
groupWithinPrec a -> Integer
getBase Int
prec
  | Int
prec Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> [a] -> [(Integer, [a])]
forall a. String -> a
err String
"negative precision"
  | Bool
otherwise = ([a] -> [a]) -> [(Integer, [a])] -> [(Integer, [a])]
forall (f :: * -> *) (g :: * -> *) a b.
(Functor f, Functor g) =>
(a -> b) -> f (g a) -> f (g b)
ffmap [a] -> [a]
forall a. [a] -> [a]
reverse ([(Integer, [a])] -> [(Integer, [a])])
-> ([a] -> [(Integer, [a])]) -> [a] -> [(Integer, [a])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> [a] -> [a] -> [(Integer, [a])]
go Integer
1 []
  where
    maxBase :: Integer
maxBase = Int -> Integer
forall a. Bits a => Int -> a
bit (Int
precInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1 -- max val with `prec` bytes
    go :: Integer -> [a] -> [a] -> [(Integer, [a])]
go Integer
base [a]
group [] = ((Integer, [a]) -> Bool) -> [(Integer, [a])] -> [(Integer, [a])]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool)
-> ((Integer, [a]) -> Bool) -> (Integer, [a]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([a] -> Bool) -> ((Integer, [a]) -> [a]) -> (Integer, [a]) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, [a]) -> [a]
forall a b. (a, b) -> b
snd) [(Integer
base,[a]
group)]
    go Integer
1 [a]
group (a
a:[a]
as) = Integer -> [a] -> [a] -> [(Integer, [a])]
go (a -> Integer
getBase a
a) (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
group) [a]
as
    go Integer
base [a]
group (a
a:[a]
as)
      | Integer
base' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
maxBase = (Integer
base,[a]
group) (Integer, [a]) -> [(Integer, [a])] -> [(Integer, [a])]
forall a. a -> [a] -> [a]
: Integer -> [a] -> [a] -> [(Integer, [a])]
go Integer
b [a
a] [a]
as
      | Bool
otherwise = Integer -> [a] -> [a] -> [(Integer, [a])]
go Integer
base' (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
group) [a]
as
      where
        b :: Integer
b = a -> Integer
getBase a
a
        base' :: Integer
base' = Integer
base Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b
{-# INLINE groupWithinPrec #-}

-- | Given a max precision in bytes, encode a series of value-base pairs
-- into a single bit vector. Bases must be at least equal to @1@ and the
-- associated values must exist in the range @[0..base-1]@.
encode :: Int -> [(Integer,Integer)] -> BitVec
encode :: Int -> [(Integer, Integer)] -> BitVec
encode = [BitVec] -> BitVec
forall a. Monoid a => [a] -> a
mconcat
         ([BitVec] -> BitVec)
-> ([(Integer, [(Integer, Integer)])] -> [BitVec])
-> [(Integer, [(Integer, Integer)])]
-> BitVec
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Integer, [(Integer, Integer)]) -> BitVec)
-> [(Integer, [(Integer, Integer)])] -> [BitVec]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([(Integer, Integer)] -> BitVec
Var.encode ([(Integer, Integer)] -> BitVec)
-> ((Integer, [(Integer, Integer)]) -> [(Integer, Integer)])
-> (Integer, [(Integer, Integer)])
-> BitVec
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, [(Integer, Integer)]) -> [(Integer, Integer)]
forall a b. (a, b) -> b
snd)
         ([(Integer, [(Integer, Integer)])] -> BitVec)
-> (Int
    -> [(Integer, Integer)] -> [(Integer, [(Integer, Integer)])])
-> Int
-> [(Integer, Integer)]
-> BitVec
forall b c a1 a2. (b -> c) -> (a1 -> a2 -> b) -> a1 -> a2 -> c
.: ((Integer, Integer) -> Integer)
-> Int -> [(Integer, Integer)] -> [(Integer, [(Integer, Integer)])]
forall a. (a -> Integer) -> Int -> [a] -> [(Integer, [a])]
groupWithinPrec (Integer, Integer) -> Integer
forall a b. (a, b) -> b
snd

-- | Return the length of the code of a sequence of values in the given
-- precision and list of bases in bits.
codeLen :: Int -> [Integer] -> Int
codeLen :: Int -> [Integer] -> Int
codeLen = Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral
          (Int -> Int)
-> ([(Integer, [Integer])] -> Int) -> [(Integer, [Integer])] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
          ([Int] -> Int)
-> ([(Integer, [Integer])] -> [Int])
-> [(Integer, [Integer])]
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Integer, [Integer]) -> Int) -> [(Integer, [Integer])] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Integer -> Int
Var.codeLen1 (Integer -> Int)
-> ((Integer, [Integer]) -> Integer) -> (Integer, [Integer]) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, [Integer]) -> Integer
forall a b. (a, b) -> a
fst)
          ([(Integer, [Integer])] -> Int)
-> (Int -> [Integer] -> [(Integer, [Integer])])
-> Int
-> [Integer]
-> Int
forall b c a1 a2. (b -> c) -> (a1 -> a2 -> b) -> a1 -> a2 -> c
.: (Integer -> Integer) -> Int -> [Integer] -> [(Integer, [Integer])]
forall a. (a -> Integer) -> Int -> [a] -> [(Integer, [a])]
groupWithinPrec Integer -> Integer
forall a. a -> a
id

-- | Try to decode the head of a bit vector given the same precision and
-- series of bases that was used to encode. If successful, returns the
-- decoded values and the remainder of the `BitVec` with the values'
-- code removed. Returns @Nothing@ if the bit vector doesn't contain
-- enough bits to specify values for each given base.
decode :: Int -> [Integer] -> BitVec -> Maybe ([Integer], BitVec)
decode :: Int -> [Integer] -> BitVec -> Maybe ([Integer], BitVec)
decode = [(Integer, [Integer])] -> BitVec -> Maybe ([Integer], BitVec)
forall {a}. [(a, [Integer])] -> BitVec -> Maybe ([Integer], BitVec)
go ([(Integer, [Integer])] -> BitVec -> Maybe ([Integer], BitVec))
-> (Int -> [Integer] -> [(Integer, [Integer])])
-> Int
-> [Integer]
-> BitVec
-> Maybe ([Integer], BitVec)
forall b c a1 a2. (b -> c) -> (a1 -> a2 -> b) -> a1 -> a2 -> c
.: (Integer -> Integer) -> Int -> [Integer] -> [(Integer, [Integer])]
forall a. (a -> Integer) -> Int -> [a] -> [(Integer, [a])]
groupWithinPrec Integer -> Integer
forall a. a -> a
id
  where
    go :: [(a, [Integer])] -> BitVec -> Maybe ([Integer], BitVec)
go [] BitVec
bv = ([Integer], BitVec) -> Maybe ([Integer], BitVec)
forall a. a -> Maybe a
Just ([], BitVec
bv)
    go ((a
_,[Integer]
bases):[(a, [Integer])]
rest) BitVec
bv = do
      ([Integer]
vals, BitVec
bv') <- [Integer] -> BitVec -> Maybe ([Integer], BitVec)
Var.decode [Integer]
bases BitVec
bv
      ([Integer] -> [Integer])
-> ([Integer], BitVec) -> ([Integer], BitVec)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ([Integer]
vals [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++) (([Integer], BitVec) -> ([Integer], BitVec))
-> Maybe ([Integer], BitVec) -> Maybe ([Integer], BitVec)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, [Integer])] -> BitVec -> Maybe ([Integer], BitVec)
go [(a, [Integer])]
rest BitVec
bv'

(.:) :: (b -> c) -> (a1 -> a2 -> b) -> a1 -> a2 -> c
.: :: forall b c a1 a2. (b -> c) -> (a1 -> a2 -> b) -> a1 -> a2 -> c
(.:) = ((a2 -> b) -> a2 -> c) -> (a1 -> a2 -> b) -> a1 -> a2 -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.)(((a2 -> b) -> a2 -> c) -> (a1 -> a2 -> b) -> a1 -> a2 -> c)
-> ((b -> c) -> (a2 -> b) -> a2 -> c)
-> (b -> c)
-> (a1 -> a2 -> b)
-> a1
-> a2
-> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(b -> c) -> (a2 -> b) -> a2 -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.)
infixr 8 .:
{-# INLINE (.:) #-}

ffmap :: (Functor f, Functor g) => (a -> b) -> f (g a) -> f (g b)
ffmap :: forall (f :: * -> *) (g :: * -> *) a b.
(Functor f, Functor g) =>
(a -> b) -> f (g a) -> f (g b)
ffmap = (g a -> g b) -> f (g a) -> f (g b)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((g a -> g b) -> f (g a) -> f (g b))
-> ((a -> b) -> g a -> g b) -> (a -> b) -> f (g a) -> f (g b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> g a -> g b
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
{-# INLINE ffmap #-}