module Data.Falsify.Internal.Integer (
    -- * Encoding
    Bit(..)
  , encIntegerEliasG
  ) where

import Data.Bits
import Numeric.Natural

{-------------------------------------------------------------------------------
  Binary encoding
-------------------------------------------------------------------------------}

data Bit = I | O
  deriving (Int -> Bit -> ShowS
[Bit] -> ShowS
Bit -> String
(Int -> Bit -> ShowS)
-> (Bit -> String) -> ([Bit] -> ShowS) -> Show Bit
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Bit -> ShowS
showsPrec :: Int -> Bit -> ShowS
$cshow :: Bit -> String
show :: Bit -> String
$cshowList :: [Bit] -> ShowS
showList :: [Bit] -> ShowS
Show, Bit -> Bit -> Bool
(Bit -> Bit -> Bool) -> (Bit -> Bit -> Bool) -> Eq Bit
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Bit -> Bit -> Bool
== :: Bit -> Bit -> Bool
$c/= :: Bit -> Bit -> Bool
/= :: Bit -> Bit -> Bool
Eq, Eq Bit
Eq Bit =>
(Bit -> Bit -> Ordering)
-> (Bit -> Bit -> Bool)
-> (Bit -> Bit -> Bool)
-> (Bit -> Bit -> Bool)
-> (Bit -> Bit -> Bool)
-> (Bit -> Bit -> Bit)
-> (Bit -> Bit -> Bit)
-> Ord Bit
Bit -> Bit -> Bool
Bit -> Bit -> Ordering
Bit -> Bit -> Bit
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Bit -> Bit -> Ordering
compare :: Bit -> Bit -> Ordering
$c< :: Bit -> Bit -> Bool
< :: Bit -> Bit -> Bool
$c<= :: Bit -> Bit -> Bool
<= :: Bit -> Bit -> Bool
$c> :: Bit -> Bit -> Bool
> :: Bit -> Bit -> Bool
$c>= :: Bit -> Bit -> Bool
>= :: Bit -> Bit -> Bool
$cmax :: Bit -> Bit -> Bit
max :: Bit -> Bit -> Bit
$cmin :: Bit -> Bit -> Bit
min :: Bit -> Bit -> Bit
Ord)

-- | Binary encoding (most significant bit first)
natToBits :: Natural -> [Bit]
natToBits :: Natural -> [Bit]
natToBits = \Natural
n -> if
  | Natural
n Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
< Natural
0     -> String -> [Bit]
forall a. HasCallStack => String -> a
error String
"toBits: negative input"
  | Natural
n Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
0    -> []
  | Bool
otherwise -> [Bit] -> [Bit]
forall a. [a] -> [a]
reverse ([Bit] -> [Bit]) -> [Bit] -> [Bit]
forall a b. (a -> b) -> a -> b
$ Natural -> [Bit]
go Natural
n
  where
    go :: Natural -> [Bit]
    go :: Natural -> [Bit]
go Natural
0 = []
    go Natural
n = (if Natural -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Natural
n Int
0 then Bit
I else Bit
O) Bit -> [Bit] -> [Bit]
forall a. a -> [a] -> [a]
: Natural -> [Bit]
go (Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
shiftR Natural
n Int
1)

{-------------------------------------------------------------------------------
  Elias γ code
-------------------------------------------------------------------------------}

-- | Elias γ code
--
-- Precondition: input @x >= 1@.
--
-- See <https://en.wikipedia.org/wiki/Elias_gamma_coding> .
encEliasG :: Natural -> [Bit]
encEliasG :: Natural -> [Bit]
encEliasG Natural
x
  | Natural
x Natural -> Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Natural
0    = String -> [Bit]
forall a. HasCallStack => String -> a
error String
"eliasG: zero"
  | Bool
otherwise = Natural -> [Bit]
zeroes Natural
x
  where
    zeroes :: Natural -> [Bit]
    zeroes :: Natural -> [Bit]
zeroes Natural
n
      | Natural
n Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
<= Natural
1    = Natural -> [Bit]
natToBits Natural
x
      | Bool
otherwise = Bit
O Bit -> [Bit] -> [Bit]
forall a. a -> [a] -> [a]
: Natural -> [Bit]
zeroes (Natural -> Int -> Natural
forall a. Bits a => a -> Int -> a
shiftR Natural
n Int
1)

-- | Extension of Elias γ coding to signed integers
--
-- This is adapted from @integerVariant@ in @Test.QuickCheck.Random@. The first
-- bit encs whether @x >= 1@ or not (this will result in @0@ and @1@ having
-- short codes).
encIntegerEliasG :: Integer -> [Bit]
encIntegerEliasG :: Integer -> [Bit]
encIntegerEliasG = \Integer
x ->
    if Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
1
      then Bit
O Bit -> [Bit] -> [Bit]
forall a. a -> [a] -> [a]
: Natural -> [Bit]
encEliasG (Integer -> Natural
forall a. Num a => Integer -> a
fromInteger          (Integer -> Natural) -> Integer -> Natural
forall a b. (a -> b) -> a -> b
$ Integer
x)
      else Bit
I Bit -> [Bit] -> [Bit]
forall a. a -> [a] -> [a]
: Natural -> [Bit]
encEliasG (Integer -> Natural
forall a. Num a => Integer -> a
fromInteger (Integer -> Natural) -> (Integer -> Integer) -> Integer -> Natural
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Integer
mangle (Integer -> Natural) -> Integer -> Natural
forall a b. (a -> b) -> a -> b
$ Integer
x)
  where
    mangle :: Integer -> Integer
    mangle :: Integer -> Integer
mangle Integer
x = Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
x