module Rattletrap.Type.CompressedWord where

import qualified Data.Bits as Bits
import qualified Rattletrap.BitGet as BitGet
import qualified Rattletrap.BitPut as BitPut
import qualified Rattletrap.Schema as Schema
import qualified Rattletrap.Utility.Json as Json

-- | Although there's no guarantee that these values will not overflow, it's
-- exceptionally unlikely. Most 'CompressedWord's are very small.
data CompressedWord = CompressedWord
  { CompressedWord -> Word
limit :: Word,
    CompressedWord -> Word
value :: Word
  }
  deriving (CompressedWord -> CompressedWord -> Bool
(CompressedWord -> CompressedWord -> Bool)
-> (CompressedWord -> CompressedWord -> Bool) -> Eq CompressedWord
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CompressedWord -> CompressedWord -> Bool
== :: CompressedWord -> CompressedWord -> Bool
$c/= :: CompressedWord -> CompressedWord -> Bool
/= :: CompressedWord -> CompressedWord -> Bool
Eq, Eq CompressedWord
Eq CompressedWord =>
(CompressedWord -> CompressedWord -> Ordering)
-> (CompressedWord -> CompressedWord -> Bool)
-> (CompressedWord -> CompressedWord -> Bool)
-> (CompressedWord -> CompressedWord -> Bool)
-> (CompressedWord -> CompressedWord -> Bool)
-> (CompressedWord -> CompressedWord -> CompressedWord)
-> (CompressedWord -> CompressedWord -> CompressedWord)
-> Ord CompressedWord
CompressedWord -> CompressedWord -> Bool
CompressedWord -> CompressedWord -> Ordering
CompressedWord -> CompressedWord -> CompressedWord
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 :: CompressedWord -> CompressedWord -> Ordering
compare :: CompressedWord -> CompressedWord -> Ordering
$c< :: CompressedWord -> CompressedWord -> Bool
< :: CompressedWord -> CompressedWord -> Bool
$c<= :: CompressedWord -> CompressedWord -> Bool
<= :: CompressedWord -> CompressedWord -> Bool
$c> :: CompressedWord -> CompressedWord -> Bool
> :: CompressedWord -> CompressedWord -> Bool
$c>= :: CompressedWord -> CompressedWord -> Bool
>= :: CompressedWord -> CompressedWord -> Bool
$cmax :: CompressedWord -> CompressedWord -> CompressedWord
max :: CompressedWord -> CompressedWord -> CompressedWord
$cmin :: CompressedWord -> CompressedWord -> CompressedWord
min :: CompressedWord -> CompressedWord -> CompressedWord
Ord, Int -> CompressedWord -> ShowS
[CompressedWord] -> ShowS
CompressedWord -> String
(Int -> CompressedWord -> ShowS)
-> (CompressedWord -> String)
-> ([CompressedWord] -> ShowS)
-> Show CompressedWord
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CompressedWord -> ShowS
showsPrec :: Int -> CompressedWord -> ShowS
$cshow :: CompressedWord -> String
show :: CompressedWord -> String
$cshowList :: [CompressedWord] -> ShowS
showList :: [CompressedWord] -> ShowS
Show)

instance Json.FromJSON CompressedWord where
  parseJSON :: Value -> Parser CompressedWord
parseJSON = String
-> (Object -> Parser CompressedWord)
-> Value
-> Parser CompressedWord
forall a. String -> (Object -> Parser a) -> Value -> Parser a
Json.withObject String
"CompressedWord" ((Object -> Parser CompressedWord)
 -> Value -> Parser CompressedWord)
-> (Object -> Parser CompressedWord)
-> Value
-> Parser CompressedWord
forall a b. (a -> b) -> a -> b
$ \Object
object -> do
    Word
limit <- Object -> String -> Parser Word
forall value. FromJSON value => Object -> String -> Parser value
Json.required Object
object String
"limit"
    Word
value <- Object -> String -> Parser Word
forall value. FromJSON value => Object -> String -> Parser value
Json.required Object
object String
"value"
    CompressedWord -> Parser CompressedWord
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CompressedWord {Word
limit :: Word
limit :: Word
limit, Word
value :: Word
value :: Word
value}

instance Json.ToJSON CompressedWord where
  toJSON :: CompressedWord -> Value
toJSON CompressedWord
x =
    [(Key, Value)] -> Value
Json.object [String -> Word -> (Key, Value)
forall value e p.
(ToJSON value, KeyValue e p) =>
String -> value -> p
Json.pair String
"limit" (Word -> (Key, Value)) -> Word -> (Key, Value)
forall a b. (a -> b) -> a -> b
$ CompressedWord -> Word
limit CompressedWord
x, String -> Word -> (Key, Value)
forall value e p.
(ToJSON value, KeyValue e p) =>
String -> value -> p
Json.pair String
"value" (Word -> (Key, Value)) -> Word -> (Key, Value)
forall a b. (a -> b) -> a -> b
$ CompressedWord -> Word
value CompressedWord
x]

schema :: Schema.Schema
schema :: Schema
schema =
  String -> Value -> Schema
Schema.named String
"compressedWord" (Value -> Schema) -> Value -> Schema
forall a b. (a -> b) -> a -> b
$
    [((Key, Value), Bool)] -> Value
