{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE BlockArguments #-}
{-|
  Given an input program 𝑝, equality saturation constructs an e-graph 𝐸 that
  represents a large set of programs equivalent to 𝑝, and then extracts the
  “best” program from 𝐸.

  The e-graph is grown by repeatedly applying pattern-based rewrites.
  Critically, these rewrites only add information to the e-graph, eliminating
  the need for careful ordering.

  Upon reaching a fixed point (saturation), 𝐸 will represent all equivalent
  ways to express 𝑝 with respect to the given rewrites.

  After saturation (or timeout), a final extraction procedure analyzes 𝐸 and
  selects the optimal program according to a user-provided cost function.
 -}
module Data.Equality.Saturation
    (
      -- * Equality saturation
      equalitySaturation, equalitySaturation'

      -- * Re-exports for equality saturation

      -- ** Writing rewrite rules
    , Rewrite(..), RewriteCondition

      -- ** Writing cost functions
      --
      -- | 'CostFunction' re-exported from 'Data.Equality.Extraction' since they are required to do equality saturation
    , CostFunction --, Cost, depthCost

      -- ** Writing expressions
      -- 
      -- | Expressions must be written in their fixed-point form, since the
      -- 'Language' must be given in its base functor form
    , Fix(..), cata

    ) where

import qualified Data.IntMap.Strict as IM

import Data.Bifunctor
import Control.Monad

import Data.Proxy

import Data.Equality.Utils
import qualified Data.Equality.Graph as G
import Data.Equality.Graph.Monad
import Data.Equality.Language
import Data.Equality.Graph.Classes
import Data.Equality.Matching
import Data.Equality.Matching.Database
import Data.Equality.Extraction

import Data.Equality.Saturation.Rewrites
import Data.Equality.Saturation.Scheduler

-- | Equality saturation with defaults
equalitySaturation :: forall l. Language l
                   => Fix l             -- ^ Expression to run equality saturation on
                   -> [Rewrite l]       -- ^ List of rewrite rules
                   -> CostFunction l    -- ^ Cost function to extract the best equivalent representation
                   -> (Fix l, EGraph l) -- ^ Best equivalent expression and resulting e-graph
equalitySaturation :: forall (l :: * -> *).
Language l =>
Fix l -> [Rewrite l] -> CostFunction l -> (Fix l, EGraph l)
equalitySaturation = Proxy BackoffScheduler
-> Fix l -> [Rewrite l] -> CostFunction l -> (Fix l, EGraph l)
forall (l :: * -> *) schd.
(Language l, Scheduler schd) =>
Proxy schd
-> Fix l -> [Rewrite l] -> CostFunction l -> (Fix l, EGraph l)
equalitySaturation' (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @BackoffScheduler)


-- | Run equality saturation on an expression given a list of rewrites, and
-- extract the best equivalent expression according to the given cost function
--
-- This variant takes all arguments instead of using defaults
equalitySaturation' :: forall l schd
                    . (Language l, Scheduler schd)
                    => Proxy schd        -- ^ Proxy for the scheduler to use
                    -> Fix l             -- ^ Expression to run equality saturation on
                    -> [Rewrite l]       -- ^ List of rewrite rules
                    -> CostFunction l    -- ^ Cost function to extract the best equivalent representation
                    -> (Fix l, EGraph l) -- ^ Best equivalent expression and resulting e-graph
equalitySaturation' :: forall (l :: * -> *) schd.
(Language l, Scheduler schd) =>
Proxy schd
-> Fix l -> [Rewrite l] -> CostFunction l -> (Fix l, EGraph l)
equalitySaturation' Proxy schd
_ Fix l
expr [Rewrite l]
rewrites CostFunction l
cost = EGraphM l (Fix l) -> (Fix l, EGraph l)
forall (l :: * -> *) a. Language l => EGraphM l a -> (a, EGraph l)
egraph (EGraphM l (Fix l) -> (Fix l, EGraph l))
-> EGraphM l (Fix l) -> (Fix l, EGraph l)
forall a b. (a -> b) -> a -> b
$ do

    -- Represent expression as an e-graph
    Int
origClass <- Fix l -> EGraphM l Int
forall (l :: * -> *). Language l => Fix l -> EGraphM l Int
represent Fix l
expr

    -- Run equality saturation (by applying non-destructively all rewrites)
    Int -> IntMap (Stat schd) -> EGraphM l ()
equalitySaturation'' Int
0 IntMap (Stat schd)
forall a. Monoid a => a
mempty -- Start at iteration 0

    -- Extract best solution from the e-class of the original expression
    (EGraph l -> Fix l) -> EGraphM l (Fix l)
forall (m :: * -> *) s a. Monad m => (s -> a) -> StateT s m a
gets ((EGraph l -> Fix l) -> EGraphM l (Fix l))
-> (EGraph l -> Fix l) -> EGraphM l (Fix l)
forall a b. (a -> b) -> a -> b
$ \EGraph l
g -> EGraph l -> CostFunction l -> Int -> Fix l
forall (lang :: * -> *).
Language lang =>
EGraph lang -> CostFunction lang -> Int -> Fix lang
extractBest EGraph l
g CostFunction l
cost Int
origClass

      where

        -- Take map each rewrite rule to stats on its usage so we can do
        -- backoff scheduling. Each rewrite rule is assigned an integer
        -- (corresponding to its position in the list of rewrite rules)
        equalitySaturation'' :: Int -> IM.IntMap (Stat schd) -> EGraphM l ()
        equalitySaturation'' :: Int -> IntMap (Stat schd) -> EGraphM l ()
equalitySaturation'' Int
30 IntMap (Stat schd)
_ = () -> EGraphM l ()
forall a. a -> StateT (EGraph l) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return () -- Stop after X iterations
        equalitySaturation'' Int
i IntMap (Stat schd)
stats = do

            egr :: EGraph l
egr@G.EGraph{ memo :: forall (l :: * -> *). EGraph l -> Memo l
G.memo = Memo l
beforeMemo, classes :: forall (l :: * -> *). EGraph l -> ClassIdMap (EClass l)
G.classes = ClassIdMap (EClass l)
beforeClasses } <- StateT (EGraph l) Identity (EGraph l)
forall (m :: * -> *) s. Monad m => StateT s m s
get

            let db :: Database l
db = EGraph l -> Database l
forall (l :: * -> *). Language l => EGraph l -> Database l
eGraphToDatabase EGraph l
egr

            -- Read-only phase, invariants are preserved
            -- With backoff scheduler
            -- ROMES:TODO parMap with chunks
            let (![(Rewrite l, Match)]
matches, IntMap (Stat schd)
newStats) = [([(Rewrite l, Match)], IntMap (Stat schd))]
-> ([(Rewrite l, Match)], IntMap (Stat schd))
forall a. Monoid a => [a] -> a
mconcat (((Int, Rewrite l) -> ([(Rewrite l, Match)], IntMap (Stat schd)))
-> [(Int, Rewrite l)]
-> [([(Rewrite l, Match)], IntMap (Stat schd))]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Database l
-> Int
-> IntMap (Stat schd)
-> (Int, Rewrite l)
-> ([(Rewrite l, Match)], IntMap (Stat schd))
matchWithScheduler Database l
db Int
i IntMap (Stat schd)
stats) ([Int] -> [Rewrite l] -> [(Int, Rewrite l)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] [Rewrite l]
rewrites))

            -- Write-only phase, temporarily break invariants
            [(Rewrite l, Match)]
