-- | [Elias codes](https://en.wikipedia.org/wiki/Elias_coding) are
-- prefix codes for positive, non-zero integers with no assumption or
-- limit to their size.
--
-- Functions of this module add @1@ at encoding time and subtract @1@ at
-- decoding time to support any natural number, including zero.
module Codec.Elias.Natural
    ( -- * Gamma coding

      -- | An Elias gamma code consists of the binary expansion of an
      -- integer, preceded by the unary encoding of the length of that
      -- expansion in zeros.

      encodeGamma
    , decodeGamma

    -- * Delta coding

    -- | An Elias delta code is like an Elias gamma code except that the
    -- length is itself coded like a gamma code instead of simply a
    -- unary encoding.

    , encodeDelta
    , decodeDelta

    -- * Omega coding

    -- | An Elias omega code is the result of recursively encoding the
    -- length of binary expansions in the prefix until a length of @1@
    -- is reached. Since binary expansions are written without any
    -- leading zeros, a single @0@ bit marks the end of the code.

    , encodeOmega
    , decodeOmega
    ) where

import Data.Bifunctor (Bifunctor(first))
import Codec.Arithmetic.Variety.BitVec (BitVec)
import qualified Codec.Elias as E

boundsError :: a
boundsError :: forall a. a
boundsError = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Elias.Natural: negative number"

-- | Encode a number in a Elias gamma code. Throws an error if the input
-- is negative.
encodeGamma :: Integer -> BitVec
encodeGamma :: Integer -> BitVec
encodeGamma Integer
n | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 = Integer -> BitVec
E.encodeGamma (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1)
              | Bool
otherwise = BitVec
forall a. a
boundsError

-- | Try to decode an Elias gamma code at the head of the given bit
-- vector. If successful, returns the decoded value and the remainder of
-- the `BitVec`, with the value code removed. Returns @Nothing@ if the
-- bit vector doesn't contain enough bits to define a number.
decodeGamma :: BitVec -> Maybe (Integer, BitVec)
decodeGamma :: BitVec -> Maybe (Integer, BitVec)
decodeGamma = ((Integer, BitVec) -> (Integer, BitVec))
-> Maybe (Integer, BitVec) -> Maybe (Integer, BitVec)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((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 -> Integer -> Integer
forall a. Num a => a -> a -> a
+(-Integer
1))) (Maybe (Integer, BitVec) -> Maybe (Integer, BitVec))
-> (BitVec -> Maybe (Integer, BitVec))
-> BitVec
-> Maybe (Integer, BitVec)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitVec -> Maybe (Integer, BitVec)
E.decodeGamma

-- | Encode a number in a Elias delta code. Throws an error if the input
-- is negative.
encodeDelta :: Integer -> BitVec
encodeDelta :: Integer -> BitVec
encodeDelta Integer
n | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 = Integer -> BitVec
E.encodeDelta (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1)
              | Bool
otherwise = BitVec
forall a. a
boundsError

-- | Try to decode an Elias delta code at the head of the given bit
-- vector. If successful, returns the decoded value and the remainder of
-- the `BitVec`, with the value code removed. Returns @Nothing@ if the
-- bit vector doesn't contain enough bits to define a number.
decodeDelta :: BitVec -> Maybe (Integer, BitVec)
decodeDelta :: BitVec -> Maybe (Integer, BitVec)
decodeDelta = ((Integer, BitVec) -> (Integer, BitVec))
-> Maybe (Integer, BitVec) -> Maybe (Integer, BitVec)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((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 -> Integer -> Integer
forall a. Num a => a -> a -> a
+(-Integer
1))) (Maybe (Integer, BitVec) -> Maybe (Integer, BitVec))
-> (BitVec -> Maybe (Integer, BitVec))
-> BitVec
-> Maybe (Integer, BitVec)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitVec -> Maybe (Integer, BitVec)
E.decodeDelta

-- | Encode a number in a Elias omega code. Throws an error if the input
-- is negative.
encodeOmega :: Integer -> BitVec
encodeOmega :: Integer -> BitVec
encodeOmega Integer
n | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 = Integer -> BitVec
E.encodeOmega Integer
n
              | Bool
otherwise = BitVec
forall a. a
boundsError

-- | Try to decode an Elias omega code at the head of the given bit
-- vector. If successful, returns the decoded value and the remainder of
-- the `BitVec`, with the value code removed. Returns @Nothing@ if the
-- bit vector doesn't contain enough bits to define a number.
decodeOmega :: BitVec -> Maybe (Integer, BitVec)
decodeOmega :: BitVec -> Maybe (Integer, BitVec)
decodeOmega = ((Integer, BitVec) -> (Integer, BitVec))
-> Maybe (Integer, BitVec) -> Maybe (Integer, BitVec)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((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 -> Integer -> Integer
forall a. Num a => a -> a -> a
+(-Integer
1))) (Maybe (Integer, BitVec) -> Maybe (Integer, BitVec))
-> (BitVec -> Maybe (Integer, BitVec))
-> BitVec
-> Maybe (Integer, BitVec)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BitVec -> Maybe (Integer, BitVec)
E.decodeOmega