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

-- NOTE this pragmas.
-- We disable some wargnings 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.
-- List of possible improvements:
-- 1. 'pushDrop', 'dupDrop', 'unitDrop' rules are essentially the
-- same. It would be good to generalize them into one rule. The same
-- applies to 'pushDip'.
-- It probably can be done more efficiently.

module Michelson.Optimizer
  ( optimize
  , optimizeWithConf
  , defaultOptimizerConf
  , defaultRules
  , defaultRulesAndPushPack
  , orRule
  , orSimpleRule
  , Rule
  , OptimizerConf (..)
  , ocGotoValuesL
  ) where

import Prelude hiding (EQ, LT, GT)

import Control.Lens (makeLensesFor)
import Data.Constraint (Dict(..))
import Data.Default (Default(def))
import Data.Singletons (sing)

import Michelson.Interpret.Pack (packValue')
import Michelson.Typed.Aliases (Value)
import Michelson.Typed.Arith
import Michelson.Typed.Instr
import Michelson.Typed.Scope (PackedValScope)
import Michelson.Typed.Sing
import Michelson.Typed.T
import Michelson.Typed.Util (DfsSettings(..), dfsInstr)
import Michelson.Typed.Value
import Util.PeanoNatural

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

data OptimizerConf = OptimizerConf
  { ocGotoValues :: Bool
  , ocRuleset    :: Rule -> Rule
  }

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

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.
optimizeWithConf :: OptimizerConf -> Instr inp out -> Instr inp out
optimizeWithConf (OptimizerConf ocGotoValues rules)
  = (fst .)
  $ dfsInstr dfsSettings
  $ (adapter .)
  $ applyOnce
  $ fixpoint rules
  where
    dfsSettings = def{ dsGoToValues = ocGotoValues }

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

-- Type of a single rewrite rule. 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`.
type Rule = forall inp out. Instr inp out -> Maybe (Instr inp out)

defaultRules :: Rule -> Rule
defaultRules =
  flattenSeqLHS
    `orSimpleRule` removeNesting
    `orSimpleRule` removeExtStackType
    `orSimpleRule` flattenFn
    `orSimpleRule` dipDrop2swapDrop
    `orSimpleRule` ifNopNop2Drop
    `orSimpleRule` nopIsNeutralForSeq
    `orSimpleRule` variousNops
    `orSimpleRule` dupSwap2dup
    `orSimpleRule` noDipNeeded
    `orSimpleRule` branchShortCut
    `orSimpleRule` compareWithZero
    `orSimpleRule` simpleDrops
    `orSimpleRule` internalNop
    `orSimpleRule` simpleDips
    `orSimpleRule` adjacentDips
    `orSimpleRule` adjacentDrops
    `orSimpleRule` isSomeOnIf
    `orSimpleRule` specificPush
    `orSimpleRule` pairUnpair
    `orSimpleRule` unpairMisc
    `orSimpleRule` swapBeforeCommutative
    `orSimpleRule` justDrops

-- | 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 :: Rule -> Rule
defaultRulesAndPushPack = defaultRules `orSimpleRule` pushPack

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

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

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

flattenFn :: Rule
flattenFn = \case
  Fn _ _ i -> Just i
  _        -> Nothing

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

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

nopIsNeutralForSeq :: Rule
nopIsNeutralForSeq = \case
  Nop :# i   -> Just i
  i   :# Nop -> Just i
  _          -> Nothing

variousNops :: Rule
variousNops = \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

  DUP                :# DROP -> Just Nop
  DUPN _             :# DROP -> Just Nop
  SWAP               :# SWAP -> Just Nop
  PUSH _             :# DROP -> Just Nop
  NONE               :# DROP -> Just Nop
  UNIT               :# DROP -> Just Nop
  NIL                :# DROP -> Just Nop
  EMPTY_SET          :# DROP -> Just Nop
  EMPTY_MAP          :# DROP -> Just Nop
  EMPTY_BIG_MAP      :# DROP -> Just Nop
  LAMBDA _           :# DROP -> Just Nop
  SELF _             :# DROP -> Just Nop
  NOW                :# DROP -> Just Nop
  AMOUNT             :# DROP -> Just Nop
  BALANCE            :# DROP -> Just Nop
  TOTAL_VOTING_POWER :# DROP -> Just Nop
  SOURCE             :# DROP -> Just Nop
  SENDER             :# DROP -> Just Nop
  CHAIN_ID           :# DROP -> Just Nop
  LEVEL              :# DROP -> Just Nop
  SELF_ADDRESS       :# DROP -> Just Nop
  READ_TICKET        :# DROP -> Just Nop
  _                          -> Nothing

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

noDipNeeded :: Rule
noDipNeeded = \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
  PUSH x    :# DIP f      -> Just $ f :# PUSH x
  UNIT      :# DIP f :# c -> Just $ f :# UNIT :# c
  UNIT      :# DIP f      -> Just $ f :# UNIT
  NOW       :# DIP f :# c -> Just $ f :# NOW :# c
  NOW       :# DIP f      -> Just $ f :# NOW
  SENDER    :# DIP f :# c -> Just $ f :# SENDER :# c
  SENDER    :# DIP f      -> Just $ f :# SENDER
  EMPTY_MAP :# DIP f :# c -> Just $ f :# EMPTY_MAP :# c
  EMPTY_MAP :# DIP f      -> Just $ f :# EMPTY_MAP
  EMPTY_SET :# DIP f :# c -> Just $ f :# EMPTY_SET :# c
  EMPTY_SET :# DIP f      -> Just $ f :# EMPTY_SET

  -- 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
  DIP f :# DROP         -> Just $ DROP :# f

  _ -> Nothing

branchShortCut :: Rule
branchShortCut = \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 (VBool True)  :# IF f _ :# c -> Just (f :# c)
  PUSH (VBool False) :# IF _ f :# c -> Just (f :# c)

  LEFT  :# IF_LEFT f _ -> Just f
  RIGHT :# IF_LEFT _ f -> Just f
  CONS  :# IF_CONS f _ -> Just f
  NIL   :# IF_CONS _ f -> Just f
  NONE  :# IF_NONE f _ -> Just f
  SOME  :# IF_NONE _ f -> Just f

  PUSH (VBool True)  :# IF f _ -> Just f
  PUSH (VBool False) :# IF _ f -> Just f

  _ -> Nothing

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

simpleDrops :: Rule
simpleDrops = \case
  -- DROP 0 is Nop
  DROPN Zero :# c -> Just c
  DROPN Zero -> Just Nop

  -- DROP 1 is DROP.
  -- @gromak: DROP seems to be cheaper (in my experiments it consumed 3 less gas).
  -- It is packed more efficiently.
  -- Unfortunately I do not know how to convince GHC that types match here.
  -- Specifically, it can not deduce that `inp` is not empty
  -- (`DROP` expects non-empty input).
  -- We have `LongerOrSameLength inp (S Z)` here, but that is not enough to
  -- convince GHC.
  -- I will leave this note and rule here in hope that someone will manage to
  -- deal with this problem one day.

  -- DROPN One :# c -> Just (DROP :# c)
  -- DROPN One -> Just DROP

  _ -> 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 = \case
  DIP Nop      -> Just Nop
  DIP Nop :# c -> Just c

  _ -> Nothing

simpleDips :: Rule
simpleDips = \case
  -- DIP 0 is redundant
  DIPN Zero i :# c -> Just (i :# c)
  DIPN Zero i -> Just i

  -- @gromak: same situation as with `DROP 1` (see above).
  -- DIPN One i :# c -> Just (DIP i :# c)
  -- DIPN One i -> Just (DIP i)

  _ -> Nothing

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

  _ -> Nothing

-- TODO (#299): optimize sequences of more than 2 DROPs.
-- | 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 = \case
  DROP :# DROP -> Just (DROPN Two)
  DROP :# DROP :# c -> Just (DROPN Two :# c)

  -- Does not compile, need to do something smart
  -- DROPN Two :# DROP -> Just (DROPN (Succ Two))

  _ -> Nothing

specificPush :: Rule
specificPush = \case
  push@PUSH{}         -> optimizePush push
  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 = \case
  IF (PUSH (VOption Just{})) (PUSH (VOption Nothing)) :# c -> case c of
    IF_NONE (PUSH (VBool False)) (DROP :# PUSH (VBool True)) :# s -> Just s
    IF_NONE (PUSH (VBool False)) (DROP :# PUSH (VBool True))      -> Just Nop
    _ -> Nothing
  _ -> Nothing

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

  UNPAIR :# PAIR :# c -> Just c
  UNPAIR :# PAIR      -> Just Nop

  _ -> Nothing

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

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

  UNPAIR :# DROP             -> Just CDR
  UNPAIR :# DROP :# c        -> Just $ CDR :# 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 = \case
  SWAP :# i :# c -> (:# c) <$> commuteArith i
  SWAP :# i -> commuteArith i

  _ -> Nothing

pushPack :: Rule
pushPack = \case
  PUSH x :# PACK -> Just (pushPacked x)
  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 = \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

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

-- | Append LHS of 'Seq' to RHS and re-run pointwise ocRuleset at each point.
--   That might cause reinvocation of this function (see 'defaultRules'),
--   but productivity ensures it will flatten any '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 '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 = \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 x = l topl x <|> r topl x

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

-- | 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.
applyOnce :: Rule -> Instr inp out -> Instr inp out
applyOnce r i = maybe i id (r i)

-- | An adapter for `dfsInstr`.
adapter :: a -> (a, ())
adapter a = (a, ())

-- | Apply a rule to the same code, until it fails.
whileApplies :: Rule -> Rule
whileApplies r = go
  where
    go i = maybe (Just i) go (r i)

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

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