Optimizer.hs

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

-- NOTE this pragmas.
-- We disable some warnings for the sake of speed up.
-- Write code with care.
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}
{-# OPTIONS_GHC -Wno-overlapping-patterns #-}

-- | Optimizer for typed instructions.
--
-- It's quite experimental and incomplete.

module Morley.Michelson.Optimizer
  ( optimize
  , optimizeWithConf
  , optimizeVerboseWithConf
  , defaultOptimizerConf
  , defaultRules
  , defaultRulesAndPushPack
  , orRule
  , orSimpleRule
  , Rule (..)
  , OptimizerConf (..)
  , ocGotoValuesL
    -- * Ruleset manipulation
  , OptimizationStage (..)
  , Ruleset
  , rulesAtPrio
  , insertRuleAtPrio
  , clearRulesAtPrio
  , alterRulesAtPrio
  , OptimizerStageStats(..)
  ) where

import Prelude hiding (EQ, GT, LT)

import Control.Lens (makeLensesFor)
import Data.Constraint (Dict(..), (\\))
import Data.Default (Default(def))
import Data.Map qualified as Map
import Data.Singletons (sing)
import Data.Type.Equality ((:~:)(Refl))
import Fmt (Buildable(..), (+|), (|+))

import Morley.Michelson.Interpret.Pack (packValue')
import Morley.Michelson.Optimizer.Internal.Proofs
import Morley.Michelson.Typed.Aliases (Value)
import Morley.Michelson.Typed.Arith
import Morley.Michelson.Typed.ClassifiedInstr
import Morley.Michelson.Typed.Instr hiding ((:#))
import Morley.Michelson.Typed.Scope (ConstantScope, PackedValScope, checkScope)
import Morley.Michelson.Typed.Sing
import Morley.Michelson.Typed.T
import Morley.Michelson.Typed.Util (DfsSettings(..), dfsFoldInstr, dfsTraverseInstr)
import Morley.Michelson.Typed.Value
import Morley.Util.Peano
import Morley.Util.PeanoNatural

pattern (:#) :: Instr inp b -> Instr b out -> Instr inp out
pattern l :# r <- (maybeSeq -> Seq l r)
  where l :# Nop = l
        Nop :# r = r
        l :# r = Seq l r
infixr 8 :#

maybeSeq :: Instr inp out -> Instr inp out
maybeSeq = \case
  x@Seq{} -> x
  x -> Seq x Nop

----------------------------------------------------------------------------
-- High level
----------------------------------------------------------------------------

-- | Optimization stages. Stages are run in first to last order, each stage has
-- an 'Int' argument, which allows splitting each stage into sub-stages, which
-- will run lowest index to highest. All default rules use sub-stage @0@.
data OptimizationStage
  = OptimizationStagePrepare Int
  | OptimizationStageMain Int
    -- ^ Main optimisation stage, except rules that would interfere with other
    -- rules.
  | OptimizationStageMainExtended Int
    -- ^ All main stage rules.
  | OptimizationStageRollAdjacent Int
    -- ^ Main stage rules unroll @DROP n@, @PAIR n@, etc into their primitive
    -- counterparts to simplify some optimisations. This stages coalesces them
    -- back.
  deriving stock (Eq, Ord)

instance Buildable OptimizationStage where
  build = \case
    OptimizationStagePrepare n -> "prepare " +| n |+ ""
    OptimizationStageMain n -> "main " +| n |+ ""
    OptimizationStageMainExtended n -> "main extended " +| n |+ ""
    OptimizationStageRollAdjacent n -> "roll adjacent " +| n |+ ""

-- | A set of optimization stages. Rules at the same sub-stage are applied in
-- arbitrary order. See 'OptimizationStage' for explanation of sub-stages.
--
-- 'Default' ruleset is empty.
newtype Ruleset = Ruleset { unRuleset :: Map OptimizationStage (NonEmpty Rule) }
  deriving newtype Default

instance Semigroup Ruleset where
  Ruleset l <> Ruleset r = Ruleset $ Map.unionWith (<>) l r

instance Monoid Ruleset where
  mempty = def

-- | Get rules for a given priority as a list.
rulesAtPrio :: OptimizationStage -> Ruleset -> [Rule]
rulesAtPrio prio = maybe [] toList . Map.lookup prio . unRuleset

-- | Insert a single rule at a given priority without touching other rules.
insertRuleAtPrio :: OptimizationStage -> Rule -> Ruleset -> Ruleset
insertRuleAtPrio = flip $ alterRulesAtPrio . (:)

-- | Remove the stage with the given priority.
clearRulesAtPrio :: OptimizationStage -> Ruleset -> Ruleset
clearRulesAtPrio = alterRulesAtPrio (const [])

-- | Alter all stage rules for a given priority.
alterRulesAtPrio :: ([Rule] -> [Rule]) -> OptimizationStage -> Ruleset -> Ruleset
alterRulesAtPrio f prio = Ruleset . Map.alter (nonEmpty . f . maybe [] toList) prio . unRuleset

data OptimizerConf = OptimizerConf
  { ocGotoValues :: Bool
  , ocRuleset    :: Ruleset
  , ocMaxIterations :: Word
  }

-- | Default config - all commonly useful rules will be applied to all the code.
defaultOptimizerConf :: OptimizerConf
defaultOptimizerConf = OptimizerConf
  { ocGotoValues = True
  , ocRuleset    = defaultRules
  , ocMaxIterations = 100
  }

instance Default OptimizerConf where
  def = defaultOptimizerConf

-- | Optimize a typed instruction by replacing some sequences of
-- instructions with smaller equivalent sequences.
-- Applies default set of rewrite rules.
optimize :: Instr inp out -> Instr inp out
optimize = optimizeWithConf def

-- | Optimize a typed instruction using a custom set of rules.
-- The set is divided into several stages, as applying
-- some rules can prevent others to be performed.
--
-- If any stage resulted in optimizations, we apply it again until we reach
-- fixpoint, but no more than 'ocMaxIterations' times.
optimizeWithConf :: OptimizerConf -> Instr inp out -> Instr inp out
optimizeWithConf = snd ... optimizeVerboseWithConf

data OptimizerStageStats = OptimizerStageStats
  { ossStage :: OptimizationStage
  , ossNumIterations :: Word
  , ossNumInstrs :: Word
  }

instance Buildable OptimizerStageStats where
  build OptimizerStageStats{..} =
    "Stage " +| ossStage
      |+ " finished after " +| ossNumIterations
      |+ " iterations. Instruction count: " +| ossNumInstrs |+ ""

-- | Returns some optimizer statistics in addition to optimized instruction.
-- Mostly useful for testing and debugging.
optimizeVerboseWithConf
  :: OptimizerConf
  -> Instr inp out
  -> ([OptimizerStageStats], Instr inp out)
optimizeVerboseWithConf OptimizerConf{..} instr =
  foldlM (performOneStage 1) instrRHS $ toPairs $ unRuleset ocRuleset
  where
    performOneStage n i stage@(stageName, stageRules)
      | not changed || n >= ocMaxIterations = ([OptimizerStageStats stageName n stats], res)
      | otherwise = performOneStage (succ n) res stage
      where
        stats = getSum $ instrCount res
        instrCount = dfsFoldInstr def { dsGoToValues = ocGotoValues } $ withClassifiedInstr \case
          SFromMichelson -> const 1
          _ -> const 0
        stageRule = fixpoint $ foldl orSimpleRule flattenSeqLHS stageRules
        (getAny -> changed, res) =
          dfsTraverseInstr def{ dsGoToValues = ocGotoValues, dsInstrStep = applyOnce stageRule } i
    instrRHS = snd $ applyOnce (fixpoint flattenSeqLHS) instr

----------------------------------------------------------------------------
-- Rewrite rules
----------------------------------------------------------------------------

{-
Note [Writing optimizer rules]
------------------------------

We locally redefine (:#) to simplify handling instruction sequence tails.
Consider that a given instruction sequence can appear in the middle, and then @a
:# b :# tail@ will match, or at the end of the sequence, and then @a :# b@ will
match.

Local definition of (:#) makes it so we can always assume there's a tail.
However, we don't need it when matching on single instructions.

Thus, the rule of thumb is this: if you're matching on a single instruction,
everything is fine. If you're matching on a sequence, i.e. using (:#), then
always match on tail.
-}

-- Type of a single rewrite rule, wrapped in `newtype`. It takes an instruction
-- and tries to optimize its head (first few instructions). If optimization
-- succeeds, it returns `Just` the optimized instruction, otherwise it returns `Nothing`.
newtype Rule = Rule {unRule :: forall inp out. Instr inp out -> Maybe (Instr inp out)}

-- | Default optimization rules.
defaultRules :: Ruleset
defaultRules = foldr ($) def $ uncurry (alterRulesAtPrio . const) <$>
  -- NB: if adding more main stages, remember to check if
  -- 'defaultRulesAndPushPack' needs updating.
  [ (mainStageRules              , OptimizationStageMain 0)
  , (dipDrop2swapDropStageRules  , OptimizationStageMainExtended 0)
  , (glueAdjacentInstrsStageRules, OptimizationStageRollAdjacent 0)
  ]
  where
    dipDrop2swapDropStageRules = dipDrop2swapDrop : mainStageRules
    glueAdjacentInstrsStageRules =
      [ adjacentDrops
      , rollPairN
      , rollUnpairN
      , rollDips
      ]

-- | We do not enable 'pushPack' rule by default because it is potentially
-- dangerous. There are various code processing functions that may depend on
-- constants, e. g. string transformations.
defaultRulesAndPushPack :: Ruleset
defaultRulesAndPushPack = defaultRules
  & insertRuleAtPrio (OptimizationStageMainExtended 0) pushPack

mainStageRules :: [Rule]
mainStageRules =
  [ removeNesting
  , removeExtStackType
  , ifNopNop2Drop
  , nopIsNeutralForSeq
  , variousNops
  , dupSwap2dup
  , noDipNeeded
  , branchShortCut
  , compareWithZero
  , internalNop
  , simpleDups
  , adjacentDips
  , isSomeOnIf
  , redundantIf
  , emptyDip
  , digDug
  , specificPush
  , pairUnpair
  , pairMisc
  , unpairMisc
  , swapBeforeCommutative
  , justDrops
  , justDoubleDrops
  , dig1AndDug1AreSwap
  , notIf
  , dropMeta
  -- strictly speaking, these rules below are de-optimisations, but they
  -- expose opportunities for other optimisations, and they are undone in the
  -- last stage.
  , unrollPairN
  , unrollUnpairN
  , unrollDropN
  , unrollDips
  ]

flattenSeqLHS :: Rule -> Rule
flattenSeqLHS toplevel = Rule $ \case
  it@(Seq (Seq _ _) _) -> Just $ linearizeAndReapply toplevel it
  _                    -> Nothing

dropMeta :: Rule
dropMeta = Rule $ \case
  Meta _ i -> Just i
  WithLoc _ i -> Just i
  _        -> Nothing

removeNesting :: Rule
removeNesting = Rule $ \case
  Nested i -> Just i
  _        -> Nothing

-- | STACKTYPE is currently a Nop and may safely be removed.
removeExtStackType :: Rule
removeExtStackType = Rule $ \case
  Ext (STACKTYPE{}) -> Just Nop
  _                 -> Nothing

dipDrop2swapDrop :: Rule
dipDrop2swapDrop = Rule $ \case
  DIP DROP -> Just $ SWAP :# DROP
  _        -> Nothing

ifNopNop2Drop :: Rule
ifNopNop2Drop = Rule $ \case
  IF Nop Nop -> Just DROP
  _          -> Nothing

nopIsNeutralForSeq :: Rule
nopIsNeutralForSeq = Rule $ \case
  -- NB: we're not using (:#) here because it'll always match when rhs is Nop
  Seq Nop i -> Just i
  Seq i Nop -> Just i
  _         -> Nothing

variousNops :: Rule
variousNops = Rule $ \case
  DUP                :# DROP :# c -> Just c
  DUPN _             :# DROP :# c -> Just c
  SWAP               :# SWAP :# c -> Just c
  PUSH _             :# DROP :# c -> Just c
  NONE               :# DROP :# c -> Just c
  UNIT               :# DROP :# c -> Just c
  NIL                :# DROP :# c -> Just c
  EMPTY_SET          :# DROP :# c -> Just c
  EMPTY_MAP          :# DROP :# c -> Just c
  EMPTY_BIG_MAP      :# DROP :# c -> Just c
  LAMBDA _           :# DROP :# c -> Just c
  SELF _             :# DROP :# c -> Just c
  NOW                :# DROP :# c -> Just c
  AMOUNT             :# DROP :# c -> Just c
  BALANCE            :# DROP :# c -> Just c
  TOTAL_VOTING_POWER :# DROP :# c -> Just c
  SOURCE             :# DROP :# c -> Just c
  SENDER             :# DROP :# c -> Just c
  CHAIN_ID           :# DROP :# c -> Just c
  LEVEL              :# DROP :# c -> Just c
  SELF_ADDRESS       :# DROP :# c -> Just c
  READ_TICKET        :# DROP :# c -> Just c
  _                          -> Nothing

dupSwap2dup :: Rule
dupSwap2dup = Rule $ \case
  DUP :# SWAP :# c -> Just $ DUP :# c
  _                -> Nothing

noDipNeeded :: Rule
noDipNeeded = Rule $ \case
  -- If we put a constant value on stack and then do something under it,
  -- we can do this "something" on original stack and then put that constant.
  PUSH x    :# DIP f :# c -> Just $ f :# PUSH x :# c
  UNIT      :# DIP f :# c -> Just $ f :# UNIT :# c
  NOW       :# DIP f :# c -> Just $ f :# NOW :# c
  SENDER    :# DIP f :# c -> Just $ f :# SENDER :# c
  EMPTY_MAP :# DIP f :# c -> Just $ f :# EMPTY_MAP :# c
  EMPTY_SET :# DIP f :# c -> Just $ f :# EMPTY_SET :# c

  -- If we do something ignoring top of the stack and then immediately
  -- drop top of the stack, we can drop that item in advance and
  -- not use 'DIP' at all.
  DIP f :# DROP :# c -> Just $ DROP :# f :# c

  _ -> Nothing

unrollDips :: Rule
unrollDips = Rule go
  where
    go :: Instr inp out -> Maybe (Instr inp out)
    go = \case
      DIPN Zero c -> Just $ c
      DIPN One c -> Just $ DIP c
      DIPN (Succ n) c -> DIP <$> go (DIPN n c)
      _ -> Nothing

rollDips :: Rule
rollDips = Rule $ \case
  DIP (DIP c) -> Just $ DIPN Two c
  DIP (DIPN n c) -> Just $ DIPN (Succ n) c
  _ -> Nothing

branchShortCut :: Rule
branchShortCut = Rule $ \case
  LEFT  :# IF_LEFT f _ :# c -> Just $ f :# c
  RIGHT :# IF_LEFT _ f :# c -> Just $ f :# c
  CONS  :# IF_CONS f _ :# c -> Just $ f :# c
  NIL   :# IF_CONS _ f :# c -> Just $ f :# c
  NONE  :# IF_NONE f _ :# c -> Just $ f :# c
  SOME  :# IF_NONE _ f :# c -> Just $ f :# c

  PUSH vOr@(VOr eitherVal) :# IF_LEFT f g :# c -> case vOr of
    (_ :: Value ('TOr l r)) -> case eitherVal of
      Left val -> case checkScope @(ConstantScope l) of
        Right Dict -> Just $ PUSH val :# f :# c
        _          -> Nothing
      Right val -> case checkScope @(ConstantScope r) of
        Right Dict -> Just $ PUSH val :# g :# c
        _          -> Nothing
  PUSH (VList (x : xs))     :# IF_CONS f _ :# c -> Just $ PUSH (VList xs) :# PUSH x :# f :# c
  PUSH (VList _)            :# IF_CONS _ f :# c -> Just $ f :# c
  PUSH (VOption Nothing)    :# IF_NONE f _ :# c -> Just $ f :# c
  PUSH (VOption (Just val)) :# IF_NONE _ f :# c -> Just $ PUSH val :# f :# c
  PUSH (VBool True)         :# IF f _      :# c -> Just $ f :# c
  PUSH (VBool False)        :# IF _ f      :# c -> Just $ f :# c
  _ -> Nothing

compareWithZero :: Rule
compareWithZero = Rule $ \case
  PUSH (VInt 0) :# COMPARE :# EQ :# c -> Just $ EQ :# c
  PUSH (VNat 0) :# COMPARE :# EQ :# c -> Just $ INT :# EQ :# c
  _                                   -> Nothing

-- If an instruction takes another instruction as an argument and that
-- internal instruction is 'Nop', sometimes the whole instruction is
-- 'Nop'.
-- For now we do it only for 'DIP', but ideally we should do it for
-- 'MAP' as well (which is harder).
internalNop :: Rule
internalNop = Rule $ \case
  DIP Nop      -> Just Nop

  _ -> Nothing

simpleDups :: Rule
simpleDups = Rule $ \case
  -- DUP 1 is just DUP
  DUPN One -> Just DUP

  _ -> Nothing

adjacentDips :: Rule
adjacentDips = Rule $ \case
  DIP f :# DIP g :# c -> Just $ DIP (f :# g) :# c

  _ -> Nothing

redundantIf :: Rule
redundantIf = Rule \case
  IF x y
    | x == y
    -> Just $ DROP :# x
  _ -> Nothing

emptyDip :: Rule
emptyDip = Rule \case
  DIP Nop -> Just Nop
  DIPN _ Nop -> Just Nop
  _ -> Nothing

digDug :: Rule
digDug = Rule \case
  DIG x :# DUG y :# c | Just Refl <- eqPeanoNat x y -> Just c
  _ -> Nothing

-- | Sequences of @DROP@s can be turned into single @DROP n@.
-- When @n@ is greater than 2 it saves size and gas.
-- When @n@ is 2 it saves gas only.
adjacentDrops :: Rule
adjacentDrops = Rule $ \case
  DROP :# DROP :# c -> Just $ DROPN Two :# c

  (DROPN n :: Instr inp out) :# (DROP :: Instr inp' out') :# c
     -> Just $ DROPN (Succ n) :# c \\ dropNDropNProof @inp @out @out' n One
                                   \\ commutativity (singPeanoNat n) (singPeanoNat One)

  (DROP :: Instr inp out) :# (DROPN n :: Instr inp' out') :# c
     -> Just $ DROPN (Succ n) :# c \\ dropNDropNProof @inp @out @out' One n

  (DROPN n :: Instr inp out) :# (DROPN m :: Instr inp' out') :# c
     -> Just $ DROPN (addPeanoNat n m) :# c \\ dropNDropNProof @inp @out @out' n m

  _ -> Nothing

specificPush :: Rule
specificPush = Rule $ \case
  push@PUSH{} :# c -> (:# c) <$> optimizePush push

  _ -> Nothing
  where
    optimizePush :: Instr inp out -> Maybe (Instr inp out)
    optimizePush = \case
      PUSH v | _ :: Value v <- v -> case v of
        VUnit -> Just UNIT
        VMap m
          | null m -> case sing @v of STMap{} -> Just EMPTY_MAP
        VSet m
          | null m -> case sing @v of STSet{} -> Just EMPTY_SET
        _ -> Nothing

      _ -> Nothing


isSomeOnIf :: Rule
isSomeOnIf = Rule $ \case
  IF (PUSH (VOption Just{})) (PUSH (VOption Nothing))
    :# IF_NONE (PUSH (VBool False)) (DROP :# PUSH (VBool True))
    :# c
    -> Just c
  _ -> Nothing

pairUnpair :: Rule
pairUnpair = Rule $ \case
  PAIR :# UNPAIR :# c -> Just c

  UNPAIR :# PAIR :# c -> Just c

  _ -> Nothing

pairMisc :: Rule
pairMisc = Rule $ \case
  PAIR :# CDR :# c -> Just $ DROP :# c

  PAIR :# CAR :# c -> Just $ (DIP DROP) :# c

  _ -> Nothing

unpairMisc :: Rule
unpairMisc = Rule $ \case
  DUP :# CAR :# DIP CDR :# c -> Just $ UNPAIR :# c

  DUP :# CDR :# DIP CAR :# c -> Just $ UNPAIR :# SWAP :# c

  UNPAIR :# DROP :# c        -> Just $ CDR :# c
  _ -> Nothing

unrollDropN :: Rule
unrollDropN = Rule go
  where
    go :: forall inp out. Instr inp out -> Maybe (Instr inp out)
    go = \case
      DROPN Zero -> Just Nop
      DROPN (Succ n) -> do
        dropn <- go $ DROPN n
        Just $ DROP :# dropn
        \\ unconsListProof @inp n
      _ -> Nothing

unrollPairN :: Rule
unrollPairN = Rule go
  where
    go :: forall inp out. Instr inp out -> Maybe (Instr inp out)
    go = \case
      PAIRN Two -> Just PAIR \\ pairN2isPairProof @inp
      PAIRN (Succ n@(Succ Succ{}) ) -> do
        pairn <- go $ PAIRN n
        Just $ DIP pairn :# PAIR
        \\ unconsListProof @inp n
      _ -> Nothing

unrollUnpairN :: Rule
unrollUnpairN = Rule go
  where
    go :: forall inp out. Instr inp out -> Maybe (Instr inp out)
    go = \case
      UNPAIRN Two -> Just UNPAIR \\ unpairN2isUnpairProof @inp @out
      UNPAIRN (Succ n@(Succ Succ{})) -> do
        unpairn <- go $ UNPAIRN n
        Just $ UNPAIR :# DIP unpairn
        \\ unpairNisUnpairDipUnpairNProof @inp @out n
      _ -> Nothing

rollPairN :: Rule
rollPairN = Rule $ \case
  DIP PAIR :# PAIR :# c -> Just $ PAIRN (Succ Two) :# c

  (DIP (PAIRN n@(Succ Succ{})) :: Instr inp out) :# PAIR :# c -> Just $ PAIRN (Succ n) :# c
    \\ dipPairNPairIsPairNProof @inp n

  _ -> Nothing

rollUnpairN :: Rule
rollUnpairN = Rule $ \case
  UNPAIR :# DIP UNPAIR :# c -> Just $ UNPAIRN (Succ Two) :# c
  UNPAIR :# DIP (UNPAIRN n@(Succ Succ{})) :# c -> Just $ UNPAIRN (Succ n) :# c
  _ -> Nothing

commuteArith ::
  forall n m s out. Instr (n ': m ': s) out -> Maybe (Instr (m ': n ': s) out)
commuteArith = \case
  ADD -> do Dict <- commutativityProof @Add @n @m; Just ADD
  MUL -> do Dict <- commutativityProof @Mul @n @m; Just MUL
  OR -> do Dict <- commutativityProof @Or @n @m; Just OR
  AND -> do Dict <- commutativityProof @And @n @m; Just AND
  XOR -> do Dict <- commutativityProof @Xor @n @m; Just XOR
  _ -> Nothing

swapBeforeCommutative :: Rule
swapBeforeCommutative = Rule $ \case
  SWAP :# i :# c -> (:# c) <$> commuteArith i

  _ -> Nothing

pushPack :: Rule
pushPack = Rule $ \case
  PUSH x :# PACK :# c -> Just $ pushPacked x :# c

  _ -> Nothing
  where
    pushPacked :: PackedValScope t => Value t -> Instr s ('TBytes ': s)
    pushPacked = PUSH . VBytes . packValue'

justDrops :: Rule
justDrops = Rule $ \case
  CAR              :# DROP :# c -> Just $ DROP :# c
  CDR              :# DROP :# c -> Just $ DROP :# c
  SOME             :# DROP :# c -> Just $ DROP :# c
  LEFT             :# DROP :# c -> Just $ DROP :# c
  RIGHT            :# DROP :# c -> Just $ DROP :# c
  SIZE             :# DROP :# c -> Just $ DROP :# c
  GETN _           :# DROP :# c -> Just $ DROP :# c
  CAST             :# DROP :# c -> Just $ DROP :# c
  RENAME           :# DROP :# c -> Just $ DROP :# c
  PACK             :# DROP :# c -> Just $ DROP :# c
  UNPACK           :# DROP :# c -> Just $ DROP :# c
  CONCAT'          :# DROP :# c -> Just $ DROP :# c
  ISNAT            :# DROP :# c -> Just $ DROP :# c
  ABS              :# DROP :# c -> Just $ DROP :# c
  NEG              :# DROP :# c -> Just $ DROP :# c
  NOT              :# DROP :# c -> Just $ DROP :# c
  EQ               :# DROP :# c -> Just $ DROP :# c
  NEQ              :# DROP :# c -> Just $ DROP :# c
  LT               :# DROP :# c -> Just $ DROP :# c
  GT               :# DROP :# c -> Just $ DROP :# c
  LE               :# DROP :# c -> Just $ DROP :# c
  GE               :# DROP :# c -> Just $ DROP :# c
  INT              :# DROP :# c -> Just $ DROP :# c
  CONTRACT _       :# DROP :# c -> Just $ DROP :# c
  SET_DELEGATE     :# DROP :# c -> Just $ DROP :# c
  IMPLICIT_ACCOUNT :# DROP :# c -> Just $ DROP :# c
  VOTING_POWER     :# DROP :# c -> Just $ DROP :# c
  SHA256           :# DROP :# c -> Just $ DROP :# c
  SHA512           :# DROP :# c -> Just $ DROP :# c
  BLAKE2B          :# DROP :# c -> Just $ DROP :# c
  SHA3             :# DROP :# c -> Just $ DROP :# c
  KECCAK           :# DROP :# c -> Just $ DROP :# c
  HASH_KEY         :# DROP :# c -> Just $ DROP :# c
  PAIRING_CHECK    :# DROP :# c -> Just $ DROP :# c
  ADDRESS          :# DROP :# c -> Just $ DROP :# c
  JOIN_TICKETS     :# DROP :# c -> Just $ DROP :# c
  _                             -> Nothing

justDoubleDrops :: Rule
justDoubleDrops = Rule $ \case
  PAIR         :# DROP :# c -> Just $ DROP :# DROP :# c
  MEM          :# DROP :# c -> Just $ DROP :# DROP :# c
  GET          :# DROP :# c -> Just $ DROP :# DROP :# c
  APPLY        :# DROP :# c -> Just $ DROP :# DROP :# c
  CONCAT       :# DROP :# c -> Just $ DROP :# DROP :# c
  ADD          :# DROP :# c -> Just $ DROP :# DROP :# c
  SUB          :# DROP :# c -> Just $ DROP :# DROP :# c
  SUB_MUTEZ    :# DROP :# c -> Just $ DROP :# DROP :# c
  MUL          :# DROP :# c -> Just $ DROP :# DROP :# c
  EDIV         :# DROP :# c -> Just $ DROP :# DROP :# c
  OR           :# DROP :# c -> Just $ DROP :# DROP :# c
  AND          :# DROP :# c -> Just $ DROP :# DROP :# c
  XOR          :# DROP :# c -> Just $ DROP :# DROP :# c
  COMPARE      :# DROP :# c -> Just $ DROP :# DROP :# c
  TICKET       :# DROP :# c -> Just $ DROP :# DROP :# c
  SPLIT_TICKET :# DROP :# c -> Just $ DROP :# DROP :# c
  SWAP :# DROP :# DROP :# c -> Just $ DROP :# DROP :# c
  _ -> Nothing

dig1AndDug1AreSwap :: Rule
dig1AndDug1AreSwap = Rule \case
  DUG One -> Just SWAP
  DIG One -> Just SWAP
  _ -> Nothing

notIf :: Rule
notIf = Rule \case
  (NOT :: Instr inp out) :# IF a b :# c | STBool <- sing @(Head inp) -> Just $ IF b a :# c
  _ -> Nothing

-- | Append LHS of v'Seq' to RHS and re-run pointwise ocRuleset at each point.
--   That might cause reinvocation of this function (see 'defaultRules'),
--   but effectively this ensures it will flatten any v'Seq'-tree right-to-left,
--   while evaling no more than once on each node.
--
--   The reason this function invokes ocRuleset is when you append an instr
--   to already-optimised RHS of v'Seq', you might get an optimisable tree.
--
--   The argument is a local, non-structurally-recursive ocRuleset.
linearizeAndReapply :: Rule -> Instr inp out -> Instr inp out
linearizeAndReapply restart = snd . \case
  Seq (Seq a b) c ->
    applyOnce restart $ Seq a (linearizeAndReapply restart (Seq b c))

  other -> applyOnce restart other

----------------------------------------------------------------------------
-- Generic functions working with rules
----------------------------------------------------------------------------

-- | Combine two rule fixpoints.
orRule :: (Rule -> Rule) -> (Rule -> Rule) -> (Rule -> Rule)
orRule l r topl = Rule $ \instr ->
  (unRule (l topl) $ instr) <|> (unRule (r topl) $ instr)

-- | Combine a rule fixpoint and a simple rule.
orSimpleRule :: (Rule -> Rule) -> Rule -> (Rule -> Rule)
orSimpleRule l r topl = Rule $ \instr ->
  (unRule (l topl) $ instr) <|> (unRule r $ instr)

-- | Turn rule fixpoint into rule.
fixpoint :: (Rule -> Rule) -> Rule
fixpoint r = go
  where
    go :: Rule
    go = whileApplies (r go)

-- | Apply the rule once, if it fails, return the instruction unmodified.
--
-- Also returns a flag showing whether the rule succeeded or not.
applyOnce :: Rule -> Instr inp out -> (Any Bool, Instr inp out)
applyOnce r i = maybe (pure i) (Any True,) (unRule r $ i)

-- | Apply a rule to the same code, until it fails.
whileApplies :: Rule -> Rule
whileApplies r = Rule $ go <=< unRule r
  -- NB: if the rule doesn't apply even once, we want to return Nothing here,
  -- hence it's first applied above and only if successful goes into recursion.
  where
    go :: Instr inp out -> Maybe (Instr inp out)
    go i = maybe (Just i) go (unRule r i)

----------------------------------------------------------------------------
-- TH
----------------------------------------------------------------------------

makeLensesFor [("ocGotoValues", "ocGotoValuesL")] ''OptimizerConf