Internal.hs

-- SPDX-FileCopyrightText: 2021 Tocqueville Group
-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA
-- SPDX-License-Identifier: LicenseRef-MIT-obsidian-systems

-- | Defines binary encoding primitives which mimic Tezos' binary encoding.
module Morley.Micheline.Binary.Internal
  ( DynamicSize(..)

  -- * Encode
  , buildWord8
  , buildText
  , buildInteger
  , buildNatural
  , buildDynamic
  , buildByteString

  -- * Decode
  , getText
  , getNatural
  , getInteger
  , getDynamic
  , getByteString
  ) where

import Control.Exception (assert)
import Data.Binary.Builder qualified as Bi
import Data.Binary.Get qualified as Bi
import Data.Bits (Bits, bit, setBit, shift, testBit, zeroBits, (.&.), (.|.))
import Data.ByteString.Lazy qualified as LBS
import Data.Text.Encoding qualified as TE
import Unsafe qualified (fromIntegral)

newtype DynamicSize a = DynamicSize { unDynamicSize :: a }

-------------------------------------------------
-- Encode
-------------------------------------------------

buildWord8 :: Word8 -> Bi.Builder
buildWord8 = Bi.singleton

buildByteString :: ByteString -> Bi.Builder
buildByteString = Bi.fromByteString

buildInteger :: Integer -> Bi.Builder
buildInteger n =
  let signBit = if n < 0 then bit 6 else zeroBits
      ab = abs n
  in
    -- Refer to: https://gitlab.com/obsidian.systems/tezos-bake-monitor-lib/-/blob/2cf12e53072bcd966d471430ead25f597db5e23f/tezos-bake-monitor-lib/src/Tezos/Common/Binary.hs#L122
    if ab < 0x40 then Bi.singleton (Unsafe.fromIntegral @Integer @Word8 ab .|. signBit)
    else Bi.singleton (Unsafe.fromIntegral @Integer @Word8 (ab .&. 0x3f) .|. signBit .|. bit 7) <> writeZ (-6) ab

writeZ :: forall a. (Integral a, Bits a) => Int -> a -> Bi.Builder
writeZ offset n =
  if n < bit (7 - offset) then Bi.singleton $ Unsafe.fromIntegral @a @Word8 $ n `shift` offset
    else Bi.singleton (Unsafe.fromIntegral @a @Word8 (((n `shift` offset) .&. 0x7f) `setBit` 7)) <> writeZ (offset - 7) n

-- | Build a binary representation of a Zarith natural. This function is
-- partial, only defined for nonnegative 'Integer's.
--
-- The reason it's not defined to accept a 'Natural' is mostly to avoid
-- 'fromIntegral'/'toInteger' conversions at the use sites, since
-- not all libraries, notably @cryptonite@, support 'Natural'.
buildNatural :: Integer -> Bi.Builder
buildNatural n = assert (n >= 0) $
  if n < 0x80 then Bi.singleton (Unsafe.fromIntegral @Integer @Word8 n)
  else Bi.singleton (Unsafe.fromIntegral @Integer @Word8 (n .&. 0xff) .|. bit 7) <> writeZ (-7) n

buildDynamic :: (a -> Bi.Builder) -> (DynamicSize a) -> Bi.Builder
buildDynamic build (DynamicSize x) =
  let b = build x
  in Bi.putWord32be (Unsafe.fromIntegral @Int64 @Word32 $ LBS.length $ Bi.toLazyByteString b) <> b

buildText :: Text -> Bi.Builder
buildText n =
  buildByteString $ TE.encodeUtf8 n

-------------------------------------------------
-- Decode
-------------------------------------------------

getInteger :: Bi.Get Integer
getInteger = do
  b <- Bi.getWord8
  n <- if b `testBit` 7 then readZ 6 (fromIntegral $ b .&. 0x3f) else pure (fromIntegral $ b .&. 0x3f)
  pure $ if b `testBit` 6 then negate n else n

getNatural :: Bi.Get Integer
getNatural = do
  b <- Bi.getWord8
  if b `testBit` 7 then readZ 7 (fromIntegral $ b .&. 0x7f) else pure (fromIntegral $ b .&. 0x7f)

readZ :: forall a. (Integral a, Bits a) => Int -> a -> Bi.Get a
readZ offset n = do
  b <- Bi.getWord8
  if (b == 0) && (offset > 0) then fail "trailing zero" else pure ()
  let n' = (Unsafe.fromIntegral @Word8 @a (b .&. 0x7f) `shift` offset) .|. n
  if b `testBit` 7 then readZ (offset + 7) n' else pure n'

getDynamic :: (Bi.Get a) -> (Bi.Get (DynamicSize a))
getDynamic getter = do
  len <- Unsafe.fromIntegral @Word32 @Int <$> Bi.getWord32be
  DynamicSize <$> Bi.isolate len getter

{-# ANN getText ("HLint: ignore Redundant fmap" :: Text) #-}
getText :: Bi.Get Text
getText =
  fmap decodeUtf8' getByteString >>= \case
    Left err -> fail $ show err
    Right answer -> pure answer

getByteString :: Bi.Get ByteString
getByteString =
  LBS.toStrict <$> Bi.getRemainingLazyByteString