-> ((Rewrite l, Match) -> EGraphM l ()) -> EGraphM l ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Rewrite l, Match)]
matches (Rewrite l, Match) -> EGraphM l ()
applyMatchesRhs

            -- Restore the invariants once per iteration
            EGraphM l ()
forall (l :: * -> *). Language l => EGraphM l ()
rebuild
            
            G.EGraph { memo :: forall (l :: * -> *). EGraph l -> Memo l
G.memo = Memo l
afterMemo, classes :: forall (l :: * -> *). EGraph l -> ClassIdMap (EClass l)
G.classes = ClassIdMap (EClass l)
afterClasses } <- StateT (EGraph l) Identity (EGraph l)
forall (m :: * -> *) s. Monad m => StateT s m s
get

            -- ROMES:TODO: Node limit...
            -- ROMES:TODO: Actual Timeout... not just iteration timeout
            -- ROMES:TODO Better saturation (see Runner)
            -- Apply rewrites until saturated or ROMES:TODO: timeout
            Bool -> EGraphM l () -> EGraphM l ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Memo l -> Int
forall (l :: * -> *) a. NodeMap l a -> Int
G.sizeNM Memo l
afterMemo Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Memo l -> Int
forall (l :: * -> *) a. NodeMap l a -> Int
G.sizeNM Memo l
beforeMemo
                      Bool -> Bool -> Bool
