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

{-# OPTIONS_HADDOCK not-home #-}

-- | Optimizer rewrite rules
module Morley.Michelson.Optimizer.Internal.Rules
  ( module Morley.Michelson.Optimizer.Internal.Rules
  ) where

import Prelude hiding (EQ, GT, LT)

import Data.Constraint (Dict(..), (\\))
import Data.Default (Default(def))
import Data.Singletons (sing)
import Data.Type.Equality ((:~:)(Refl))

import Morley.Michelson.Interpret.Pack (packValue')
import Morley.Michelson.Optimizer.Internal.Proofs
import Morley.Michelson.Optimizer.Internal.Ruleset
import Morley.Michelson.Optimizer.Utils
import Morley.Michelson.Typed.Aliases (Value)
import Morley.Michelson.Typed.Arith
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.Value
import Morley.Util.Peano
import Morley.Util.PeanoNatural

{-
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.
-}

-- | 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)
  , (fixupStageRules             , OptimizationStageFixup 0)
  , (glueAdjacentInstrsStageRules, OptimizationStageRollAdjacent 0)
  ]
  where
    dipDrop2swapDropStageRules = dipDrop2swapDrop : mainStageRules
    fixupStageRules = [dipSwapDrop]
    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
  , dupDugDrop
  -- 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

-- | NB: This rule MUST be applied AFTER all main stages, but BEFORE nested dips
-- are rolled up
dipSwapDrop :: Rule
dipSwapDrop = Rule \case
  DIP (SWAP :# DROP) -> Just $ DIPN Two DROP
  _ -> Nothing

dupDugDrop :: Rule
dupDugDrop = Rule \case
  DUP :# DUG (Succ m) :# DROP :# xs -> Just $ DUG m :# xs
  _ -> Nothing