-- |
-- Module      :  Codec.Scale.Compact
-- Copyright   :  Aleksandr Krupenkin 2016-2021
-- License     :  Apache-2.0
--
-- Maintainer  :  mail@akru.me
-- Stability   :  experimental
-- Portability :  noportable
--
-- Efficient general integer codec.
--

module Codec.Scale.Compact (Compact(..)) where

import           Control.Monad      (replicateM)
import           Data.Bits          (shiftL, shiftR, (.&.), (.|.))
import           Data.List          (unfoldr)
import           Data.Serialize.Get (getWord16le, getWord32le, getWord8,
                                     lookAhead)
import           Data.Serialize.Put (putWord16le, putWord32le, putWord8)

import           Codec.Scale.Class  (Decode (..), Encode (..))

-- | A "compact" or general integer encoding is sufficient for encoding
-- large integers (up to 2**536) and is more efficient at encoding most
-- values than the fixed-width version.
newtype Compact a = Compact { Compact a -> a
unCompact :: a }
  deriving (Compact a -> Compact a -> Bool
(Compact a -> Compact a -> Bool)
-> (Compact a -> Compact a -> Bool) -> Eq (Compact a)
forall a. Eq a => Compact a -> Compact a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Compact a -> Compact a -> Bool
$c/= :: forall a. Eq a => Compact a -> Compact a -> Bool
== :: Compact a -> Compact a -> Bool
$c== :: forall a. Eq a => Compact a -> Compact a -> Bool
Eq, Eq (Compact a)
Eq (Compact a)
-> (Compact a -> Compact a -> Ordering)
-> (Compact a -> Compact a -> Bool)
-> (Compact a -> Compact a -> Bool)
-> (Compact a -> Compact a -> Bool)
-> (Compact a -> Compact a -> Bool)
-> (Compact a -> Compact a -> Compact a)
-> (Compact a -> Compact a -> Compact a)
-> Ord (Compact a)
Compact a -> Compact a -> Bool
Compact a -> Compact a -> Ordering
Compact a -> Compact a -> Compact a
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
forall a. Ord a => Eq (Compact a)
forall a. Ord a => Compact a -> Compact a -> Bool
forall a. Ord a => Compact a -> Compact a -> Ordering
forall a. Ord a => Compact a -> Compact a -> Compact a
min :: Compact a -> Compact a -> Compact a
$cmin :: forall a. Ord a => Compact a -> Compact a -> Compact a
max :: Compact a -> Compact a -> Compact a
$cmax :: forall a. Ord a => Compact a -> Compact a -> Compact a
>= :: Compact a -> Compact a -> Bool
$c>= :: forall a. Ord a => Compact a -> Compact a -> Bool
> :: Compact a -> Compact a -> Bool
$c> :: forall a. Ord a => Compact a -> Compact a -> Bool
<= :: Compact a -> Compact a -> Bool
$c<= :: forall a. Ord a => Compact a -> Compact a -> Bool
< :: Compact a -> Compact a -> Bool
$c< :: forall a. Ord a => Compact a -> Compact a -> Bool
compare :: Compact a -> Compact a -> Ordering
$ccompare :: forall a. Ord a => Compact a -> Compact a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (Compact a)
Ord)

instance Show a => Show (Compact a) where
    show :: Compact a -> String
