-- | [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.
--
-- For codes that include the value @0@, see
-- [Elias.Natural](https://hackage-content.haskell.org/package/variety/docs/Codec-Elias-Natural.html).
module Codec.Elias
    ( -- * 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.
      --
      -- For example, while the binary expansion of @21@ is:
      --
      -- > import qualified Codec.Arithmetic.Variety.BitVec as BV
      -- > BV.toString $ BV.fromInteger 21
      -- > "10101"
      --
      -- its Elias code is:
      --
      -- > BV.toString $ enodeGamma 21
      -- > "000010101"
      --
      -- where an expansion of \(i\) is always preceeded by \(i-1\)
      -- 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.
    --
    -- For example:
    --
    -- > BV.toString $ BV.fromInteger (10^6)
    -- > "11110100001001000000"
    -- >
    -- > length "11110100001001000000"
    -- > 20
    --
    -- is prefixed with the gamma encoding of @20@ and loses its leading
    -- bit which begins every binary expansion:
    --
    -- > BV.toString <$> [encodeGamma 20, BV.fromInteger (10^6)]
    -- > ["000010100","11110100001001000000"]
    -- >
    -- > BV.toString $ encodeDelta 1000000
    -- > "0000101001110100001001000000"
    -- >
    -- > length "0000101001110100001001000000"
    -- > 28

    , 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.
    --
    -- For example:
    --
    -- > BV.toString . BV.fromInteger <$> [2,4,19,10^6]
    -- > ["10","100","10011","11110100001001000000"]
    -- >
    -- > length <$> ["10","100","10011","11110100001001000000"]
    -- > [2,3,5,20]
    -- >
    -- > BV.toString $ encodeOmega (10^6)
    -- > "1010010011111101000010010000000"
    -- >
    -- > length $ "1010010011111101000010010000000"
    -- > 31
    --
    -- Notice that, while /asymptotically/ more efficient, omega codes
    -- are longer than delta codes until around 1 googol, or @10^100@.

    , encodeOmega
    , decodeOmega
    ) where

import qualified Data.Bits as Bits
import Data.Bifunctor (Bifunctor(first))
import Codec.Arithmetic.Variety.BitVec (BitVec)
import qualified Codec.Arithmetic.Variety.BitVec as BV

boundsError :: a
boundsError :: forall a. a
boundsError = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Elias: Number must be positive and non-zero"

-- | Encode a number in a Elias gamma code. Throws an error if the input
-- is not positive and non-zero.
encodeGamma :: Integer -> BitVec
encodeGamma :: Integer -> BitVec
encodeGamma Integer
x | Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0 = Int -> Bool -> BitVec
BV.replicate Int
n Bool
False BitVec -> BitVec -> BitVec
forall a. Semigroup a => a -> a -> a
<> BitVec
xBits
              | Bool
otherwise = BitVec
forall a. a
boundsError
  where
    xBits :: BitVec
xBits = Integer -> BitVec
BV.fromInteger Integer
x
    n :: Int
n = BitVec -> Int
BV.length BitVec
xBits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

-- | 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 BitVec
bv | BitVec -> Int
BV.length BitVec
xBits Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
xLen = Maybe (Integer, BitVec)
forall a. Maybe a
Nothing
               | Bool
otherwise = (Integer, BitVec) -> Maybe (Integer, BitVec)
forall a. a -> Maybe a
Just (Integer
x, BitVec
bv'')
  where
    n :: Int
n = BitVec -> Int
BV.countLeadingZeros BitVec
bv
    xLen :: Int
xLen = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    bv' :: BitVec
bv' = Int -> Integer -> BitVec
BV.bitVec (BitVec -> Int
BV.length BitVec
bv Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) (Integer -> BitVec) -> Integer -> BitVec
forall a b. (a -> b) -> a -> b
$ BitVec -> Integer
BV.toInteger BitVec
bv -- truncate
    (BitVec
xBits, BitVec
bv'') = Int -> BitVec -> (BitVec, BitVec)
BV.splitAt Int
xLen BitVec
bv'
    x :: Integer
x = BitVec -> Integer
BV.toInteger BitVec
xBits

-- | Encode a number in a Elias delta code. Throws an error if the input
-- is not positive and non-zero.
encodeDelta :: Integer -> BitVec
encodeDelta :: Integer -> BitVec
encodeDelta Integer
x | Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0 = Integer -> BitVec
encodeGamma (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
xLen) BitVec -> BitVec -> BitVec
forall a. Semigroup a => a -> a -> a
<> BitVec
tailBits
              | Bool
otherwise = BitVec
forall a. a
boundsError
  where
    xBits :: BitVec
xBits = Integer -> BitVec
BV.fromInteger Integer
x
    xLen :: Int
xLen = BitVec -> Int
BV.length BitVec
xBits
    n :: Int
n = Int
xLen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    tailBits :: BitVec
tailBits = Int -> Integer -> BitVec
BV.bitVec Int
n (Integer -> BitVec) -> Integer -> BitVec
forall a b. (a -> b) -> a -> b
$ Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
Bits.clearBit Integer
x Int
n -- without leading bit

-- | 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 BitVec
bv = do
  (Int
xLen, BitVec
bv') <- (Integer -> Int) -> (Integer, BitVec) -> (Int, 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 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Integer, BitVec) -> (Int, BitVec))
-> Maybe (Integer, BitVec) -> Maybe (Int, BitVec)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BitVec -> Maybe (Integer, BitVec)
decodeGamma BitVec
bv
  let n :: Int
n = Int
xLen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
      (BitVec
xTail, BitVec
bv'') = Int -> BitVec -> (BitVec, BitVec)
BV.splitAt Int
n BitVec
bv'
      xBits :: BitVec
xBits = Bool -> BitVec
BV.singleton Bool
True BitVec -> BitVec -> BitVec
forall a. Semigroup a => a -> a -> a
<> BitVec
xTail
  if BitVec -> Int
BV.length BitVec
xBits Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
xLen then Maybe (Integer, BitVec)
forall a. Maybe a
Nothing
    else (Integer, BitVec) -> Maybe (Integer, BitVec)
forall a. a -> Maybe a
Just (BitVec -> Integer
BV.toInteger BitVec
xBits, BitVec
bv'')

-- | Encode a number in a Elias omega code. Throws an error if the input
-- is not positive and non-zero.
encodeOmega :: Integer -> BitVec
encodeOmega :: Integer -> BitVec
encodeOmega Integer
x0 | Integer
x0 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0 = Integer -> BitVec -> BitVec
go Integer
x0 BitVec
eom
               | Bool
otherwise = BitVec
forall a. a
boundsError
  where
    eom :: BitVec
eom = Int -> Integer -> BitVec
BV.bitVec Int
1 Integer
0 -- "0"

    go :: Integer -> BitVec -> BitVec
go Integer
1 = BitVec -> BitVec
forall a. a -> a
id -- end
    go Integer
n = Integer -> BitVec -> BitVec
go (Integer
lenInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) (BitVec -> BitVec) -> (BitVec -> BitVec) -> BitVec -> BitVec
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (BitVec
bv BitVec -> BitVec -> BitVec
forall a. Semigroup a => a -> a -> a
<>)
      where
        bv :: BitVec
bv = Integer -> BitVec
BV.fromInteger Integer
n
        len :: Integer
len = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ BitVec -> Int
BV.length BitVec
bv

-- | 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 -> Maybe (Integer, BitVec)
go Integer
1
  where
    go :: Integer -> BitVec -> Maybe (Integer, BitVec)
go Integer
n BitVec
bv = do
      Bool
b <- BitVec
bv BitVec -> Int -> Maybe Bool
BV.!? Int
0 -- head
      case Bool
b of
        Bool
True | BitVec -> Int
BV.length BitVec
valBits Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
len -> Maybe (Integer, BitVec)
forall a. Maybe a
Nothing
             | Bool
otherwise -> Integer -> BitVec -> Maybe (Integer, BitVec)
go Integer
n' BitVec
bv'
          where
            len :: Int
len = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            (BitVec
valBits, BitVec
bv') = Int -> BitVec -> (BitVec, BitVec)
BV.splitAt Int
len BitVec
bv
            n' :: Integer
n' = BitVec -> Integer
BV.toInteger BitVec
valBits

        Bool
False -> (Integer, BitVec) -> Maybe (Integer, BitVec)
forall a. a -> Maybe a
Just (Integer
n, Int -> BitVec -> BitVec
BV.drop Int
1 BitVec
bv) -- eom