&& ClassIdMap (EClass l) -> Int
forall a. IntMap a -> Int
IM.size ClassIdMap (EClass l)
afterClasses Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ClassIdMap (EClass l) -> Int
forall a. IntMap a -> Int
IM.size ClassIdMap (EClass l)
beforeClasses)
                (Int -> IntMap (Stat schd) -> EGraphM l ()
equalitySaturation'' (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) IntMap (Stat schd)
newStats)

        matchWithScheduler :: Database l -> Int -> IM.IntMap (Stat schd) -> (Int, Rewrite l) -> ([(Rewrite l, Match)], IM.IntMap (Stat schd))
        matchWithScheduler :: Database l
-> Int
-> IntMap (Stat schd)
-> (Int, Rewrite l)
-> ([(Rewrite l, Match)], IntMap (Stat schd))
matchWithScheduler Database l
db Int
i IntMap (Stat schd)
stats = \case
            (Int
rw_id, Rewrite l
rw :| RewriteCondition l
cnd) -> ([(Rewrite l, Match)] -> [(Rewrite l, Match)])
-> ([(Rewrite l, Match)], IntMap (Stat schd))
-> ([(Rewrite l, Match)], IntMap (Stat schd))
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (((Rewrite l, Match) -> (Rewrite l, Match))
-> [(Rewrite l, Match)] -> [(Rewrite l, Match)]
forall a b. (a -> b) -> [a] -> [b]
map ((Rewrite l -> Rewrite l)
-> (Rewrite l, Match) -> (Rewrite l, Match)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Rewrite l -> RewriteCondition l -> Rewrite l
forall (lang :: * -> *).
Rewrite lang -> RewriteCondition lang -> Rewrite lang
:| RewriteCondition l
cnd))) (([(Rewrite l, Match)], IntMap (Stat schd))
 -> ([(Rewrite l, Match)], IntMap (Stat schd)))
-> ([(Rewrite l, Match)], IntMap (Stat schd))
-> ([(Rewrite l, Match)], IntMap (Stat schd))
forall a b. (a -> b) -> a -> b
$ Database l
-> Int
-> IntMap (Stat schd)
-> (Int, Rewrite l)
-> ([(Rewrite l, Match)], IntMap (Stat schd))
matchWithScheduler Database l
db Int
i IntMap (Stat schd)
stats (Int
rw_id, Rewrite l
rw)
            (Int
rw_id, Pattern l
lhs := Pattern l
rhs) -> do
                case Int -> IntMap (Stat schd) -> Maybe (Stat schd)
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
rw_id IntMap (Stat schd)
stats of
                  -- If it's banned until some iteration, don't match this rule
                  -- against anything.
                  Just Stat schd
s | forall s. Scheduler s => Int -> Stat s -> Bool
isBanned @schd Int
i Stat schd
s -> ([], IntMap (Stat schd)
stats)

                  -- Otherwise, match and update stats
                  Maybe (Stat schd)
x -> do

                      -- Match pattern
                      let matches' :: [Match]
matches' = Database l -> Pattern l -> [Match]
forall (l :: * -> *).
Language l =>
Database l -> Pattern l -> [Match]
ematch Database l
db Pattern l
lhs -- Add rewrite to the e-match substitutions

                      -- Backoff scheduler: update stats
                      let newStats :: IntMap (Stat schd)
newStats = forall s.
Scheduler s =>
Int
-> Int
-> Maybe (Stat s)
-> IntMap (Stat s)
-> [Match]
-> IntMap (Stat s)
updateStats @schd Int
i Int
rw_id Maybe (Stat schd)
x IntMap (Stat schd)
stats [Match]
matches'

                      ((Match -> (Rewrite l, Match)) -> [Match] -> [(Rewrite l, Match)]
forall a b. (a -> b) -> [a] -> [b]
map (Pattern l
lhs Pattern l -> Pattern l -> Rewrite l
forall (lang :: * -> *).
Pattern lang -> Pattern lang -> Rewrite lang
:= Pattern l
rhs,) [Match]
matches', IntMap (Stat schd)
newStats)

        applyMatchesRhs :: (Rewrite l, Match) -> EGraphM l ()
        applyMatchesRhs :: (Rewrite l, Match) -> EGraphM l ()
