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 Morley.Micheline.Expression
import 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 (fromIntegral ab .|. signBit)
    else Bi.singleton (fromIntegral (ab .&. 0x3f) .|. signBit .|. bit 7) <> writeZ (-6) ab

writeZ :: (Integral a, Bits a) => Int -> a -> Bi.Builder
writeZ offset n =
  if n < bit (7 - offset) then Bi.singleton $ fromIntegral $ n `shift` offset
    else Bi.singleton (fromIntegral (((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 (fromIntegral $ 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 (fromIntegral 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 :: (Num 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' = (fromIntegral (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 <- fromIntegral <$> 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)