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

{-# OPTIONS_HADDOCK not-home #-}

-- | Optimizer rule and ruleset definitions.
module Morley.Michelson.Optimizer.Internal.Ruleset
  ( module Morley.Michelson.Optimizer.Internal.Ruleset
  ) where

import Prelude

import Data.Default (Default(def))
import Data.Map qualified as Map
import Fmt (Buildable(..), (+|), (|+))

import Morley.Michelson.Typed.Instr

-- | Type of a single rewrite rule, wrapped in `newtype`. 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`.
newtype Rule = Rule {unRule :: forall inp out. Instr inp out -> Maybe (Instr inp out)}

-- | Optimization stages. Stages are run in first to last order, each stage has
-- an 'Int' argument, which allows splitting each stage into sub-stages, which
-- will run lowest index to highest. All default rules use sub-stage @0@.
data OptimizationStage
  = OptimizationStagePrepare Int
  | OptimizationStageMain Int
    -- ^ Main optimisation stage, except rules that would interfere with other
    -- rules.
  | OptimizationStageMainExtended Int
    -- ^ All main stage rules.
  | OptimizationStageFixup Int
    -- ^ Post main stage fixups.
  | OptimizationStageRollAdjacent Int
    -- ^ Main stage rules unroll @DROP n@, @PAIR n@, etc into their primitive
    -- counterparts to simplify some optimisations. This stages coalesces them
    -- back.
  deriving stock (Eq, Ord)

instance Buildable OptimizationStage where
  build = \case
    OptimizationStagePrepare n -> "prepare " +| n |+ ""
    OptimizationStageMain n -> "main " +| n |+ ""
    OptimizationStageMainExtended n -> "main extended " +| n |+ ""
    OptimizationStageFixup n -> "fixup " +| n |+ ""
    OptimizationStageRollAdjacent n -> "roll adjacent " +| n |+ ""

-- | A set of optimization stages. Rules at the same sub-stage are applied in
-- arbitrary order. See 'OptimizationStage' for explanation of sub-stages.
--
-- 'Default' ruleset is empty.
newtype Ruleset = Ruleset { unRuleset :: Map OptimizationStage (NonEmpty Rule) }
  deriving newtype Default

instance Semigroup Ruleset where
  Ruleset l <> Ruleset r = Ruleset $ Map.unionWith (<>) l r

instance Monoid Ruleset where
  mempty = def

-- | Get rules for a given priority as a list.
rulesAtPrio :: OptimizationStage -> Ruleset -> [Rule]
rulesAtPrio prio = maybe [] toList . Map.lookup prio . unRuleset

-- | Insert a single rule at a given priority without touching other rules.
insertRuleAtPrio :: OptimizationStage -> Rule -> Ruleset -> Ruleset
insertRuleAtPrio = flip $ alterRulesAtPrio . (:)

-- | Remove the stage with the given priority.
clearRulesAtPrio :: OptimizationStage -> Ruleset -> Ruleset
clearRulesAtPrio = alterRulesAtPrio (const [])

-- | Alter all stage rules for a given priority.
alterRulesAtPrio :: ([Rule] -> [Rule]) -> OptimizationStage -> Ruleset -> Ruleset
alterRulesAtPrio f prio = Ruleset . Map.alter (nonEmpty . f . maybe [] toList) prio . unRuleset