Binary.hs

-- SPDX-FileCopyrightText: 2020 Tocqueville Group
-- SPDX-FileCopyrightText: 2018 obsidian.systems
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ
-- SPDX-License-Identifier: LicenseRef-MIT-obsidian-systems

-- | Module that define encoding and decoding function from Expression type
-- to binary format.
module Morley.Micheline.Binary
  ( decodeExpression
  , eitherDecodeExpression
  , encodeExpression
  , encodeExpression'
  ) where

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

import Morley.Micheline.Expression
import Morley.Util.Binary (UnpackError(..), ensureEnd, launchGet)

newtype DynamicSize a = DynamicSize { unDynamicSize :: a }

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

-- | Encode 'Expression' to 'ByteString'.
encodeExpression :: Expression -> LByteString
encodeExpression = Bi.toLazyByteString . buildExpr

-- | Same as 'encodeExpression', for strict bytestring.
encodeExpression' :: Expression -> BS.ByteString
encodeExpression' = LBS.toStrict . encodeExpression

buildExpr :: Expression -> Bi.Builder
buildExpr = \case
  ExpressionSeq xs -> buildWord8 2 <> buildDynamic buildList (DynamicSize xs)
  ExpressionPrim (MichelinePrimAp prim args annots) -> case (args, annots) of
    ([], []) -> buildWord8 3 <> buildPrim prim
    ([], _) -> buildWord8 4 <> buildPrim prim <> buildAnnotationList annots
    ([arg1], []) -> buildWord8 5 <> buildPrim prim <> buildExpr arg1
    ([arg1], _) -> buildWord8 6 <> buildPrim prim <> buildExpr arg1 <> buildAnnotationList annots
    ([arg1, arg2], []) -> buildWord8 7 <> buildPrim prim <> buildExpr arg1 <> buildExpr arg2
    ([arg1, arg2], _) -> buildWord8 8 <> buildPrim prim <> buildExpr arg1 <> buildExpr arg2 <> buildAnnotationList annots
    _ -> buildWord8 9 <> buildPrim prim <> buildDynamic buildList (DynamicSize args) <> buildAnnotationList annots
  ExpressionString x -> buildWord8 1 <> buildDynamic buildText (DynamicSize x)
  ExpressionInt x -> buildWord8 0 <> buildInteger x
  ExpressionBytes x -> buildWord8 10 <> buildDynamic buildByteString (DynamicSize x)

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

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

buildList :: [Expression] -> Bi.Builder
buildList = foldMap buildExpr

buildPrim :: MichelinePrimitive -> Bi.Builder
buildPrim (MichelinePrimitive p) = case Seq.elemIndexL p michelsonPrimitive of
  Nothing -> error $ "unknown Michelson/Micheline primitive: " <> p
  Just ix -> buildWord8 (Unsafe.fromIntegral @Int @Word8 ix)

buildAnnotationList :: [Annotation] -> Bi.Builder
buildAnnotationList listAnn = buildDynamic buildText (DynamicSize $ unwords . fmap annotToText $ listAnn)

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

-- | Decode 'Expression' from 'ByteString'.
eitherDecodeExpression :: BS.ByteString -> Either UnpackError Expression
eitherDecodeExpression x = launchGet (getExpr <* ensureEnd) $ LBS.fromStrict x

-- | Partial version of 'eitherDecodeExpression'.
decodeExpression :: HasCallStack => BS.ByteString -> Expression
decodeExpression = either (error . unUnpackError) id . eitherDecodeExpression

getExpr :: Bi.Get Expression
getExpr = Bi.getWord8 >>= \case
  0 -> ExpressionInt <$> getInteger
  1 -> ExpressionString . unDynamicSize <$> (getDynamic getText)
  2 -> ExpressionSeq . unDynamicSize <$> (getDynamic getList)
  3 -> ExpressionPrim . (\pn -> MichelinePrimAp pn [] []) <$> getPrim
  4 -> ExpressionPrim <$> (flip MichelinePrimAp [] <$> getPrim <*> getAnnotationList)
  5 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> (one <$> getExpr) <*> pure [])
  6 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> (one <$> getExpr) <*> getAnnotationList)
  7 -> ExpressionPrim <$> ((\n a -> MichelinePrimAp n a []) <$> getPrim <*> replicateM 2 getExpr)
  8 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> replicateM 2 getExpr <*> getAnnotationList)
  9 -> ExpressionPrim <$> (MichelinePrimAp <$> getPrim <*> (unDynamicSize <$> (getDynamic getList)) <*> getAnnotationList)
  10 -> ExpressionBytes . unDynamicSize <$> (getDynamic getByteString)
  _ -> fail "invalid Micheline expression tag"

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

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

getList :: Bi.Get [Expression]
getList = many getExpr

getPrim :: Bi.Get MichelinePrimitive
getPrim =
  Bi.getWord8 >>= \ix -> case Seq.lookup (fromIntegral ix) michelsonPrimitive of
    Nothing -> fail "unknown Michelson/Micheline opcode"
    Just str -> pure $ MichelinePrimitive str

getAnnotationList :: Bi.Get [Annotation]
getAnnotationList = mapM annotFromText . words . unDynamicSize =<< (getDynamic getText)