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

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