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

-- | General-purpose utility functions for typed types.

module Morley.Michelson.Typed.Util
  ( -- * Instruction analysis
    DfsSettings (..)
  , CtorEffectsApp (..)
  , dfsTraverseInstr
  , dfsFoldInstr
  , dfsModifyInstr

  , isMichelsonInstr

  -- * Changing instruction tree structure
  , linearizeLeft
  , linearizeLeftDeep

  -- * Value analysis
  , dfsFoldMapValue
  , dfsFoldMapValueM
  , dfsMapValue
  , dfsTraverseValue
  , isStringValue
  , isBytesValue
  , allAtomicValues

  -- * Instruction generation
  , PushableStorageSplit (..)
  , splitPushableStorage

  -- * Working with 'RemFail'
  , analyzeInstrFailure

  -- * Annotations
  , SomeAnns(..)
  , instrAnns
  ) where

import Debug qualified (show)
import Prelude hiding (Ordering(..))

import Control.Monad.Writer.Strict (Writer, execWriterT, runWriter, tell, writer)
import Data.Constraint (Dict(..))
import Data.Default (Default(..))
import Data.Map qualified as M
import Data.Set qualified as S
import Data.Text.Internal.Builder (fromText)
import Fmt (Buildable(..))

import Morley.Michelson.Text (MText)
import Morley.Michelson.Typed.Aliases
import Morley.Michelson.Typed.Annotation
import Morley.Michelson.Typed.Contract
import Morley.Michelson.Typed.Instr
import Morley.Michelson.Typed.Scope
import Morley.Michelson.Typed.T qualified as T
import Morley.Michelson.Typed.Value
import Morley.Michelson.Typed.View
import Morley.Michelson.Untyped (AnyAnn)

-- $setup
-- >>> :m + Morley.Michelson.Typed.Instr

