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
/= :: CompressedWord -> CompressedWord -> Bool
$c/= :: CompressedWord -> CompressedWord -> Bool
== :: CompressedWord -> CompressedWord -> Bool
$c== :: 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
min :: CompressedWord -> CompressedWord -> CompressedWord
$cmin :: CompressedWord -> CompressedWord -> CompressedWord
max :: CompressedWord -> CompressedWord -> CompressedWord
$cmax :: CompressedWord -> CompressedWord -> CompressedWord
>= :: CompressedWord -> CompressedWord -> Bool
$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
compare :: CompressedWord -> CompressedWord -> Ordering
$ccompare :: CompressedWord -> CompressedWord -> Ordering
$cp1Ord :: Eq 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
showList :: [CompressedWord] -> ShowS
$cshowList :: [CompressedWord] -> ShowS
show :: CompressedWord -> String
$cshow :: CompressedWord -> String
showsPrec :: Int -> CompressedWord -> ShowS
$cshowsPrec :: Int -> 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 (f :: * -> *) a. Applicative f => a -> f a
pure CompressedWord :: Word -> Word -> CompressedWord
CompressedWord { Word
limit :: Word
limit :: Word
limit, Word
value :: Word
value :: Word
value }

instance Json.ToJSON CompressedWord where
  toJSON :: CompressedWord -> Value
toJSON CompressedWord
x =
    [Pair] -> Value
Json.object [String -> Word -> Pair
forall value pair.
(ToJSON value, KeyValue pair) =>
String -> value -> pair
Json.pair String
"limit" (Word -> Pair) -> Word -> Pair
forall a b. (a -> b) -> a -> b
$ CompressedWord -> Word
limit CompressedWord
x, String -> Word -> Pair
forall value pair.
(ToJSON value, KeyValue pair) =>
String -> value -> pair
Json.pair String
"value" (Word -> Pair) -> Word -> Pair
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
$ [(Pair, Bool)] -> Value
Schema.object
  [ (String -> Value -> Pair
forall value pair.
(ToJSON value, KeyValue pair) =>
String -> value -> pair
Json.pair String
"limit" (Value -> Pair) -> Value -> Pair
forall a b. (a -> b) -> a -> b
$ [Pair] -> Value
Json.object [String -> String -> Pair
forall value pair.
(ToJSON value, KeyValue pair) =>
String -> value -> pair
Json.pair String
"type" String
"integer"], Bool
True)
  , (String -> Value -> Pair
forall value pair.
(ToJSON value, KeyValue pair) =>
String -> value -> pair
Json.pair String
"value" (Value -> Pair) -> Value -> Pair
forall a b. (a -> b) -> a -> b
$ [Pair] -> Value
Json.object [String -> String -> Pair
forall value pair.
(ToJSON value, KeyValue pair) =>
String -> value -> pair
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 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
limit = do
  Word
value <- Word -> Word -> Word -> Word -> BitGet Word
step Word
limit (Word -> Word
getMaxBits_ Word
limit) Word
0 Word
0
  CompressedWord -> BitGet CompressedWord
forall (f :: * -> *) a. Applicative f => a -> f a
pure CompressedWord :: Word -> Word -> CompressedWord
CompressedWord { Word
limit :: Word
limit :: Word
limit, Word
value :: Word
value :: Word
value }

getMaxBits_ :: Word -> Word
getMaxBits_ :: Word -> Word
getMaxBits_ Word
x = do
  let
    n :: Word
    n :: Word
n = Word -> Word -> Word
forall a. Ord a => a -> a -> a
max Word
1 (Double -> Word
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))))
  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 -> Word -> Word
forall a b. (Num a, Integral b) => a -> b -> a
^ Word
n then Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1 else Word
n

step :: Word -> Word -> Word -> Word -> BitGet.BitGet Word
step :: Word -> Word -> Word -> Word -> BitGet Word
step Word
limit_ Word
maxBits Word
position Word
value_ = do
  let x :: Word
x = Word -> Int -> Word
forall a. Bits a => a -> Int -> a
Bits.shiftL Word
1 (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
position) :: Word
  if Word
position Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
< Word
maxBits Bool -> Bool -> Bool
&& Word
value_ 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 do
      Bool
bit <- BitGet Bool
BitGet.bool
      let newValue :: Word
newValue = if Bool
bit then Word
value_ Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
x else Word
value_
      Word -> Word -> Word -> Word -> BitGet Word
step Word
limit_ Word
maxBits (Word
position Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1) Word
newValue
    else Word -> BitGet Word
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word
value_