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

-- | Utilities for writing optimizer rules.
module Morley.Michelson.Optimizer.Utils
  ( pattern (:#)
  , orRule
  , orSimpleRule
  , fixpoint
  , applyOnce
  , whileApplies
  , linearizeAndReapply
  ) where

import Morley.Michelson.Optimizer.Internal.Ruleset (Rule(..))
import Morley.Michelson.Typed.Instr hiding ((:#))

{- | This is a redefinition of @(:#)@ from "Morley.Michelson.Typed.Instr" that
is particularly useful for writing optimizer rules. When matching on an
instruction @x@ that isn't v'Seq', it behaves as if it matched on @x :# Nop@.

When constructing instructions using this pattern, @Nop@ is automatically
removed.

To understand why this is useful, consider that a given instruction sequence can
appear in the middle of a sequence, and then @a :# b :# tail@ will match, or at
the end of the sequence, and then @a :# b@ will match. Thus, to cover all cases
one would have to duplicate most rules.

This definition of @(:#)@ makes it so we can always assume there's a @tail@.
However, we don't need it when matching on single instructions.

Thus, the rule of thumb is this: if you're matching on a single instruction,
everything is fine. If you're matching on a sequence, i.e. using (:#), then
always match on tail, e.g.

@
dupSwap2dup :: Rule
dupSwap2dup = Rule $ \case
  DUP :# SWAP :# c -> Just $ DUP :# c
  _                -> Nothing
@

But this works, too:

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

-}
pattern (:#) :: Instr inp b -> Instr b out -> Instr inp out
pattern l :# r <- (\case { x@Seq{} -> x; x -> Seq x Nop } -> Seq l r)
  where l :# Nop = l
        Nop :# r = r
        l :# r = Seq l r
infixr 8 :#

-- | Combine two rule fixpoints.
orRule :: (Rule -> Rule) -> (Rule -> Rule) -> (Rule -> Rule)
orRule l r topl = Rule $ \instr ->
  (unRule (l topl) $ instr) <|> (unRule (r topl) $ instr)

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

-- | 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.
--
-- Also returns a flag showing whether the rule succeeded or not.
applyOnce :: Rule -> Instr inp out -> (Any Bool, Instr inp out)
applyOnce r i = maybe (pure i) (Any True,) (unRule r $ i)

-- | Apply a rule to the same code, until it fails.
whileApplies :: Rule -> Rule
whileApplies r = Rule $ go <=< unRule r
  -- NB: if the rule doesn't apply even once, we want to return Nothing here,
  -- hence it's first applied above and only if successful goes into recursion.
  where
    go :: Instr inp out -> Maybe (Instr inp out)
    go i = maybe (Just i) go (unRule r i)

-- | Append LHS of v'Seq' to RHS and re-run pointwise ocRuleset at each point.
--   That might cause reinvocation of this function (see @defaultRule@'),
--   but effectively this ensures it will flatten any v'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 v'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 = snd . \case
  Seq (Seq a b) c ->
    applyOnce restart $ Seq a (linearizeAndReapply restart (Seq b c))

  other -> applyOnce restart other