-- 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)

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.Peano (SingNat(..))

----------------------------------------------------------------------------
-- 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` 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` specificPush
    `orSimpleRule` pairUnpair
    `orSimpleRule` unpairMisc
    `orSimpleRule` swapBeforeCommutative

-- | 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

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

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

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

variousNops :: Rule
variousNops = \case
  Seq  SWAP    (Seq SWAP c) -> Just c
  Seq (PUSH _) (Seq DROP c) -> Just c
  Seq  DUP     (Seq DROP c) -> Just c
  Seq  UNIT    (Seq DROP c) -> Just c
  Seq  NOW     (Seq DROP c) -> Just c
  Seq  SENDER  (Seq DROP c) -> Just c
  Seq  EMPTY_MAP (Seq DROP c) -> Just c
  Seq  EMPTY_SET (Seq DROP c) -> Just c

  Seq  SWAP         SWAP    -> Just Nop
  Seq (PUSH _)      DROP    -> Just Nop
  Seq  DUP          DROP    -> Just Nop
  Seq  UNIT         DROP    -> Just Nop
  Seq  NOW          DROP    -> Just Nop
  Seq  SENDER       DROP    -> Just Nop
  Seq  EMPTY_MAP   DROP    -> Just Nop
  Seq  EMPTY_SET   DROP    -> Just Nop
  _                         -> Nothing

dupSwap2dup :: Rule
dupSwap2dup = \case
  Seq DUP (Seq SWAP c) -> Just (Seq DUP c)
  Seq 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.
  Seq (PUSH x) (Seq (DIP f) c) -> Just (Seq f (Seq (PUSH x) c))
  Seq (PUSH x)      (DIP f)    -> Just (Seq f      (PUSH x))
  Seq UNIT     (Seq (DIP f) c) -> Just (Seq f (Seq UNIT c))
  Seq UNIT          (DIP f)    -> Just (Seq f      UNIT)
  Seq NOW     (Seq (DIP f) c) -> Just (Seq f (Seq NOW c))
  Seq NOW          (DIP f)    -> Just (Seq f      NOW)
  Seq SENDER  (Seq (DIP f) c) -> Just (Seq f (Seq SENDER c))
  Seq SENDER       (DIP f)    -> Just (Seq f      SENDER)
  Seq EMPTY_MAP     (Seq (DIP f) c) -> Just (Seq f (Seq EMPTY_MAP c))
  Seq EMPTY_MAP          (DIP f)    -> Just (Seq f      EMPTY_MAP)
  Seq EMPTY_SET     (Seq (DIP f) c) -> Just (Seq f (Seq EMPTY_SET c))
  Seq EMPTY_SET          (DIP f)    -> Just (Seq 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.
  Seq (DIP f) (Seq DROP c)     -> Just (Seq DROP (Seq f c))
  Seq (DIP f) DROP             -> Just (Seq DROP f)

  _                            -> Nothing

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

  Seq (PUSH (VBool True))  (Seq (IF f _) c) -> Just (Seq f c)
  Seq (PUSH (VBool False)) (Seq (IF _ f) c) -> Just (Seq f c)

  Seq LEFT       (IF_LEFT f _)    -> Just f
  Seq RIGHT      (IF_LEFT _ f)    -> Just f
  Seq CONS       (IF_CONS f _)    -> Just f
  Seq NIL        (IF_CONS _ f)    -> Just f
  Seq NONE       (IF_NONE f _)    -> Just f
  Seq SOME       (IF_NONE _ f)    -> Just f

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

  _ -> Nothing

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

simpleDrops :: Rule
simpleDrops = \case
  -- DROP 0 is Nop
  Seq (DROPN SZ) c -> Just c
  DROPN SZ -> 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.

  -- Seq (DROPN (SS SZ)) c -> Just (Seq DROP c)
  -- DROPN (SS SZ) -> 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
  Seq (DIP Nop) c -> Just c

  _ -> Nothing

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

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

  _ -> Nothing

adjacentDips :: Rule
adjacentDips = \case
  Seq (DIP f) (DIP g) -> Just (DIP (Seq f g))
  Seq (DIP f) (Seq (DIP g) c) -> Just (DIP (Seq f g) `Seq` 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
  Seq DROP DROP -> Just (DROPN (SS $ SS SZ))
  Seq DROP (Seq DROP c) -> Just (DROPN (SS $ SS SZ) `Seq` c)

  -- Does not compile, need to do something smart
  -- Seq (DROPN (SS (SS SZ))) DROP -> Just (DROPN (SS $ SS $ SS SZ))

  _ -> Nothing

specificPush :: Rule
specificPush = \case
  push@PUSH{} -> optimizePush push
  Seq (push@PUSH{}) c -> (`Seq` 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

-- UNPAIR with continuation
pattern UNPAIR_c
  :: ()
  => (i ~ ('TPair a b : s), i' ~ (a : b : s), o' ~ o)
  => Instr i' o' -> Instr i o
pattern UNPAIR_c c = Seq DUP (Seq CAR (Seq (DIP CDR) c))

pairUnpair :: Rule
pairUnpair = \case
  Seq PAIR (UNPAIR_c c) -> Just c
  Seq PAIR UNPAIR       -> Just Nop

  UNPAIR_c (Seq PAIR c) -> Just c
  UNPAIR_c      PAIR    -> Just Nop

  _ -> Nothing

unpairMisc :: Rule
unpairMisc = \case
  Seq UNPAIR SWAP -> Just (DUP `Seq` CDR `Seq` DIP CAR)
  Seq UNPAIR (Seq SWAP c) -> Just (DUP `Seq` CDR `Seq` DIP CAR `Seq` 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
  Seq SWAP (Seq i c) -> (`Seq` c) <$> commuteArith i
  Seq SWAP i -> commuteArith i

  _ -> Nothing

pushPack :: Rule
pushPack = \case
  Seq (PUSH x) PACK -> Just (pushPacked x)
  Seq (PUSH x) (Seq PACK c) -> Just (pushPacked x `Seq` c)

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

-- | 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