-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Measuring operation size.
--
-- When originating a contract or making a transfer, tezos node forms operation
-- which is submitted over network. Size of this operation depends on content
-- of originated contract or transfer parameter resp., and tezos has a hard
-- limit on operation size thus it has to be accounted.
--
-- Functions declared in this module allow assessing size of origination or
-- transfer operation with up to constant precision because it yet accounts only
-- for Michelson primitives participating in the operation.
-- Other stuff which affects op size include parameters which user passes to
-- origination or transfer themselves, for instance, amount of mutez carried
-- to the contract. ATM we don't have necessary primitives in Haskell to be
-- able to handle those parameters here, probably waiting for [TM-89].
-- Currently, we can assess overall transfer size only approximatelly, like
-- in 'smallTransferOpSize'.
module Michelson.OpSize
  ( OpSize (..)
  , opSizeHardLimit
  , smallTransferOpSize
  , instrOpSize
  , expandedInstrsOpSize
  , valueOpSize
  ) where

import Prelude hiding (Ordering(..))

import Michelson.Interpret.Pack
import Michelson.Untyped

-- | Operation size in bytes.
--
-- We use newtype wrapper because there are different units of measure
-- (another one is gas, and we don't want to confuse them).
newtype OpSize = OpSize { unOpSize :: Word }
  deriving stock (Show, Eq, Ord)

instance Semigroup OpSize where
  OpSize a <> OpSize b = OpSize (a + b)
instance Monoid OpSize where
  mempty = OpSize 0

-- | Maximal operation size allowed by Tezos production nodes.
opSizeHardLimit :: OpSize
opSizeHardLimit = OpSize 16384

-- | Base cost of any transfer of 0 mutez with no extra parameters.
-- (Add 'valueOpSize <param>' to it to get assessment of actual transfer op size)
smallTransferOpSize :: OpSize
smallTransferOpSize = OpSize 162

instrOpSize :: InstrAbstract ExpandedOp -> OpSize
instrOpSize = (OpSize 2 <>) . \case
  EXT{} -> mempty
  DROPN n -> stackDepthOpSize n
  DROP -> mempty
  DUP va -> annsOpSize va
  SWAP -> mempty
  DIG n -> stackDepthOpSize n
  DUG n -> stackDepthOpSize n
  PUSH va t v -> annsOpSize va <> typeOpSize t <> valueOpSize v
  SOME ta va -> annsOpSize ta va
  NONE ta va t -> annsOpSize ta va <> typeOpSize t
  UNIT ta va -> annsOpSize ta va
  IF_NONE l r -> ifOpSize l r
  PAIR ta va fal far -> annsOpSize ta va fal far
  CAR va fa -> annsOpSize va fa
  CDR va fa -> annsOpSize va fa
  LEFT ta va fal far t -> annsOpSize ta va fal far <> typeOpSize t
  RIGHT ta va fal far t -> annsOpSize ta va fal far <> typeOpSize t
  IF_LEFT l r -> ifOpSize l r
  NIL ta va t -> annsOpSize ta va <> typeOpSize t
  CONS va -> annsOpSize va
  IF_CONS l r -> ifOpSize l r
  SIZE va -> annsOpSize va
  EMPTY_SET ta va ct -> annsOpSize ta va <> innerOpSize ct
  EMPTY_MAP ta va ct t -> annsOpSize ta va <> innerOpSize ct <> typeOpSize t
  EMPTY_BIG_MAP ta va ct t -> annsOpSize ta va <> innerOpSize ct <> typeOpSize t
  MAP va is -> annsOpSize va <> subcodeOpSize is
  ITER is -> subcodeOpSize is
  MEM va -> annsOpSize va
  GET va -> annsOpSize va
  UPDATE va -> annsOpSize va
  IF l r -> ifOpSize l r
  LOOP is -> expandedInstrsOpSize is
  LOOP_LEFT is -> expandedInstrsOpSize is
  LAMBDA va ti to is ->
    annsOpSize va <> typeOpSize ti <> typeOpSize to <> expandedInstrsOpSize is
  EXEC va -> annsOpSize va
  APPLY va -> annsOpSize va
  DIP is -> subcodeOpSize is
  DIPN n is -> stackDepthOpSize n <> subcodeOpSize is
  FAILWITH -> mempty
  CAST va t -> annsOpSize va <> typeOpSize t
  RENAME va -> annsOpSize va
  PACK va -> annsOpSize va
  UNPACK ta va t -> annsOpSize ta va <> typeOpSize t
  CONCAT va -> annsOpSize va
  SLICE va -> annsOpSize va
  ISNAT va -> annsOpSize va
  ADD va -> annsOpSize va
  SUB va -> annsOpSize va
  MUL va -> annsOpSize va
  EDIV va -> annsOpSize va
  ABS va -> annsOpSize va
  NEG va -> annsOpSize va
  LSL va -> annsOpSize va
  LSR va -> annsOpSize va
  OR va -> annsOpSize va
  AND va -> annsOpSize va
  XOR va -> annsOpSize va
  NOT va -> annsOpSize va
  COMPARE va -> annsOpSize va
  EQ va -> annsOpSize va
  NEQ va -> annsOpSize va
  LT va -> annsOpSize va
  LE va -> annsOpSize va
  GT va -> annsOpSize va
  GE va -> annsOpSize va
  INT va -> annsOpSize va
  SELF va fa -> annsOpSize va fa
  CONTRACT va fa t -> annsOpSize va fa <> typeOpSize t
  TRANSFER_TOKENS va -> annsOpSize va
  SET_DELEGATE va -> annsOpSize va
  CREATE_CONTRACT va1 va2 c -> annsOpSize va1 va2 <> contractOpSize c
  IMPLICIT_ACCOUNT va -> annsOpSize va
  NOW va -> annsOpSize va
  AMOUNT va -> annsOpSize va
  BALANCE va -> annsOpSize va
  CHECK_SIGNATURE va -> annsOpSize va
  SHA256 va -> annsOpSize va
  SHA512 va -> annsOpSize va
  BLAKE2B va -> annsOpSize va
  HASH_KEY va -> annsOpSize va
  SOURCE va -> annsOpSize va
  SENDER va -> annsOpSize va
  ADDRESS va -> annsOpSize va
  CHAIN_ID va -> annsOpSize va
  where
    subcodeOpSize is = expandedInstrOpSize (SeqEx is)
    ifOpSize l r = expandedInstrOpSize (SeqEx l) <> expandedInstrOpSize (SeqEx r)
    stackDepthOpSize n = OpSize 1 <> numOpSize n

expandedInstrOpSize :: ExpandedOp -> OpSize
expandedInstrOpSize = \case
  PrimEx i -> instrOpSize i
  SeqEx is -> OpSize 5 <> expandedInstrsOpSize is
  WithSrcEx _ i -> expandedInstrOpSize i

expandedInstrsOpSize :: [ExpandedOp] -> OpSize
expandedInstrsOpSize = foldMap expandedInstrOpSize

contractOpSize :: Contract -> OpSize
contractOpSize (Contract (ParameterType cp rootAnn) st is _) =
  OpSize 16 <> typeOpSize' [rootAnn] cp <> typeOpSize st <> expandedInstrsOpSize is

numOpSize :: Integral i => i -> OpSize
numOpSize = OpSize . fromIntegral . length . encodeIntPayload . fromIntegral

valueOpSize :: Value -> OpSize
valueOpSize = (OpSize 1 <>) . \case
  ValueInt i -> numOpSize i
  ValueString s -> (seqOpSize <>) . OpSize . fromIntegral $ length s
  ValueBytes (InternalByteString b) -> (seqOpSize <>) . OpSize . fromIntegral $ length b
  ValueUnit -> baseOpSize
  ValueTrue -> baseOpSize
  ValueFalse -> baseOpSize
  ValuePair l r -> baseOpSize <> valueOpSize l <> valueOpSize r
  ValueLeft x -> baseOpSize <> valueOpSize x
  ValueRight x -> baseOpSize <> valueOpSize x
  ValueSome x -> baseOpSize <> valueOpSize x
  ValueNone -> baseOpSize
  ValueNil -> seqOpSize
  ValueSeq l -> seqOpSize <> foldMap valueOpSize l
  ValueMap m -> seqOpSize <> foldMap eltOpSize m
  ValueLambda m -> seqOpSize <> foldMap expandedInstrOpSize m
  where
    baseOpSize = OpSize 1
    seqOpSize = OpSize 4
    eltOpSize (Elt k v) = OpSize 2 <> valueOpSize k <> valueOpSize v

typeOpSize :: Type -> OpSize
typeOpSize = typeOpSize' []

typeOpSize' :: [FieldAnn] -> Type -> OpSize
typeOpSize' anns (Type t ta) =
  tOpSize t <> annsOpSize ta anns

tOpSize :: T -> OpSize
tOpSize t =
  OpSize 2 <> case t of
    TKey -> mempty
    TUnit -> mempty
    TSignature -> mempty
    TChainId -> mempty
    TOption a -> typeOpSize a
    TList a -> typeOpSize a
    TSet a -> innerOpSize a
    TOperation -> mempty
    TContract a -> typeOpSize a
    TPair al ar l r -> typeOpSize' [al] l <> typeOpSize' [ar] r
    TOr al ar l r -> typeOpSize' [al] l <> typeOpSize' [ar] r
    TLambda i o -> typeOpSize i <> typeOpSize o
    TMap k v -> innerOpSize k <> typeOpSize v
    TBigMap k v -> innerOpSize k <> typeOpSize v
    TInt -> mempty
    TNat -> mempty
    TString -> mempty
    TBytes -> mempty
    TMutez -> mempty
    TBool -> mempty
    TKeyHash -> mempty
    TTimestamp -> mempty
    TAddress -> mempty

innerOpSize :: Type -> OpSize
innerOpSize (Type _ a) =
  (OpSize 2) <> annsOpSize a

-- | Accepts an arbitrary number of 'TypeAnn' 'FieldAnn' and/or 'VarAnn' that
-- belong to the same entity and returns their total operation size.
--
-- Note that annotations which belong to the same entity (type or instruction)
-- __must be__ considered in aggregate using one call of this function and to be
-- specified in order. See 'AnnotationSet' for details.
annsOpSize :: AnnsOpSizeVararg x => x
annsOpSize = annsOpSizeVararg emptyAnnSet

class AnnsOpSizeVararg x where
  annsOpSizeVararg :: AnnotationSet -> x

instance (KnownAnnTag t, AnnsOpSizeVararg x) => AnnsOpSizeVararg (Annotation t -> x) where
  annsOpSizeVararg acc an = annsOpSizeVararg (acc <> singleAnnSet an)

instance (KnownAnnTag t, AnnsOpSizeVararg x) => AnnsOpSizeVararg ([Annotation t] -> x) where
  annsOpSizeVararg acc ans = annsOpSizeVararg (acc <> singleGroupAnnSet ans)

instance AnnsOpSizeVararg OpSize where
  annsOpSizeVararg = annsOpSizeImpl

annsOpSizeImpl :: AnnotationSet -> OpSize
annsOpSizeImpl annSet
  | isNoAnnSet annSet = mempty
  | otherwise = OpSize . fromIntegral $ 3 * (minAnnSetSize annSet + 1)