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

-- | Optimizer for typed instructions.
--
-- It's quite experimental and incomplete.
module Morley.Michelson.Optimizer
  ( optimize
  , optimizeWithConf
  , optimizeVerboseWithConf
  , defaultOptimizerConf
  , defaultRules
  , defaultRulesAndPushPack
  , 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.Default (Default(def))
import Fmt (Buildable(..), (+|), (|+))

import Morley.Michelson.Optimizer.Internal.Rules
import Morley.Michelson.Optimizer.Internal.Ruleset
import Morley.Michelson.Optimizer.Utils
import Morley.Michelson.Typed.ClassifiedInstr
import Morley.Michelson.Typed.Instr hiding ((:#))
import Morley.Michelson.Typed.Util (DfsSettings(..), dfsFoldInstr, dfsTraverseInstr)

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

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

instance Default OptimizerConf where
  def :: OptimizerConf
def = OptimizerConf
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 :: forall (inp :: [T]) (out :: [T]). Instr inp out -> Instr inp out
optimize = OptimizerConf -> Instr inp out -> Instr inp out
forall (inp :: [T]) (out :: [T]).
OptimizerConf -> Instr inp out -> Instr inp out
optimizeWithConf OptimizerConf
forall a. Default a => a
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 :: forall (inp :: [T]) (out :: [T]).
OptimizerConf -> Instr inp out -> Instr inp out
optimizeWithConf = ([OptimizerStageStats], Instr inp out) -> Instr inp out
forall a b. (a, b) -> b
snd (([OptimizerStageStats], Instr inp out) -> Instr inp out)
-> (OptimizerConf
    -> Instr inp out -> ([OptimizerStageStats], Instr inp out))
-> OptimizerConf
-> Instr inp out
-> Instr inp out
forall a b c. SuperComposition a b c => a -> b -> c
... OptimizerConf
-> Instr inp out -> ([OptimizerStageStats], Instr inp out)
forall (inp :: [T]) (out :: [T]).
OptimizerConf
-> Instr inp out -> ([OptimizerStageStats], Instr inp out)
optimizeVerboseWithConf

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

instance Buildable OptimizerStageStats where
  build :: OptimizerStageStats -> Doc
build OptimizerStageStats{Word
OptimizationStage
ossStage :: OptimizerStageStats -> OptimizationStage
ossNumIterations :: OptimizerStageStats -> Word
ossNumInstrs :: OptimizerStageStats -> Word
ossStage :: OptimizationStage
ossNumIterations :: Word
ossNumInstrs :: Word
..} =
    Doc
"Stage " Doc -> Doc -> Doc
forall b. FromDoc b => Doc -> Doc -> b
+| OptimizationStage
ossStage
      OptimizationStage -> Doc -> Doc
forall a b. (Buildable a, FromDoc b) => a -> Doc -> b
|+ Doc
" finished after " Doc -> Doc -> Doc
forall b. FromDoc b => Doc -> Doc -> b
+| Word
ossNumIterations
      Word -> Doc -> Doc
forall a b. (Buildable a, FromDoc b) => a -> Doc -> b
|+ Doc
" iterations. Instruction count: " Doc -> Doc -> Doc
forall b. FromDoc b => Doc -> Doc -> b
+| Word
ossNumInstrs Word -> Doc -> Doc
forall a b. (Buildable a, FromDoc b) => a -> Doc -> b
|+ Doc
""

-- | 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 :: forall (inp :: [T]) (out :: [T]).
OptimizerConf
-> Instr inp out -> ([OptimizerStageStats], Instr inp out)
optimizeVerboseWithConf OptimizerConf{Bool
Word
Ruleset
ocGotoValues :: OptimizerConf -> Bool
ocRuleset :: OptimizerConf -> Ruleset
ocMaxIterations :: OptimizerConf -> Word
ocGotoValues :: Bool
ocRuleset :: Ruleset
ocMaxIterations :: Word
..} Instr inp out
instr =
  (Instr inp out
 -> (OptimizationStage, NonEmpty Rule)
 -> ([OptimizerStageStats], Instr inp out))
-> Instr inp out
-> [(OptimizationStage, NonEmpty Rule)]
-> ([OptimizerStageStats], Instr inp out)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM (Word
-> Instr inp out
-> (OptimizationStage, NonEmpty Rule)
-> ([OptimizerStageStats], Instr inp out)
performOneStage Word
1) Instr inp out
instrRHS ([(OptimizationStage, NonEmpty Rule)]
 -> ([OptimizerStageStats], Instr inp out))
-> [(OptimizationStage, NonEmpty Rule)]
-> ([OptimizerStageStats], Instr inp out)
forall a b. (a -> b) -> a -> b
$ Map OptimizationStage (NonEmpty Rule)
-> [(Key (Map OptimizationStage (NonEmpty Rule)),
     Val (Map OptimizationStage (NonEmpty Rule)))]
forall t. ToPairs t => t -> [(Key t, Val t)]
toPairs (Map OptimizationStage (NonEmpty Rule)
 -> [(Key (Map OptimizationStage (NonEmpty Rule)),
      Val (Map OptimizationStage (NonEmpty Rule)))])
-> Map OptimizationStage (NonEmpty Rule)
-> [(Key (Map OptimizationStage (NonEmpty Rule)),
     Val (Map OptimizationStage (NonEmpty Rule)))]
forall a b. (a -> b) -> a -> b
$ Ruleset -> Map OptimizationStage (NonEmpty Rule)
unRuleset Ruleset
ocRuleset
  where
    performOneStage :: Word
-> Instr inp out
-> (OptimizationStage, NonEmpty Rule)
-> ([OptimizerStageStats], Instr inp out)
performOneStage Word
n Instr inp out
i stage :: (OptimizationStage, NonEmpty Rule)
stage@(OptimizationStage
stageName, NonEmpty Rule
stageRules)
      | Bool -> Bool
forall a. Boolean a => a -> a
not Bool
changed Bool -> Bool -> Bool
forall a. Boolean a => a -> a -> a
|| Word
n Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Word
ocMaxIterations = ([OptimizationStage -> Word -> Word -> OptimizerStageStats
OptimizerStageStats OptimizationStage
stageName Word
n Word
stats], Instr inp out
res)
      | Bool
otherwise = Word
-> Instr inp out
-> (OptimizationStage, NonEmpty Rule)
-> ([OptimizerStageStats], Instr inp out)
performOneStage (Word -> Word
forall a. Enum a => a -> a
succ Word
n) Instr inp out
res (OptimizationStage, NonEmpty Rule)
stage
      where
        stats :: Word
stats = Sum Word -> Word
forall a. Sum a -> a
getSum (Sum Word -> Word) -> Sum Word -> Word
forall a b. (a -> b) -> a -> b
$ Instr inp out -> Sum Word
instrCount Instr inp out
res
        instrCount :: Instr inp out -> Sum Word
instrCount = DfsSettings (Writer (Sum Word))
-> (forall (i :: [T]) (o :: [T]). Instr i o -> Sum Word)
-> Instr inp out
-> Sum Word
forall x (inp :: [T]) (out :: [T]).
Monoid x =>
DfsSettings (Writer x)
-> (forall (i :: [T]) (o :: [T]). Instr i o -> x)
-> Instr inp out
-> x
dfsFoldInstr DfsSettings (Writer (Sum Word))
forall a. Default a => a
def { dsGoToValues :: Bool
dsGoToValues = Bool
ocGotoValues } ((forall (i :: [T]) (o :: [T]). Instr i o -> Sum Word)
 -> Instr inp out -> Sum Word)
-> (forall (i :: [T]) (o :: [T]). Instr i o -> Sum Word)
-> Instr inp out
-> Sum Word
forall a b. (a -> b) -> a -> b
$ (forall (cls :: InstrClass).
 (SingI cls, WCIConstraint Instr cls) =>
 Sing (GetClassified cls) -> ClassifiedInstr cls i o -> Sum Word)
-> Instr i o -> Sum Word
forall t (inp :: [T]) (out :: [T]) r.
ClassifyInstr t =>
(forall (cls :: InstrClass).
 (SingI cls, WCIConstraint Instr cls) =>
 Sing (GetClassified cls) -> ClassifiedInstr cls inp out -> r)
-> Instr inp out -> r
forall (instr :: [T] -> [T] -> *) t (inp :: [T]) (out :: [T]) r.
(WithClassifiedInstr instr, ClassifyInstr t) =>
(forall (cls :: InstrClass).
 (SingI cls, WCIConstraint instr cls) =>
 Sing (GetClassified cls) -> ClassifiedInstr cls inp out -> r)
-> instr inp out -> r
withClassifiedInstr \case
          Sing (GetClassified cls)
SingIsMichelson (GetClassified cls)
SFromMichelson -> Sum Word -> ClassifiedInstr cls i o -> Sum Word
forall a b. a -> b -> a
const Sum Word
1
          Sing (GetClassified cls)
_ -> Sum Word -> ClassifiedInstr cls i o -> Sum Word
forall a b. a -> b -> a
const Sum Word
0
        stageRule :: Rule
stageRule = (Rule -> Rule) -> Rule
fixpoint ((Rule -> Rule) -> Rule) -> (Rule -> Rule) -> Rule
forall a b. (a -> b) -> a -> b
$ ((Rule -> Rule) -> Element (NonEmpty Rule) -> Rule -> Rule)
-> (Rule -> Rule) -> NonEmpty Rule -> Rule -> Rule
forall t b. Container t => (b -> Element t -> b) -> b -> t -> b
forall b.
(b -> Element (NonEmpty Rule) -> b) -> b -> NonEmpty Rule -> b
foldl (Rule -> Rule) -> Element (NonEmpty Rule) -> Rule -> Rule
(Rule -> Rule) -> Rule -> Rule -> Rule
orSimpleRule Rule -> Rule
flattenSeqLHS NonEmpty Rule
stageRules
        (Any Bool -> Bool
forall a. Any a -> a
getAny -> Bool
changed, Instr inp out
res) =
          DfsSettings ((,) (Any Bool))
-> Instr inp out -> (Any Bool, Instr inp out)
forall (m :: * -> *) (inp :: [T]) (out :: [T]).
Monad m =>
DfsSettings m -> Instr inp out -> m (Instr inp out)
dfsTraverseInstr DfsSettings ((,) (Any Bool))
forall a. Default a => a
def{ dsGoToValues :: Bool
dsGoToValues = Bool
ocGotoValues, dsInstrStep :: forall (i :: [T]) (o :: [T]). Instr i o -> (Any Bool, Instr i o)
dsInstrStep = Rule -> Instr i o -> (Any Bool, Instr i o)
forall (inp :: [T]) (out :: [T]).
Rule -> Instr inp out -> (Any Bool, Instr inp out)
applyOnce Rule
stageRule } Instr inp out
i
    instrRHS :: Instr inp out
instrRHS = (Any Bool, Instr inp out) -> Instr inp out
forall a b. (a, b) -> b
snd ((Any Bool, Instr inp out) -> Instr inp out)
-> (Any Bool, Instr inp out) -> Instr inp out
forall a b. (a -> b) -> a -> b
$ Rule -> Instr inp out -> (Any Bool, Instr inp out)
forall (inp :: [T]) (out :: [T]).
Rule -> Instr inp out -> (Any Bool, Instr inp out)
applyOnce ((Rule -> Rule) -> Rule
fixpoint Rule -> Rule
flattenSeqLHS) Instr inp out
instr

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