show = (String
"Compact " String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> (Compact a -> String) -> Compact a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> String
forall a. Show a => a -> String
show (a -> String) -> (Compact a -> a) -> Compact a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Compact a -> a
forall a. Compact a -> a
unCompact

instance Integral a => Encode (Compact a) where
    put :: Putter (Compact a)
put (Compact a
x)
      | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 = String -> Put
forall a. HasCallStack => String -> a
error String
"negatives not supported by compact codec"
      | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
64 = Put
singleByteMode
      | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
14 = Put
twoByteMode
      | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
30 = Put
fourByteMode
      | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
2Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
536 = Put
bigIntegerMode
      | Bool
otherwise = String -> Put
forall a. HasCallStack => String -> a
error (String -> Put) -> String -> Put
forall a b. (a -> b) -> a -> b
$ String
"unable to encode " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" as compact"
      where
        n :: Integer
n = a -> Integer
forall a. Integral a => a -> Integer
toInteger a
x
        singleByteMode :: Put
singleByteMode = Putter Word8
putWord8 (a -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
2)
        twoByteMode :: Put
twoByteMode = Putter Word16
putWord16le (a -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
`shiftL` Int
2 Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Word16
1)
        fourByteMode :: Put
fourByteMode = Putter Word32
putWord32le (a -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int
2 Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32
2)
        bigIntegerMode :: Put
bigIntegerMode = do
            let step :: b -> Maybe (a, b)
step b
0 = Maybe (a, b)
forall a. Maybe a
Nothing
                step b
i = (a, b) -> Maybe (a, b)
forall a. a -> Maybe a
Just (b -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral b
i, b
i b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftR` Int
8)
                unroll :: [Word8]
unroll = (Integer -> Maybe (Word8, Integer)) -> Integer -> [Word8]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr Integer -> Maybe (Word8, Integer)
forall b a. (Integral b, Num a, Bits b) => b -> Maybe (a, b)
step Integer
n
            Putter Word8
putWord8 (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Word8] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
unroll Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
4) Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
`shiftL` Int
2 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word8
3)
            Putter Word8 -> [Word8] -> Put
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Putter Word8
putWord8 [Word8]
unroll

instance Integral a => Decode (Compact a) where
    get :: Get (Compact a)
get = do
        Word8
mode <- Get Word8 -> Get Word8
forall a. Get a -> Get a
lookAhead ((Word8
3 Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&.) (Word8 -> Word8) -> Get Word8 -> Get Word8
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
getWord8)
        a -> Compact a
forall a. a -> Compact a
Compact (a -> Compact a) -> Get a -> Get (Compact a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Word8
mode of
          Word8
0 -> Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> a) -> Get Word8 -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
singleByteMode
          Word8
1 -> Word16 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16 -> a) -> Get Word16 -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word16
twoByteMode
          Word8
2 -> Word32 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> a) -> Get Word32 -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word32
fourByteMode
          Word8
3 -> Get a
bigIntegerMode
          Word8
_ -> String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"unexpected prefix decoding compact number"
      where
        singleByteMode :: Get Word8
singleByteMode = (Word8 -> Int -> Word8) -> Int -> Word8 -> Word8
forall a b c. (a -> b -> c) -> b -> a -> c
flip Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
shiftR Int
2 (Word8 -> Word8) -> Get Word8 -> Get Word8
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
getWord8
        twoByteMode :: Get Word16
twoByteMode = (Word16 -> Int -> Word16) -> Int -> Word16 -> Word16
forall a b c. (a -> b -> c) -> b -> a -> c
flip Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
shiftR Int
2 (Word16 -> Word16) -> Get Word16 -> Get Word16
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word16
getWord16le
        fourByteMode :: Get Word32
fourByteMode = (Word32 -> Int -> Word32) -> Int -> Word32 -> Word32
forall a b c. (a -> b -> c) -> b -> a -> c
flip Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
shiftR Int
2 (Word32 -> Word32) -> Get Word32 -> Get Word32
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word32
getWord32le
        bigIntegerMode :: Get a
bigIntegerMode = do
            let unstep :: a -> a -> a
unstep a
b a
a = a
a a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
8 a -> a -> a
forall a. Bits a => a -> a -> a
.|. a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
b
                roll :: [Word8] -> a
roll = Integer -> a
forall a. Num a => Integer -> a
fromInteger (Integer -> a) -> ([Word8] -> Integer) -> [Word8] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word8 -> Integer -> Integer) -> Integer -> [Word8] -> Integer
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Word8 -> Integer -> Integer
forall a a. (Bits a, Integral a, Num a) => a -> a -> a
unstep Integer
0
            Word8
len <- ((Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+Word8
4) (Word8 -> Word8) -> (Word8 -> Word8) -> Word8 -> Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word8 -> Int -> Word8) -> Int -> Word8 -> Word8
forall a b c. (a -> b -> c) -> b -> a -> c
flip Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
shiftR Int
2) (Word8 -> Word8) -> Get Word8 -> Get Word8
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
getWord8
            [Word8] -> a
roll ([Word8] -> a) -> Get [Word8] -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get Word8 -> Get [Word8]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
len) Get Word8
getWord8