Operation.hs

-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Implementation of generic operations submission.
module Morley.Client.Action.Operation
  ( Result
  , runOperations
  , runOperationsNonEmpty
  -- helpers
  , dryRunOperationsNonEmpty
  ) where

import Control.Lens (has, (%=), (&~))
import Data.List (zipWith4)
import Data.List.NonEmpty qualified as NE
import Data.Singletons (Sing, SingI, sing)
import Data.Text qualified as T
import Fmt (blockListF, blockListF', listF, nameF, pretty, (+|), (|+))

import Morley.Client.Action.Common
import Morley.Client.Logging
import Morley.Client.RPC.Class
import Morley.Client.RPC.Error
import Morley.Client.RPC.Getters
import Morley.Client.RPC.Types
import Morley.Client.TezosClient
import Morley.Client.Types
import Morley.Micheline (StringEncode(..), TezosInt64, TezosMutez(..))
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias (AddressOrAlias(..))
import Morley.Tezos.Core
import Morley.Tezos.Crypto
import Morley.Util.ByteString

-- | Designates output of an operation.
data Result
instance OperationInfoDescriptor Result where
  type TransferInfo Result = ()
  type OriginationInfo Result = Address
  type RevealInfo Result = ()

logOperations
  :: forall (runMode :: RunMode) env m.
     ( WithClientLog env m
     , HasTezosClient m
     , SingI runMode -- We don't ask aliases with 'tezos-client' in 'DryRun' mode
     )
  => AddressOrAlias
  -> NonEmpty (OperationInfo ClientInput)
  -> m ()
logOperations sender ops = do
  let runMode = sing @runMode

      opName =
        if | all (has _OpTransfer) ops -> "transactions"
           | all (has _OpOriginate) ops -> "originations"
           | all (has _OpReveal) ops -> "reveals"
           | otherwise -> "operations"

      buildOp = \case
        (OpTransfer tx, mbAlias) ->
          buildTxDataWithAlias mbAlias tx
        (OpOriginate orig, _) ->
          odName orig |+ " (temporary alias)"
        (OpReveal rv, mbAlias) ->
          "Key " +| rdPublicKey rv |+ maybe "" (\a -> " (" +| a |+ ")") mbAlias

  sender' <- case sender of
    addr@AddressResolved{} -> case runMode of
      SRealRun -> AddressAlias <$> getAlias sender
      SDryRun  -> pure addr
    alias -> pure alias

  aliases <- case runMode of
    SRealRun ->
      forM ops $ \case
        OpTransfer (TransactionData tx) ->
          Just <$> (getAlias . AddressResolved $ tdReceiver tx)
        OpOriginate _ ->
          pure Nothing
        OpReveal r ->
          Just <$> (getAlias . AddressResolved . mkKeyAddress $ rdPublicKey r)
    SDryRun -> pure $ ops $> Nothing

  logInfo $ T.strip $ -- strip trailing newline
    "Running " +| opName +| " by " +| sender' |+ ":\n" +|
    blockListF' "-" buildOp (ops `NE.zip` aliases)

-- | Perform sequence of operations.
--
-- Returns operation hash (or @Nothing@ in case empty list was provided) and result of
-- each operation (nothing for transactions and an address for originated contracts
runOperations
  :: forall m env.
     ( HasTezosRpc m
     , HasTezosClient m
     , WithClientLog env m
     )
  => AddressOrAlias
  -> [OperationInfo ClientInput]
  -> m (Maybe OperationHash, [OperationInfo Result])
runOperations sender operations = case operations of
  [] -> return (Nothing, [])
  op : ops -> do
    (opHash, res) <- runOperationsNonEmpty sender $ op :| ops
    return $ (Just opHash, toList res)

-- | How many times to retry if an operation fails after injection
injectionRetryCount :: Natural
injectionRetryCount = 2

-- | Perform non-empty sequence of operations.
--
-- Returns operation hash and result of each operation
-- (nothing for transactions and an address for originated contracts).
runOperationsNonEmpty
  :: forall m env.
     ( HasTezosRpc m
     , HasTezosClient m
     , WithClientLog env m
     )
  => AddressOrAlias
  -> NonEmpty (OperationInfo ClientInput)
  -> m (OperationHash, NonEmpty (OperationInfo Result))
runOperationsNonEmpty sender operations =
  runOperationsNonEmptyHelper @'RealRun injectionRetryCount sender operations

-- | Flag that is used to determine @runOperationsNonEmptyHelper@ behaviour.
data RunMode = DryRun | RealRun

isRealRun :: forall (runMode :: RunMode). (SingI runMode) => Bool
isRealRun = case sing @runMode of
  SRealRun -> True
  SDryRun  -> False

-- | Type family which is used to determine the output type of the
-- @runOperationsNonEmptyHelper@.
type family RunResult (a :: RunMode) where
  RunResult 'DryRun = NonEmpty (AppliedResult, TezosMutez)
  RunResult 'RealRun = (OperationHash, NonEmpty (OperationInfo Result))

data SingRunResult :: RunMode -> Type where
  SDryRun :: SingRunResult 'DryRun
  SRealRun :: SingRunResult 'RealRun

type instance Sing = SingRunResult

instance SingI 'DryRun where
  sing = SDryRun

instance SingI 'RealRun where
  sing = SRealRun

-- | Perform dry-run for sequence of operations.
--
-- Returned @AppliedResult@ contains information about estimated limits,
-- storage changes, etc. Additionally, estimated fees are returned.
dryRunOperationsNonEmpty
  :: forall m env.
     ( HasTezosRpc m
     , HasTezosClient m
     , WithClientLog env m
     )
  => AddressOrAlias
  -> NonEmpty (OperationInfo ClientInput)
  -> m (NonEmpty (AppliedResult, TezosMutez))
dryRunOperationsNonEmpty sender operations =
  runOperationsNonEmptyHelper @'DryRun 0 sender operations

-- | Perform non-empty sequence of operations and either dry-run
-- and return estimated limits and fees or perform operation injection.
-- Behaviour is defined via @RunMode@ flag argument.
runOperationsNonEmptyHelper
  :: forall (runMode :: RunMode) m env.
     ( HasTezosRpc m
     , HasTezosClient m
     , WithClientLog env m
     , SingI runMode
     )
  => Natural
  -> AddressOrAlias
  -> NonEmpty (OperationInfo ClientInput)
  -> m (RunResult runMode)
runOperationsNonEmptyHelper retryCount sender operations = do
  logOperations @runMode sender operations
  senderAddress <- resolveAddress sender
  prohibitContractSender senderAddress $ head operations
  mbPassword <- getKeyPassword senderAddress
  when (isRealRun @runMode && mayNeedSenderRevealing (toList operations)) $
    revealKeyUnlessRevealed senderAddress mbPassword

  pp <- getProtocolParameters
  OperationConstants{..} <- preProcessOperation senderAddress

  let convertOps i = WithCommonOperationData commonData . \case
        OpTransfer (TransactionData TD {..}) ->
          OpTransfer TransactionOperation
          { toDestination = tdReceiver
          , toAmount = TezosMutez tdAmount
          , toParameters = toParametersInternals tdEpName tdParam
          }
        OpOriginate OriginationData{..} ->
          OpOriginate OriginationOperation
          { ooBalance = TezosMutez odBalance
          , ooScript = mkOriginationScript odContract odStorage
          }
        OpReveal RevealData{..} ->
          OpReveal RevealOperation
          { roPublicKey = rdPublicKey
          }
        where
          commonData = mkCommonOperationData senderAddress (ocCounter + i) pp

  let opsToRun = NE.zipWith convertOps (1 :| [(2 :: TezosInt64)..]) operations
      mbFees = operations <&> \case
        OpTransfer (TransactionData TD {..}) -> tdMbFee
        OpOriginate OriginationData{..} -> odMbFee
        OpReveal RevealData{..} -> rdMbFee

  -- Perform run_operation with dumb signature in order
  -- to estimate gas cost, storage size and paid storage diff
  let runOp = RunOperation
        { roOperation = RunOperationInternal
          { roiBranch = ocLastBlockHash
          , roiContents = opsToRun
          , roiSignature = stubSignature
          }
        , roChainId = bcChainId ocBlockConstants
        }
  results <- getAppliedResults (Left runOp)

  let -- Learn how to forge given operations
      forgeOp :: NonEmpty OperationInput -> m ByteString
      forgeOp ops =
        fmap unHexJSONByteString . forgeOperation $ ForgeOperation
          { foBranch = ocLastBlockHash
          , foContents = ops
          }

  let -- Attach a signature to forged operation + return the signature itself
      signForgedOp :: ByteString -> m (Signature, ByteString)
      signForgedOp op = do
        signature' <- signBytes sender mbPassword (addOperationPrefix op)
        return (signature', prepareOpForInjection op signature')

  -- Fill in fees
  let
    updateOp :: OperationInput -> Maybe Mutez -> AppliedResult -> Bool -> m (OperationInput, Maybe (Signature, ByteString))
    updateOp opToRun@(WithCommonOperationData _ internalOp) mbFee ar isFirst = do
      let gasSafetyGuard = 100
          -- @gasSafetyGuard@ is added to origination operations and transfers to non-implicit
          -- accounts, see https://gitlab.com/tezos/tezos/-/blob/v13.0/src/proto_013_PtJakart/lib_client/injection.ml#L750
          additionalGas = case internalOp of
            OpOriginate _ -> gasSafetyGuard
            OpTransfer (TransactionOperation {..}) -> case toDestination of
              KeyAddress _ -> 0
              ContractAddress _ -> gasSafetyGuard
              TransactionRollupAddress _ -> gasSafetyGuard
            OpReveal _ -> 0
      let storageLimit = computeStorageLimit [ar] pp + 20 -- similarly to tezos-client, we add 20 for safety
      let gasLimit = arConsumedGas ar + additionalGas
          updateCommonDataForFee fee =
            updateCommonData gasLimit storageLimit (TezosMutez fee)

      (_fee, op, mReadySignedOp) <- convergingFee
        @OperationInput
        @(Maybe (Signature, ByteString))  -- ready operation and its signature
        (\fee ->
          return $ opToRun &~
            wcoCommonDataL %= updateCommonDataForFee fee
          )
        (\op -> do
          forgedOp <- forgeOp $ one op
          -- In the Tezos implementation the first transaction
          -- in the series pays for signature.
          -- Signature of hash should be constant in size,
          -- so we can pass any signature, not necessarily the final one
          (fullForgedOpLength, mExtra) <-
            if isFirst
              then do
                res@(_signature, signedOp) <- signForgedOp forgedOp
                return (length signedOp, Just res)
              else
                -- Forge output automatically includes additional 32-bytes header
                -- which should be ommited for all operations in batch except the first one.
                pure (length forgedOp - 32, Nothing)
          return
            ( maybe (computeFee ocFeeConstants fullForgedOpLength gasLimit) id mbFee
            , mExtra
            )
          )

      return (op, mReadySignedOp)

  let
    zipWith4NE
      :: (a -> b -> c -> d -> e) -> NonEmpty a -> NonEmpty b -> NonEmpty c -> NonEmpty d
      -> NonEmpty e
    zipWith4NE f (a :| as) (b :| bs) (c :| cs) (d :| ds) =
      (f a b c d) :| zipWith4 f as bs cs ds
  -- These two lists must have the same length here.
  -- @opsToRun@ is constructed directly from @params@.
  -- The length of @results@ is checked in @getAppliedResults@.
  (updOps, readySignedOps) <- fmap NE.unzip . sequenceA $
    zipWith4NE updateOp opsToRun mbFees results (True :| repeat False)

  -- Forge operation with given limits and get its hexadecimal representation
  (signature', signedOp) <- case readySignedOps of
    -- Save one forge + sign call pair in case of one operation
    Just readyOp :| [] -> pure readyOp
    -- In case of batch we have to reforge the full operation
    _ -> forgeOp updOps >>= signForgedOp

  -- Operation still can fail due to insufficient gas or storage limit, so it's required
  -- to preapply it before injecting
  let preApplyOp = PreApplyOperation
        { paoProtocol = bcProtocol ocBlockConstants
        , paoBranch = ocLastBlockHash
        , paoContents = updOps
        , paoSignature = signature'
        }
  ars2 <- getAppliedResults (Right preApplyOp)
  case sing @runMode of
    SDryRun -> do
      let fees = map (codFee . wcoCommon) updOps
      return $ NE.zip ars2 fees
    SRealRun -> do
      operationHash <- waitForOperation $ injectOperation (HexJSONByteString signedOp)
      let contractAddrs = arOriginatedContracts <$> ars2
      opsRes <- forM (NE.zip operations contractAddrs) $ \case
        (OpTransfer _, []) ->
          return $ OpTransfer ()
        (OpTransfer _, addrs) -> do
          logInfo . T.strip $
            "The following contracts were originated during transactions: " +|
            listF addrs |+ ""
          return $ OpTransfer ()
        (OpOriginate _, []) ->
          throwM RpcOriginatedNoContracts
        (OpOriginate OriginationData{..}, [addr]) -> do
          logDebug $ "Saving " +| addr |+ " for " +| odName |+ "\n"
          rememberContract odReplaceExisting addr odName
          alias <- getAlias $ AddressResolved addr
          logInfo $ "Originated contract: " <> pretty alias
          return $ OpOriginate addr
        (OpOriginate _, addrs@(_ : _ : _)) ->
          throwM $ RpcOriginatedMoreContracts addrs
        (OpReveal _, _) ->
          return $ OpReveal ()
      forM_ ars2 logStatistics
      return (operationHash, opsRes)
      `catch` \case
        (UnexpectedRunErrors errs) | retryCount > 0 -> do
          logError $ pretty $ nameF "When injecting operations, there were unexpected errors" $
            blockListF errs
          runOperationsNonEmptyHelper @runMode (retryCount - 1) sender operations
        e -> throwM e
  where
    mayNeedSenderRevealing :: [OperationInfo i] -> Bool
    mayNeedSenderRevealing = any \case
      OpTransfer{} -> True
      OpOriginate{} -> True
      OpReveal{} -> False

    logStatistics :: AppliedResult -> m ()
    logStatistics ar = do
      let showTezosInt64 = show . unStringEncode
      logInfo $ "Consumed gas: " <> showTezosInt64 (arConsumedGas ar)
      logInfo $ "Storage size: " <> showTezosInt64 (arStorageSize ar)
      logInfo $ "Paid storage size diff: " <> showTezosInt64 (arPaidStorageDiff ar)

    prohibitContractSender :: Address -> OperationInfo ClientInput -> m ()
    prohibitContractSender addr op = case (addr, op) of
      (KeyAddress _, _) -> pass
      (ContractAddress _, op') -> throwM $ ContractSender addr (opName op')
      (TransactionRollupAddress _, op') -> throwM $ ContractSender addr (opName op')

    opName = \case
      OpTransfer _ -> "transfer"
      OpOriginate _ -> "origination"
      OpReveal _ -> "reveal"