applyMatchesRhs =
            \case
                (Rewrite l
rw :| RewriteCondition l
cond, m :: Match
m@(Match Subst
subst Int
_)) -> do
                    -- If the rewrite condition is satisfied, applyMatchesRhs on the rewrite rule.
                    EGraph l
egr <- StateT (EGraph l) Identity (EGraph l)
forall (m :: * -> *) s. Monad m => StateT s m s
get
                    Bool -> EGraphM l () -> EGraphM l ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (RewriteCondition l
cond Subst
subst EGraph l
egr) (EGraphM l () -> EGraphM l ()) -> EGraphM l () -> EGraphM l ()
forall a b. (a -> b) -> a -> b
$
                       (Rewrite l, Match) -> EGraphM l ()
applyMatchesRhs (Rewrite l
rw, Match
m)

                (Pattern l
_ := VariablePattern Int
v, Match Subst
subst Int
eclass) -> do
                    -- rhs is equal to a variable, simply merge class where lhs
                    -- pattern was found (@eclass@) and the eclass the pattern
                    -- variable matched (@lookup v subst@)
                    case Int -> Subst -> Maybe Int
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
v Subst
subst of
                      Maybe Int
Nothing -> [Char] -> EGraphM l ()
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible: couldn't find v in subst"
                      Just Int
n  -> do
                          Int
_ <- Int -> Int -> EGraphM l Int
forall (l :: * -> *). Language l => Int -> Int -> EGraphM l Int
merge Int
n Int
eclass
                          () -> EGraphM l ()
forall a. a -> StateT (EGraph l) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

                (Pattern l
_ := NonVariablePattern l (Pattern l)
rhs, Match Subst
subst Int
eclass) -> do
                    -- rhs is (at the top level) a non-variable pattern, so substitute
                    -- all pattern variables in the pattern and create a new e-node (and
                    -- e-class that represents it), then merge the e-class of the
                    -- substituted rhs with the class that matched the left hand side
                    Int
eclass' <- Subst -> l (Pattern l) -> EGraphM l Int
reprPat Subst
subst l (Pattern l)
rhs
                    Int
_ <- Int -> Int -> EGraphM l Int
forall (l :: * -> *). Language l => Int -> Int -> EGraphM l Int
merge Int
eclass Int
eclass'
                    () -> EGraphM l ()
forall a. a -> StateT (EGraph l) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

        -- | Represent a pattern in the e-graph a pattern given substitions
        reprPat :: Subst -> l (Pattern l) -> EGraphM l ClassId
        reprPat :: Subst -> l (Pattern l) -> EGraphM l Int
reprPat Subst
subst = ENode l -> EGraphM l Int
forall (l :: * -> *). Language l => ENode l -> EGraphM l Int
add (ENode l -> EGraphM l Int)
-> (l Int -> ENode l) -> l Int -> EGraphM l Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. l Int -> ENode l
forall (l :: * -> *). l Int -> ENode l
G.Node (l Int -> EGraphM l Int)
-> (l (Pattern l) -> StateT (EGraph l) Identity (l Int))
-> l (Pattern l)
-> EGraphM l Int
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< (Pattern l -> EGraphM l Int)
-> l (Pattern l) -> StateT (EGraph l) Identity (l Int)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> l a -> f (l b)
traverse \case
            VariablePattern Int
v ->
                case Int -> Subst -> Maybe Int
forall a. Int -> IntMap a -> Maybe a
IM.lookup Int
v Subst
subst of
                    Maybe Int
Nothing -> [Char] -> EGraphM l Int
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible: couldn't find v in subst?"
                    Just Int
i  -> Int -> EGraphM l Int
forall a. a -> StateT (EGraph l) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i
            NonVariablePattern l (Pattern l)
p -> Subst -> l (Pattern l) -> EGraphM l Int
reprPat Subst
subst l (Pattern l)
p
{-# SCC equalitySaturation' #-}