-- | Options for 'dfsTraverseInstr' family of functions.
data DfsSettings m = DfsSettings
  { dsGoToValues :: Bool
    -- ^ Whether @dfsTraverseInstr@ function should go into values which contain other
    -- instructions: lambdas and constant contracts
    -- (which can be passed to @CREATE_CONTRACT@).
  , dsCtorEffectsApp :: CtorEffectsApp m
    -- ^ How do we handle intermediate nodes in instruction tree.
  , dsInstrStep :: (forall i o. Instr i o -> m (Instr i o))
  , dsValueStep :: (forall t'. Value t' -> m (Value t'))
  }

-- | Describes how intermediate nodes in instruction tree are accounted.
data CtorEffectsApp m = CtorEffectsApp
  { ceaName :: Text
    -- ^ Name of this way.
  , ceaPostStep
      :: forall i o. Monad m
      => Instr i o
      -> m (Instr i o)
      -> m (Instr i o)
    -- ^ This transformation is applied after the step.
    -- It will be provided with the old instruction and the action gathered
    -- with the recursive traversal for the instruction subtree,
    -- and the result will go to the parent node.
  }

instance Buildable (CtorEffectsApp x) where
  build CtorEffectsApp{..} = fromText ceaName

instance (Applicative x) => Default (DfsSettings x) where
  def = DfsSettings
    { dsGoToValues = False
    , dsCtorEffectsApp = CtorEffectsApp
        { ceaName = "Do nothing"
        , ceaPostStep = const id
        }
    , dsInstrStep = pure
    , dsValueStep = pure
    }

-- | Traverse a typed instruction in depth-first order.
--
-- The 'dsInstrStep' and 'dsValueStep' actions will be applied in bottom-to-top order, i.e.
-- first to the children of a node, then to the node itself.
dfsTraverseInstr ::
     forall m inp out.
     (Monad m)
  => DfsSettings m
  -> Instr inp out
  -> m (Instr inp out)
dfsTraverseInstr settings@DfsSettings{..} i =
  case i of
    Seq i1 i2 -> recursion2 Seq i1 i2
    WithLoc loc i1 -> recursion1 (WithLoc loc) i1
    Meta meta i1 -> recursion1 (Meta meta) i1
    FrameInstr p i1 -> recursion1 (FrameInstr p) i1
    Nested i1 -> recursion1 Nested i1
    DocGroup dg i1 -> recursion1 (DocGroup dg) i1
    IF_NONE i1 i2 -> recursion2 IF_NONE i1 i2
    IF_LEFT i1 i2 -> recursion2 IF_LEFT i1 i2
    IF_CONS i1 i2 -> recursion2 IF_CONS i1 i2
    IF i1 i2 -> recursion2 IF i1 i2
    AnnMAP ann i1 -> recursion1 (AnnMAP ann) i1
    ITER i1 -> recursion1 ITER i1
    LOOP i1 -> recursion1 LOOP i1
    LOOP_LEFT i1 -> recursion1 LOOP_LEFT i1
    DIP i1 -> recursion1 DIP i1
    DIPN s i1 -> recursion1 (DIPN s) i1

    -- This case is more complex so we duplicate @recursion1@ a bit.
    -- We have to traverse the pushed value because a lambda can be
    -- somewhere inside of it (e. g. one item of a pair).
    AnnPUSH ann v
      | dsGoToValues -> ceaPostStep dsCtorEffectsApp i do
        innerV <- dfsTraverseValue settings v
        dsInstrStep $ AnnPUSH ann innerV
     | otherwise -> recursion0 i

    AnnLAMBDA ann (VLam i1)
      | dsGoToValues ->
          recursion1 (AnnLAMBDA ann . VLam . analyzeInstrFailure) (rfAnyInstr i1)
      | otherwise -> recursion0 i
    AnnCREATE_CONTRACT ann contract
      | dsGoToValues ->
          ceaPostStep dsCtorEffectsApp i
            case cCode contract of
              ContractCode c -> do
                codeI <- dfsTraverseInstr settings c
                viewsI <- forM (toList $ cViews contract) \(SomeView v) -> do
                    code <- dfsTraverseInstr settings $ vCode v
                    return $ SomeView v{ vCode = code }
                dsInstrStep $ AnnCREATE_CONTRACT ann $ contract
                  { cCode = ContractCode codeI
                  , cViews = UnsafeViewsSet $ fromList viewsI
                  }
      | otherwise -> recursion0 i

    Nop{} -> recursion0 i
    Ext (TEST_ASSERT (TestAssert nm pc i1)) ->
      recursion1 (Ext . TEST_ASSERT . TestAssert nm pc) i1
    Ext{} -> recursion0 i
    AnnCAR{} -> recursion0 i
    AnnCDR{} -> recursion0 i
    DROP{} -> recursion0 i
    DROPN{} -> recursion0 i
    AnnDUP{} -> recursion0 i
    AnnDUPN{} -> recursion0 i
    SWAP{} -> recursion0 i
    DIG{} -> recursion0 i
    DUG{} -> recursion0 i
    AnnSOME{} -> recursion0 i
    AnnNONE{} -> recursion0 i
    AnnUNIT{} -> recursion0 i
    AnnPAIR{} -> recursion0 i
    AnnUNPAIR{} -> recursion0 i
    AnnPAIRN{} -> recursion0 i
    UNPAIRN{} -> recursion0 i
    AnnLEFT{} -> recursion0 i
    AnnRIGHT{} -> recursion0 i
    AnnNIL{} -> recursion0 i
    AnnCONS{} -> recursion0 i
    AnnSIZE{} -> recursion0 i
    AnnEMPTY_SET{} -> recursion0 i
    AnnEMPTY_MAP{} -> recursion0 i
    AnnEMPTY_BIG_MAP{} -> recursion0 i
    AnnMEM{} -> recursion0 i
    AnnGET{} -> recursion0 i
    AnnGETN{} -> recursion0 i
    AnnUPDATE{} -> recursion0 i
    AnnUPDATEN{} -> recursion0 i
    AnnGET_AND_UPDATE{} -> recursion0 i
    AnnEXEC{} -> recursion0 i
    AnnAPPLY{} -> recursion0 i
    FAILWITH{} -> recursion0 i
    AnnCAST{} -> recursion0 i
    AnnRENAME{} -> recursion0 i
    AnnPACK{} -> recursion0 i
    AnnUNPACK{} -> recursion0 i
    AnnCONCAT{} -> recursion0 i
    AnnCONCAT'{} -> recursion0 i
    AnnSLICE{} -> recursion0 i
    AnnISNAT{} -> recursion0 i
    AnnADD{} -> recursion0 i
    AnnSUB{} -> recursion0 i
    AnnSUB_MUTEZ{} -> recursion0 i
    AnnMUL{} -> recursion0 i
    AnnEDIV{} -> recursion0 i
    AnnABS{} -> recursion0 i
    AnnNEG{} -> recursion0 i
    AnnLSL{} -> recursion0 i
    AnnLSR{} -> recursion0 i
    AnnOR{} -> recursion0 i
    AnnAND{} -> recursion0 i
    AnnXOR{} -> recursion0 i
    AnnNOT{} -> recursion0 i
    AnnCOMPARE{} -> recursion0 i
    AnnEQ{} -> recursion0 i
    AnnNEQ{} -> recursion0 i
    AnnLT{} -> recursion0 i
    AnnGT{} -> recursion0 i
    AnnLE{} -> recursion0 i
    AnnGE{} -> recursion0 i
    AnnINT{} -> recursion0 i
    AnnVIEW{} -> recursion0 i
    AnnSELF{} -> recursion0 i
    AnnCONTRACT{} -> recursion0 i
    AnnTRANSFER_TOKENS{} -> recursion0 i
    AnnSET_DELEGATE{} -> recursion0 i
    AnnIMPLICIT_ACCOUNT{} -> recursion0 i
    AnnNOW{} -> recursion0 i
    AnnAMOUNT{} -> recursion0 i
    AnnBALANCE{} -> recursion0 i
    AnnVOTING_POWER{} -> recursion0 i
    AnnTOTAL_VOTING_POWER{} -> recursion0 i
    AnnCHECK_SIGNATURE{} -> recursion0 i
    AnnSHA256{} -> recursion0 i
    AnnSHA512{} -> recursion0 i
    AnnBLAKE2B{} -> recursion0 i
    AnnSHA3{} -> recursion0 i
    AnnKECCAK{} -> recursion0 i
    AnnHASH_KEY{} -> recursion0 i
    AnnPAIRING_CHECK{} -> recursion0 i
    AnnSOURCE{} -> recursion0 i
    AnnSENDER{} -> recursion0 i
    AnnADDRESS{} -> recursion0 i
    AnnCHAIN_ID{} -> recursion0 i
    AnnLEVEL{} -> recursion0 i
    AnnSELF_ADDRESS{} -> recursion0 i
    NEVER{} -> recursion0 i
    AnnTICKET{} -> recursion0 i
    AnnREAD_TICKET{} -> recursion0 i
    AnnSPLIT_TICKET{} -> recursion0 i
    AnnJOIN_TICKETS{} -> recursion0 i
    AnnOPEN_CHEST{} -> recursion0 i
    AnnSAPLING_EMPTY_STATE{} -> recursion0 i
    AnnSAPLING_VERIFY_UPDATE{} -> recursion0 i
    AnnMIN_BLOCK_TIME{} -> recursion0 i
    AnnEMIT{} -> recursion0 i
  where
    recursion0 ::
      forall a b. Instr a b -> m (Instr a b)
    recursion0 i0 =
      ceaPostStep dsCtorEffectsApp i0 $
      dsInstrStep i0

    recursion1 ::
      forall a b c d. (Instr a b -> Instr c d) -> Instr a b -> m (Instr c d)
    recursion1 constructor i0 =
      ceaPostStep dsCtorEffectsApp (constructor i0) do
        innerI <- dfsTraverseInstr settings i0
        dsInstrStep $ constructor innerI

    recursion2 ::
      forall i o i1 o1 i2 o2.
      (Instr i1 o1 -> Instr i2 o2 -> Instr i o) ->
      Instr i1 o1 -> Instr i2 o2 -> m (Instr i o)
    recursion2 constructor i1 i2 =
      ceaPostStep dsCtorEffectsApp (constructor i1 i2) do
        i1' <- dfsTraverseInstr settings i1
        i2' <- dfsTraverseInstr settings i2
        dsInstrStep $ constructor i1' i2'

-- | Specialization of 'dfsTraverseInstr' for case when changing the instruction is
-- not required.
dfsFoldInstr
  :: forall x inp out.
      (Monoid x)
  => DfsSettings (Writer x)
  -> (forall i o. Instr i o -> x)
  -> Instr inp out
  -> x
dfsFoldInstr settings step =
  snd . runWriter . dfsTraverseInstr settings{dsInstrStep = writer . (id &&& step)}

-- | Specialization of 'dfsTraverseInstr' which only modifies given instruction.
dfsModifyInstr
  :: DfsSettings Identity
  -> (forall i o. Instr i o -> Instr i o)
  -> Instr inp out
  -> Instr inp out
dfsModifyInstr settings step =
  runIdentity . dfsTraverseInstr settings{dsInstrStep = (pure . step)}

-- | Check whether instruction fails at each execution path or have at least one
-- non-failing path.
--
-- This function assumes that given instruction contains no dead code
-- (contract with dead code cannot be valid Michelson contract) and may behave
-- in unexpected way if such is present. Term "dead code" includes instructions
-- which render into empty Michelson, like Morley extensions.
-- On the other hand, this function does not traverse the whole instruction tree;
-- performs fastest on left-growing combs.
--
-- Often we already have information about instruction failure, use this
-- function only in cases when this info is actually unavailable or hard
-- to use.
analyzeInstrFailure :: HasCallStack => Instr i o -> RemFail Instr i o
analyzeInstrFailure = go
  where
  go :: Instr i o -> RemFail Instr i o
  go = \case
    WithLoc loc i -> case go i of
      RfNormal i0 ->
        RfNormal (WithLoc loc i0)
      r -> r
    Meta meta i -> case go i of
      RfNormal i0 ->
        RfNormal (Meta meta i0)
      r -> r
    FrameInstr s i -> case go i of
      RfNormal i0 ->
        RfNormal (FrameInstr s i0)
      RfAlwaysFails i0 ->
        error $ "FrameInstr wraps always-failing instruction: " <> Debug.show i0
    Seq a b -> Seq a `rfMapAnyInstr` go b
    Nop -> RfNormal Nop
    Ext e -> RfNormal (Ext e)
    Nested i -> Nested `rfMapAnyInstr` go i
    DocGroup g i -> DocGroup g `rfMapAnyInstr` go i

    IF_NONE l r -> rfMerge IF_NONE (go l) (go r)
    IF_LEFT l r -> rfMerge IF_LEFT (go l) (go r)
    IF_CONS l r -> rfMerge IF_CONS (go l) (go r)
    IF l r -> rfMerge IF (go l) (go r)

    i@AnnMAP{} -> RfNormal i
    i@ITER{} -> RfNormal i
    i@LOOP{} -> RfNormal i
    i@LOOP_LEFT{} -> RfNormal i
    i@AnnLAMBDA{} -> RfNormal i
    i@DIP{} -> RfNormal i
    i@DIPN{} -> RfNormal i

    i@AnnCAR{} -> RfNormal i
    i@AnnCDR{} -> RfNormal i
    i@DROP{} -> RfNormal i
    i@DROPN{} -> RfNormal i
    i@AnnDUP{} -> RfNormal i
    i@AnnDUPN{} -> RfNormal i
    i@SWAP{} -> RfNormal i
    i@DIG{} -> RfNormal i
    i@DUG{} -> RfNormal i
    i@AnnPUSH{} -> RfNormal i
    i@AnnSOME{} -> RfNormal i
    i@AnnNONE{} -> RfNormal i
    i@AnnUNIT{} -> RfNormal i
    i@AnnPAIR{} -> RfNormal i
    i@AnnUNPAIR{} -> RfNormal i
    i@AnnPAIRN{} -> RfNormal i
    i@UNPAIRN{} -> RfNormal i
    i@AnnLEFT{} -> RfNormal i
    i@AnnRIGHT{} -> RfNormal i
    i@AnnNIL{} -> RfNormal i
    i@AnnCONS{} -> RfNormal i
    i@AnnSIZE{} -> RfNormal i
    i@AnnEMPTY_SET{} -> RfNormal i
    i@AnnEMPTY_MAP{} -> RfNormal i
    i@AnnEMPTY_BIG_MAP{} -> RfNormal i
    i@AnnMEM{} -> RfNormal i
    i@AnnGET{} -> RfNormal i
    i@AnnGETN{} -> RfNormal i
    i@AnnUPDATE{} -> RfNormal i
    i@AnnUPDATEN{} -> RfNormal i
    i@AnnGET_AND_UPDATE{} -> RfNormal i
    i@AnnEXEC{} -> RfNormal i
    i@AnnAPPLY{} -> RfNormal i
    FAILWITH -> RfAlwaysFails FAILWITH
    i@AnnCAST{} -> RfNormal i
    i@AnnRENAME{} -> RfNormal i
    i@AnnPACK{} -> RfNormal i
    i@AnnUNPACK{} -> RfNormal i
    i@AnnCONCAT{} -> RfNormal i
    i@AnnCONCAT'{} -> RfNormal i
    i@AnnSLICE{} -> RfNormal i
    i@AnnISNAT{} -> RfNormal i
    i@AnnADD{} -> RfNormal i
    i@AnnSUB{} -> RfNormal i
    i@AnnSUB_MUTEZ{} -> RfNormal i
    i@AnnMUL{} -> RfNormal i
    i@AnnEDIV{} -> RfNormal i
    i@AnnABS{} -> RfNormal i
    i@AnnNEG{} -> RfNormal i
    i@AnnLSL{} -> RfNormal i
    i@AnnLSR{} -> RfNormal i
    i@AnnOR{} -> RfNormal i
    i@AnnAND{} -> RfNormal i
    i@AnnXOR{} -> RfNormal i
    i@AnnNOT{} -> RfNormal i
    i@AnnCOMPARE{} -> RfNormal i
    i@AnnEQ{} -> RfNormal i
    i@AnnNEQ{} -> RfNormal i
    i@AnnLT{} -> RfNormal i
    i@AnnGT{} -> RfNormal i
    i@AnnLE{} -> RfNormal i
    i@AnnGE{} -> RfNormal i
    i@AnnINT{} -> RfNormal i
    i@AnnVIEW{} -> RfNormal i
    i@AnnSELF{} -> RfNormal i
    i@AnnCONTRACT{} -> RfNormal i
    i@AnnTRANSFER_TOKENS{} -> RfNormal i
    i@AnnSET_DELEGATE{} -> RfNormal i
    i@AnnCREATE_CONTRACT{} -> RfNormal i
    i@AnnIMPLICIT_ACCOUNT{} -> RfNormal i
    i@AnnNOW{} -> RfNormal i
    i@AnnAMOUNT{} -> RfNormal i
    i@AnnBALANCE{} -> RfNormal i
    i@AnnVOTING_POWER{} -> RfNormal i
    i@AnnTOTAL_VOTING_POWER{} -> RfNormal i
    i@AnnCHECK_SIGNATURE{} -> RfNormal i
    i@AnnSHA256{} -> RfNormal i
    i@AnnSHA512{} -> RfNormal i
    i@AnnBLAKE2B{} -> RfNormal i
    i@AnnSHA3{} -> RfNormal i
    i@AnnKECCAK{} -> RfNormal i
    i@AnnHASH_KEY{} -> RfNormal i
    i@AnnPAIRING_CHECK{} -> RfNormal i
    i@AnnSOURCE{} -> RfNormal i
    i@AnnSENDER{} -> RfNormal i
    i@AnnADDRESS{} -> RfNormal i
    i@AnnCHAIN_ID{} -> RfNormal i
    i@AnnLEVEL{} -> RfNormal i
    i@AnnSELF_ADDRESS{} -> RfNormal i
    NEVER -> RfAlwaysFails NEVER
    i@AnnTICKET{} -> RfNormal i
    i@AnnREAD_TICKET{} -> RfNormal i
    i@AnnSPLIT_TICKET{} -> RfNormal i
    i@AnnJOIN_TICKETS{} -> RfNormal i
    i@AnnOPEN_CHEST{} -> RfNormal i
    i@AnnSAPLING_EMPTY_STATE{} -> RfNormal i
    i@AnnSAPLING_VERIFY_UPDATE{} -> RfNormal i
    i@AnnMIN_BLOCK_TIME{} -> RfNormal i
    i@AnnEMIT{} -> RfNormal i

-- | There are many ways to represent a sequence of more than 2 instructions.
-- E. g. for @i1; i2; i3@ it can be @Seq i1 $ Seq i2 i3@ or @Seq (Seq i1 i2) i3@.
-- This function enforces a particular structure. Specifically, it makes each
-- v'Seq' have a single instruction (i. e. not v'Seq') in its second argument.
-- This function also erases redundant 'Nop's.
--
-- Please note that this function is not recursive, it does not
-- linearize contents of @IF@ and similar instructions.
linearizeLeft :: Instr inp out -> Instr inp out
linearizeLeft = linearizeLeftHelper False
  where
    -- In order to avoid quadratic performance we make a simple optimization.
    -- We track whether left argument of `Seq` is already linearized.
    -- If it is, we do not need to ever linearize it again.
    linearizeLeftHelper :: Bool -> Instr inp out -> Instr inp out
    linearizeLeftHelper isLeftInstrAlreadyLinear =
      \case
        Seq i1 (Seq i2 i3) ->
          linearizeLeftHelper True $
          Seq (linearizeLeftHelper isLeftInstrAlreadyLinear (Seq i1 i2)) i3
        -- `i2` is not a `Seq`, so we only need to linearize `i1`
        -- and connect it with `i2`.
        Seq i1 i2
          | isLeftInstrAlreadyLinear
          , Nop <- i2 -> i1
          | isLeftInstrAlreadyLinear -> Seq i1 i2
          | Nop <- i2 -> linearizeLeft i1
          | otherwise -> Seq (linearizeLeft i1) i2
        i -> i

-- | \"Deep\" version of 'linearizeLeft'. It recursively linearizes
-- instructions stored in other instructions.
linearizeLeftDeep :: Instr inp out -> Instr inp out
linearizeLeftDeep = dfsModifyInstr def linearizeLeft

----------------------------------------------------------------------------
-- Value analysis
----------------------------------------------------------------------------

-- | Traverse a value in depth-first order.
dfsMapValue ::
     forall t.
     DfsSettings Identity
  -> Value t
  -> Value t
dfsMapValue settings v = runIdentity $ dfsTraverseValue settings v

-- | Traverse a value in depth-first order.
dfsTraverseValue ::
     forall t m.
     (Monad m)
  => DfsSettings m
  -> Value t
  -> m (Value t)
dfsTraverseValue settings@DfsSettings{..} i = case i of
  -- Atomic
  VKey{} -> dsValueStep i
  VUnit -> dsValueStep i
  VSignature{} -> dsValueStep i
  VChainId{} -> dsValueStep i
  VOp{} -> dsValueStep i
  VContract{} -> dsValueStep i
  VTicket{} -> dsValueStep i  -- cannot appear as constant in a contract
  VLam lambda -> do
    v <- fmap (VLam . analyzeInstrFailure) $ dfsTraverseInstr settings (rfAnyInstr lambda)
    dsValueStep v
  VInt{} -> dsValueStep i
  VNat{} -> dsValueStep i
  VString{} -> dsValueStep i
  VBytes{} -> dsValueStep i
  VMutez{} -> dsValueStep i
  VBool{} -> dsValueStep i
  VKeyHash{} -> dsValueStep i
  VBls12381Fr{} -> dsValueStep i
  VBls12381G1{} -> dsValueStep i
  VBls12381G2{} -> dsValueStep i
  VTimestamp{} -> dsValueStep i
  VAddress{} -> dsValueStep i
  VChestKey{} -> dsValueStep i
  VChest{} -> dsValueStep i
  VTxRollupL2Address{} -> dsValueStep i

  -- Non-atomic
  VOption mVal -> case mVal of
    Nothing -> dsValueStep i
    Just val -> recursion1 (VOption . Just) val
  VList vals -> do
    vs <- traverse (dfsTraverseValue settings) vals
    dsValueStep $ VList vs
  VSet vals -> do
    cs <- S.fromList <$> traverse (dfsTraverseValue settings) (S.toList vals)
    dsValueStep (VSet cs)
  VPair (v1, v2) -> do
    v1' <- dfsTraverseValue settings v1
    v2' <- dfsTraverseValue settings v2
    dsValueStep $ VPair (v1', v2')
  VOr vEither -> case vEither of
    Left v -> recursion1 (VOr . Left) v
    Right v -> recursion1 (VOr . Right) v
  VMap vmap -> mapRecursion VMap vmap
  VBigMap bmId vmap -> mapRecursion (VBigMap bmId) vmap
  where
    recursion1 ::
         forall t'.
         (Value t' -> Value t)
      -> Value t'
      -> m (Value t)
    recursion1 constructor v = do
      v' <- dfsTraverseValue settings v
      dsValueStep $ constructor v'

    mapRecursion
      :: forall k v. Comparable k
      => (M.Map (Value k) (Value v) -> Value t)
      -> M.Map (Value k) (Value v)
      -> m (Value t)
    mapRecursion constructor vmap = do
      vmap' <-
        M.fromList <$> forM (M.toList vmap) \(k, v) -> do
          k' <- dfsTraverseValue settings k
          v' <- dfsTraverseValue settings v
          pure (k', v')
      dsValueStep $ constructor vmap'

-- | Specialization of 'dfsMapValue' for case when changing the value is
-- not required.
dfsFoldMapValue
  :: Monoid x
  => (forall t'. Value t' -> x)
  -> Value t
  -> x
dfsFoldMapValue step v =
  runIdentity $ dfsFoldMapValueM (pure . step) v

-- | Specialization of 'dfsMapValue' for case when changing the value is
-- not required.
dfsFoldMapValueM
  :: (Monoid x, Monad m)
  => (forall t'. Value t' -> m x)
  -> Value t
  -> m x
dfsFoldMapValueM step v = do
  execWriterT $
    dfsTraverseValue
      (def{ dsValueStep =
        (\val -> do
            x <- lift $ step val
            tell x
            pure val
        )
      }) v

-- | If value is a string, return the stored string.
isStringValue :: Value t -> Maybe MText
isStringValue =
  \case
    VString str -> Just str
    _ -> Nothing

-- | If value is a bytestring, return the stored bytestring.
isBytesValue :: Value t -> Maybe ByteString
isBytesValue =
  \case
    VBytes bytes -> Just bytes
    _ -> Nothing

-- | Takes a selector which checks whether a value can be converted
-- to something. Recursively applies it to all values. Collects extracted
-- values in a list.
allAtomicValues ::
  forall t a. (forall t'. Value t' -> Maybe a) -> Value t -> [a]
allAtomicValues selector = dfsFoldMapValue (maybeToList . selector)


--------------------------------------------------------------------------------
-- Instruction generation
--------------------------------------------------------------------------------

-- | Result of splitting a storage 'Value' of @st@ on the stack @s@.
--
-- The idea behind this is to either: prove that the whole 'Value' can be put on
-- the stack without containing a single @big_map@ or to split it into:
-- a 'Value' containing its @big_map@s and an instruction to reconstruct the
-- storage.
--
-- The main idea behind this is to create a large storage in Michelson code to
-- then create a contract using @CREATE_CONTRACT@.
-- Note: a simpler solution would have been to replace @big_map@ 'Value's with
-- an 'EMPTY_BIG_MAP' followed by many 'UPDATE' to push its content, but sadly
-- a bug (tezos/tezos/1154) prevents this from being done.
data PushableStorageSplit s st where
  -- | The type of the storage is fully constant.
  ConstantStorage
    :: ConstantScope st
    => Value st
    -> PushableStorageSplit s st

  -- | The type of the storage is not a constant, but its value does not contain
  -- @big_map@s. E.g. A 'Right ()' value of type 'Either (BigMap k v) ()'.
  PushableValueStorage
    :: StorageScope st
    => Instr s (st ': s)
    -> PushableStorageSplit s st

  -- | The type of the storage and part of its value (here @heavy@) contain one or
  -- more @big_map@s or @ticket@s. The instruction can take the non-pushable
  -- 'Value heavy' and reconstruct the original 'Value st' without using any
  -- 'EMPTY_BIG_MAP'.
  PartlyPushableStorage
    :: (StorageScope heavy, StorageScope st)
    => Value heavy -> Instr (heavy ': s) (st ': s)
    -> PushableStorageSplit s st

-- | Splits the given storage 'Value' into a 'PushableStorageSplit'.
--
-- This is based off the fact that the only storages that cannot be directly
-- 'PUSH'ed are the ones that contain 'Morley.Michelson.Typed.Haskell.Value.BigMap's and tickets.
-- See difference between 'StorageScope' and 'ConstantScope'.
--
-- So what we do here is to create a 'Value' as small as possible with all the
-- @big_map@s in it (if any) and an 'Instr' that can use it to rebuild the original
-- storage 'Value'.
--
-- Note: This is done this way to avoid using 'EMPTY_BIG_MAP' instructions, see
-- 'PushableStorageSplit' for motivation.
splitPushableStorage :: StorageScope t => Value t -> PushableStorageSplit s t
splitPushableStorage v = case v of
  -- Atomic (except op and contract)
  VKey{}        -> ConstantStorage v
  VUnit         -> ConstantStorage v
  VSignature{}  -> ConstantStorage v
  VChainId{}    -> ConstantStorage v
  VLam{}       -> ConstantStorage v
  VInt{}        -> ConstantStorage v
  VNat{}        -> ConstantStorage v
  VString{}     -> ConstantStorage v
  VBytes{}      -> ConstantStorage v
  VMutez{}      -> ConstantStorage v
  VBool{}       -> ConstantStorage v
  VKeyHash{}    -> ConstantStorage v
  VBls12381Fr{} -> ConstantStorage v
  VBls12381G1{} -> ConstantStorage v
  VBls12381G2{} -> ConstantStorage v
  VTimestamp{}  -> ConstantStorage v
  VAddress{}    -> ConstantStorage v
  VChest{}      -> ConstantStorage v
  VChestKey{}   -> ConstantStorage v
  VTxRollupL2Address{} -> ConstantStorage v
  VTicket{}     -> PartlyPushableStorage v Nop

  -- Non-atomic
  VOption (Nothing :: Maybe (Value tm)) -> case checkScope @(ConstantScope tm) of
    Right Dict -> ConstantStorage $ VOption Nothing
    Left _     -> PushableValueStorage NONE
  VOption (Just jVal :: Maybe (Value tm)) -> case splitPushableStorage jVal of
    ConstantStorage _ -> ConstantStorage . VOption $ Just jVal
    PushableValueStorage instr -> PushableValueStorage $ instr `Seq` SOME
    PartlyPushableStorage val instr -> PartlyPushableStorage val $ instr `Seq` SOME

  VList (vals :: [Value tl]) -> case checkScope @(ConstantScope tl) of
    Right Dict -> ConstantStorage v
    Left _     ->
      -- Here we check that even tho the type contains big_maps, we actually
      -- have big_maps in (one or more of) the values too.
      let handleList
            :: Instr s ('T.TList tl ': s) -> Value tl
            -> Maybe (Instr s ('T.TList tl ': s))
          handleList instr ele = case splitPushableStorage ele of
            ConstantStorage val ->
              Just $ instr `Seq` PUSH val `Seq` CONS
            PushableValueStorage eleInstr ->
              Just $ instr `Seq` eleInstr `Seq` CONS
            PartlyPushableStorage _ _ -> Nothing
      in  maybe (PartlyPushableStorage v Nop) PushableValueStorage $
            foldM handleList NIL vals

  VSet{} -> ConstantStorage v

  VPair (v1 :: Value t1, v2 :: Value t2) ->
    withValueTypeSanity v1 $ withValueTypeSanity v2 $
      withDeMorganScope @StorageScope @'T.TPair @t1 @t2 $
        let handlePair
              :: PushableStorageSplit s t2
              -> PushableStorageSplit (t2 ': s) t1
              -> PushableStorageSplit s ('T.TPair t1 t2)
            handlePair psp2 psp1 = case (psp2, psp1) of
              -- at least one side is a constant
              (ConstantStorage _, ConstantStorage _) ->
                ConstantStorage v
              (ConstantStorage val2, _) ->
                handlePair (PushableValueStorage $ PUSH val2) psp1
              (_, ConstantStorage val1) ->
                handlePair psp2 (PushableValueStorage $ PUSH val1)
              -- at least one side is a constant or has no big_map values
              (PushableValueStorage instr2, PushableValueStorage instr1) ->
                PushableValueStorage $ instr2 `Seq` instr1 `Seq` PAIR
              (PushableValueStorage instr2, PartlyPushableStorage val1 instr1) ->
                PartlyPushableStorage val1 $ DIP instr2 `Seq` instr1 `Seq` PAIR
              (PartlyPushableStorage val2 instr2, PushableValueStorage instr1) ->
                PartlyPushableStorage val2 $ instr2 `Seq` instr1 `Seq` PAIR
              -- both sides contain a big_map
              (PartlyPushableStorage val2 instr2, PartlyPushableStorage val1 instr1) ->
                PartlyPushableStorage (VPair (val1, val2)) $
                  UNPAIR `Seq` DIP instr2 `Seq` instr1 `Seq` PAIR
        in handlePair (splitPushableStorage v2) (splitPushableStorage v1)

  VOr (Left orVal :: Either (Value t1) (Value t2)) ->
    withValueTypeSanity orVal $ withDeMorganScope @StorageScope @'T.TOr @t1 @t2 $
      case splitPushableStorage orVal of
        ConstantStorage val -> case checkScope @(ConstantScope t2) of
          -- note: here we need to check for the opposite branch too
          Right Dict -> ConstantStorage v
          Left _ -> PushableValueStorage $ PUSH val `Seq` LEFT
        PushableValueStorage instr -> PushableValueStorage $ instr `Seq` LEFT
        PartlyPushableStorage val instr -> PartlyPushableStorage val $ instr `Seq` LEFT
  VOr (Right orVal :: Either (Value t1) (Value t2)) ->
    withValueTypeSanity orVal $ withDeMorganScope @StorageScope @'T.TOr @t1 @t2 $
      case splitPushableStorage orVal of
        ConstantStorage val -> case checkScope @(ConstantScope t1) of
          -- note: here we need to check for the opposite branch too
          Right Dict -> ConstantStorage v
          Left _ -> PushableValueStorage $ PUSH val `Seq` RIGHT
        PushableValueStorage instr -> PushableValueStorage $ instr `Seq` RIGHT
        PartlyPushableStorage val instr -> PartlyPushableStorage val $ instr `Seq` RIGHT

  VMap (vMap :: (Map (Value tk) (Value tv))) -> case checkScope @(ConstantScope tk) of
    Left _ ->
      -- NOTE: all keys for a map need to be comparable and (even tho it's
      -- not a superclass) that means they have to be constants.
      -- I (@pasqu4le) found no exception to this rule, perhaps it should be
      -- enforced in the type system as well, but it may have more
      -- downsides than it's worth accepting.
      error "impossible: all map keys should be PUSHable"
    Right Dict -> case checkScope @(ConstantScope tv) of
      Right Dict -> ConstantStorage v
      _ -> withDeMorganScope @HasNoOp @'T.TMap @tk @tv $
        -- Similarly as for lists, here we check that even tho the value type
        -- contains a big_map, we actually have big_maps in (one or more of) them.
        let handleMap
              :: Instr s ('T.TMap tk tv ': s) -> (Value tk, Value tv)
              -> Maybe (Instr s ('T.TMap tk tv ': s))
            handleMap instr (key, ele) = case splitPushableStorage (VOption $ Just ele) of
              ConstantStorage val ->
                Just $ instr `Seq` PUSH val `Seq` PUSH key `Seq` UPDATE
              PushableValueStorage eleInstr ->
                Just $ instr `Seq` eleInstr `Seq` PUSH key `Seq` UPDATE
              PartlyPushableStorage _ _ -> Nothing
        in  maybe (PartlyPushableStorage v Nop) PushableValueStorage $
              foldM handleMap EMPTY_MAP $ M.toList vMap

  VBigMap _ _ -> PartlyPushableStorage v Nop

-- | Whether this instruction is a real Michelson instruction.
--
-- Only the root is in question, children in the instruction tree are not
-- accounted for.
--
-- >>> isMichelsonInstr (Seq Nop Nop)
-- True
--
-- >>> isMichelsonInstr (Ext $ COMMENT_ITEM "comment")
-- False
--
-- This function is helpful e.g. in debugger.
isMichelsonInstr :: Instr i o -> Bool
isMichelsonInstr = \case
  WithLoc{} -> False
  Meta{} -> False
  FrameInstr{} -> False
  Seq{} -> True
  Nop -> False
  Ext{} -> False
  Nested{} -> True
  DocGroup{} -> False
  AnnCAR{} -> True
  AnnCDR{} -> True
  DROP{} -> True
  DROPN{} -> True
  AnnDUP{} -> True
  AnnDUPN{} -> True
  SWAP{} -> True
  DIG{} -> True
  DUG{} -> True
  AnnPUSH{} -> True
  AnnSOME{} -> True
  AnnNONE{} -> True
  AnnUNIT{} -> True
  IF_NONE{} -> True
  AnnPAIR{} -> True
  AnnUNPAIR{} -> True
  AnnPAIRN{} -> True
  UNPAIRN{} -> True
  AnnLEFT{} -> True
  AnnRIGHT{} -> True
  IF_LEFT{} -> True
  AnnNIL{} -> True
  AnnCONS{} -> True
  IF_CONS{} -> True
  AnnSIZE{} -> True
  AnnEMPTY_SET{} -> True
  AnnEMPTY_MAP{} -> True
  AnnEMPTY_BIG_MAP{} -> True
  AnnMAP{} -> True
  ITER{} -> True
  AnnMEM{} -> True
  AnnGET{} -> True
  AnnGETN{} -> True
  AnnUPDATE{} -> True
  AnnUPDATEN{} -> True
  AnnGET_AND_UPDATE{} -> True
  IF{} -> True
  LOOP{} -> True
  LOOP_LEFT{} -> True
  AnnLAMBDA{} -> True
  AnnEXEC{} -> True
  AnnAPPLY{} -> True
  DIP{} -> True
  DIPN{} -> True
  FAILWITH{} -> True
  AnnCAST{} -> True
  AnnRENAME{} -> True
  AnnPACK{} -> True
  AnnUNPACK{} -> True
  AnnCONCAT{} -> True
  AnnCONCAT'{} -> True
  AnnSLICE{} -> True
  AnnISNAT{} -> True
  AnnADD{} -> True
  AnnSUB{} -> True
  AnnSUB_MUTEZ{} -> True
  AnnMUL{} -> True
  AnnEDIV{} -> True
  AnnABS{} -> True
  AnnNEG{} -> True
  AnnLSL{} -> True
  AnnLSR{} -> True
  AnnOR{} -> True
  AnnAND{} -> True
  AnnXOR{} -> True
  AnnNOT{} -> True
  AnnCOMPARE{} -> True
  AnnEQ{} -> True
  AnnNEQ{} -> True
  AnnLT{} -> True
  AnnGT{} -> True
  AnnLE{} -> True
  AnnGE{} -> True
  AnnINT{} -> True
  AnnVIEW{} -> True
  AnnSELF{} -> True
  AnnCONTRACT{} -> True
  AnnTRANSFER_TOKENS{} -> True
  AnnSET_DELEGATE{} -> True
  AnnCREATE_CONTRACT{} -> True
  AnnIMPLICIT_ACCOUNT{} -> True
  AnnNOW{} -> True
  AnnAMOUNT{} -> True
  AnnBALANCE{} -> True
  AnnVOTING_POWER{} -> True
  AnnTOTAL_VOTING_POWER{} -> True
  AnnCHECK_SIGNATURE{} -> True
  AnnSHA256{} -> True
  AnnSHA512{} -> True
  AnnBLAKE2B{} -> True
  AnnSHA3{} -> True
  AnnKECCAK{} -> True
  AnnHASH_KEY{} -> True
  AnnPAIRING_CHECK{} -> True
  AnnSOURCE{} -> True
  AnnSENDER{} -> True
  AnnADDRESS{} -> True
  AnnCHAIN_ID{} -> True
  AnnLEVEL{} -> True
  AnnSELF_ADDRESS{} -> True
  NEVER{} -> True
  AnnTICKET{} -> True
  AnnREAD_TICKET{} -> True
  AnnSPLIT_TICKET{} -> True
  AnnJOIN_TICKETS{} -> True
  AnnOPEN_CHEST{} -> True
  AnnSAPLING_EMPTY_STATE{} -> True
  AnnSAPLING_VERIFY_UPDATE{} -> True
  AnnMIN_BLOCK_TIME{} -> True
  AnnEMIT{} -> True

-- | A wrapper around either typechecked 'Anns' or unchecked 'NonEmpty' of
-- 'AnyAnn'. Annotations on some instructions aren't typechecked, hence these
-- two constructors.
--
-- Helper for 'instrAnns'.
data SomeAnns where
  SomeAnns :: Anns xs -> SomeAnns
  SomeUncheckedAnns :: NonEmpty AnyAnn -> SomeAnns

-- | Get annotations from a typed 'Instr'. This doesn't recurse, use with
-- 'dfsFoldInstr' to collect all annotations in a tree/sequence.
instrAnns :: Instr i o -> Maybe SomeAnns
instrAnns = \case
  WithLoc{} -> mzero
  Meta{} -> mzero
  FrameInstr{} -> mzero
  Seq{} -> mzero
  Nop -> mzero
  Ext{} -> mzero
  Nested{} -> mzero
  DocGroup{} -> mzero
  AnnCAR anns -> pure $ SomeAnns anns
  AnnCDR anns -> pure $ SomeAnns anns
  DROP -> mzero
  DROPN _ -> mzero
  AnnDUP anns -> pure $ SomeAnns anns
  AnnDUPN anns _ -> pure $ SomeAnns anns
  SWAP -> mzero
  DIG _ -> mzero
  DUG _ -> mzero
  AnnPUSH anns _ -> pure $ SomeAnns anns
  AnnSOME anns -> pure $ SomeAnns anns
  AnnNONE anns -> pure $ SomeAnns anns
  AnnUNIT anns -> pure $ SomeAnns anns
  IF_NONE _ _ -> mzero
  AnnPAIR anns -> pure $ SomeAnns anns
  AnnUNPAIR anns -> pure $ SomeAnns anns
  AnnPAIRN anns _ -> pure $ SomeAnns anns
  UNPAIRN _ -> mzero
  AnnLEFT anns -> pure $ SomeAnns anns
  AnnRIGHT anns -> pure $ SomeAnns anns
  IF_LEFT _ _ -> mzero
  AnnNIL anns -> pure $ SomeAnns anns
  AnnCONS anns -> pure $ SomeAnns anns
  IF_CONS _ _ -> mzero
  AnnSIZE anns -> pure $ SomeAnns anns
  AnnEMPTY_SET anns -> pure $ SomeAnns anns
  AnnEMPTY_MAP anns -> pure $ SomeAnns anns
  AnnEMPTY_BIG_MAP anns -> pure $ SomeAnns anns
  AnnMAP anns _ -> pure $ SomeAnns anns
  ITER _ -> mzero
  AnnMEM anns -> pure $ SomeAnns anns
  AnnGET anns -> pure $ SomeAnns anns
  AnnGETN anns _ -> pure $ SomeAnns anns
  AnnUPDATE anns -> pure $ SomeAnns anns
  AnnUPDATEN anns _ -> pure $ SomeAnns anns
  AnnGET_AND_UPDATE anns -> pure $ SomeAnns anns
  IF _ _ -> mzero
  LOOP _ -> mzero
  LOOP_LEFT _ -> mzero
  AnnLAMBDA anns _ -> pure $ SomeAnns anns
  AnnEXEC anns -> pure $ SomeAnns anns
  AnnAPPLY anns -> pure $ SomeAnns anns
  DIP _ -> mzero
  DIPN _ _ -> mzero
  FAILWITH -> mzero
  AnnCAST anns -> pure $ SomeAnns anns
  AnnRENAME anns -> pure $ SomeAnns anns
  AnnPACK anns -> pure $ SomeAnns anns
  AnnUNPACK anns -> pure $ SomeAnns anns
  AnnCONCAT anns -> pure $ SomeAnns anns
  AnnCONCAT' anns -> pure $ SomeAnns anns
  AnnSLICE anns -> pure $ SomeAnns anns
  AnnISNAT anns -> pure $ SomeAnns anns
  AnnADD anns -> pure $ SomeAnns anns
  AnnSUB anns -> pure $ SomeAnns anns
  AnnSUB_MUTEZ anns -> pure $ SomeAnns anns
  AnnMUL anns -> pure $ SomeAnns anns
  AnnEDIV anns -> pure $ SomeAnns anns
  AnnABS anns -> pure $ SomeAnns anns
  AnnNEG anns -> pure $ SomeAnns anns
  AnnLSL anns -> pure $ SomeAnns anns
  AnnLSR anns -> pure $ SomeAnns anns
  AnnOR anns -> pure $ SomeAnns anns
  AnnAND anns -> pure $ SomeAnns anns
  AnnXOR anns -> pure $ SomeAnns anns
  AnnNOT anns -> pure $ SomeAnns anns
  AnnCOMPARE anns -> pure $ SomeAnns anns
  AnnEQ anns -> pure $ SomeAnns anns
  AnnNEQ anns -> pure $ SomeAnns anns
  AnnLT anns -> pure $ SomeAnns anns
  AnnGT anns -> pure $ SomeAnns anns
  AnnLE anns -> pure $ SomeAnns anns
  AnnGE anns -> pure $ SomeAnns anns
  AnnINT anns -> pure $ SomeAnns anns
  AnnVIEW anns _ -> pure $ SomeAnns anns
  AnnSELF anns _ -> pure $ SomeAnns anns
  AnnCONTRACT anns _ -> pure $ SomeAnns anns
  AnnTRANSFER_TOKENS anns -> pure $ SomeAnns anns
  AnnSET_DELEGATE anns -> pure $ SomeAnns anns
  AnnCREATE_CONTRACT anns _ -> pure $ SomeAnns anns
  AnnIMPLICIT_ACCOUNT anns -> pure $ SomeAnns anns
  AnnNOW anns -> pure $ SomeAnns anns
  AnnAMOUNT anns -> pure $ SomeAnns anns
  AnnBALANCE anns -> pure $ SomeAnns anns
  AnnVOTING_POWER anns -> pure $ SomeAnns anns
  AnnTOTAL_VOTING_POWER anns -> pure $ SomeAnns anns
  AnnCHECK_SIGNATURE anns -> pure $ SomeAnns anns
  AnnSHA256 anns -> pure $ SomeAnns anns
  AnnSHA512 anns -> pure $ SomeAnns anns
  AnnBLAKE2B anns -> pure $ SomeAnns anns
  AnnSHA3 anns -> pure $ SomeAnns anns
  AnnKECCAK anns -> pure $ SomeAnns anns
  AnnHASH_KEY anns -> pure $ SomeAnns anns
  AnnPAIRING_CHECK anns -> pure $ SomeAnns anns
  AnnSOURCE anns -> pure $ SomeAnns anns
  AnnSENDER anns -> pure $ SomeAnns anns
  AnnADDRESS anns -> pure $ SomeAnns anns
  AnnCHAIN_ID anns -> pure $ SomeAnns anns
  AnnLEVEL anns -> pure $ SomeAnns anns
  AnnSELF_ADDRESS anns -> pure $ SomeAnns anns
  NEVER -> mzero
  AnnTICKET anns -> pure $ SomeAnns anns
  AnnREAD_TICKET anns -> pure $ SomeAnns anns
  AnnSPLIT_TICKET anns -> pure $ SomeAnns anns
  AnnJOIN_TICKETS anns -> pure $ SomeAnns anns
  AnnOPEN_CHEST anns -> pure $ SomeAnns anns
  AnnSAPLING_EMPTY_STATE anns _ -> pure $ SomeAnns anns
  AnnSAPLING_VERIFY_UPDATE anns -> pure $ SomeAnns anns
  AnnMIN_BLOCK_TIME anns -> SomeUncheckedAnns <$> nonEmpty anns
  AnnEMIT anns tag ty -> pure $
    case ty of
      Just ty' -> SomeAnns $ tag `AnnsCons` ty' `AnnsTyCons` anns
      Nothing -> SomeAnns $ tag `AnnsCons` anns