Schema.object
      [ (String -> Value -> (Key, Value)
forall value e p.
(ToJSON value, KeyValue e p) =>
String -> value -> p
Json.pair String
"limit" (Value -> (Key, Value)) -> Value -> (Key, Value)
forall a b. (a -> b) -> a -> b
$ [(Key, Value)] -> Value
Json.object [String -> String -> (Key, Value)
forall value e p.
(ToJSON value, KeyValue e p) =>
String -> value -> p
Json.pair String
"type" String
"integer"], Bool
True),
        (String -> Value -> (Key, Value)
forall value e p.
(ToJSON value, KeyValue e p) =>
String -> value -> p
Json.pair String
"value" (Value -> (Key, Value)) -> Value -> (Key, Value)
forall a b. (a -> b) -> a -> b
$ [(Key, Value)] -> Value
Json.object [String -> String -> (Key, Value)
forall value e p.
(ToJSON value, KeyValue e p) =>
String -> value -> p
Json.pair String
"type" String
"integer"], Bool
True)
      ]

bitPut :: CompressedWord -> BitPut.BitPut
bitPut :: CompressedWord -> BitPut
bitPut CompressedWord
compressedWord =
  let limit_ :: Word
limit_ = CompressedWord -> Word
limit CompressedWord
compressedWord
      value_ :: Word
value_ = CompressedWord -> Word
value CompressedWord
compressedWord
      maxBits :: Int
maxBits = Word -> Int
getMaxBits Word
limit_
   in Word -> Word -> Int -> Int -> Word -> BitPut
putCompressedWordStep Word
limit_ Word
value_ Int
maxBits Int
0 Word
0

putCompressedWordStep :: Word -> Word -> Int -> Int -> Word -> BitPut.BitPut
putCompressedWordStep :: Word -> Word -> Int -> Int -> Word -> BitPut
putCompressedWordStep Word
limit_ Word
value_ Int
maxBits Int
position Word
soFar =
  if Int
position Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxBits
    then do
      let x :: Word
x = Word -> Int -> Word
forall a. Bits a => a -> Int -> a
Bits.shiftL Word
1 Int
position :: Word
      if Int
maxBits Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 Bool -> Bool -> Bool
&& Int
position Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maxBits Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Bool -> Bool -> Bool
&& Word
soFar Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
x Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
limit_
        then BitPut
forall a. Monoid a => a
mempty
        else
          let bit :: Bool
bit = Word -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
Bits.testBit Word
value_ Int
position
              delta :: Word
delta = if Bool
bit then Word
x else Word
0
           in Bool -> BitPut
BitPut.bool Bool
bit
                BitPut -> BitPut -> BitPut
forall a. Semigroup a => a -> a -> a
<> Word -> Word -> Int -> Int -> Word -> BitPut
putCompressedWordStep
                  Word
limit_
                  Word
value_
                  Int
maxBits
                  (Int
position Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                  (Word
soFar Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
delta)
    else BitPut
forall a. Monoid a => a
mempty

getMaxBits :: Word -> Int
getMaxBits :: Word -> Int
getMaxBits Word
x =
  let n :: Int
      n :: Int
n = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase (Double
2 :: Double) (Word -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word -> Word -> Word
forall a. Ord a => a -> a -> a
max Word
1 Word
x))))
   in if Word
x Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
1024 Bool -> Bool -> Bool
&& Word
x Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
2 Word -> Int -> Word
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
n then Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 else Int
n

bitGet :: Word -> BitGet.BitGet CompressedWord
bitGet :: Word -> BitGet CompressedWord
bitGet = Word -> BitGet CompressedWord
bitGetNew

bitGetNew :: Word -> BitGet.BitGet CompressedWord
bitGetNew :: Word -> BitGet CompressedWord
bitGetNew Word
limit = do
  Word
value <-
    if Word
limit Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
1
      then Word -> Get BitString Identity Word
forall a. a -> Get BitString Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word
0
      else do
        let numBits :: Int
numBits =
              Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
0 :: Int)
                (Int -> Int) -> (Double -> Int) -> Double -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1
                (Int -> Int) -> (Double -> Int) -> Double -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling
                (Double -> Int) -> (Double -> Double) -> Double -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase (Double
2 :: Double)
                (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Word -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
limit
        Word
partial <- Int -> Get BitString Identity Word
forall a. Bits a => Int -> BitGet a
BitGet.bits Int
numBits
        let next :: Word
next = Word
partial Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word -> Int -> Word
forall a. Bits a => a -> Int -> a
Bits.shiftL Word
1 Int
numBits
        if Word
next Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
limit
          then Word -> Get BitString Identity Word
forall a. a -> Get BitString Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word
partial
          else do
            Bool
x <- BitGet Bool
BitGet.bool
            Word -> Get BitString Identity Word
forall a. a -> Get BitString Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word -> Get BitString Identity Word)
-> Word -> Get BitString Identity Word
forall a b. (a -> b) -> a -> b
$ if Bool
x then Word
next else Word
partial
  CompressedWord -> BitGet CompressedWord
forall a. a -> Get BitString Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CompressedWord {Word
limit :: Word
limit :: Word
limit, Word
value :: Word
value :: Word
value}