{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE BangPatterns #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Algorithm.EqSat.Queries
-- Copyright   :  (c) Fabricio Olivetti 2021 - 2024
-- License     :  BSD3
-- Maintainer  :  fabricio.olivetti@gmail.com
-- Stability   :  experimental
-- Portability :
--
-- Query functions for e-graphs
-- Heavily based on hegg (https://github.com/alt-romes/hegg by alt-romes)
--
-----------------------------------------------------------------------------

module Algorithm.EqSat.Queries where

import Algorithm.EqSat.Egraph
import qualified Data.IntMap as IntMap
import qualified Data.Map as Map
import qualified Data.HashSet as Set
import Control.Monad.State ( gets, modify' )
import Control.Monad ( filterM )
import Control.Lens ( over )
import Data.Maybe
import Data.Sequence ( Seq(..) )

import Debug.Trace

-- this is too slow for now, it needs a db of its own
-- basically a db for each query we need
getEClassesThat :: Monad m => (EClass -> Bool) -> EGraphST m [EClassId]
getEClassesThat :: forall (m :: * -> *).
Monad m =>
(EClass -> Bool) -> EGraphST m [EClassId]
getEClassesThat EClass -> Bool
p = do
    (EGraph -> [EClassId]) -> StateT EGraph m [EClassId]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (((EClassId, EClass) -> EClassId)
-> [(EClassId, EClass)] -> [EClassId]
forall a b. (a -> b) -> [a] -> [b]
map (EClassId, EClass) -> EClassId
forall a b. (a, b) -> a
fst ([(EClassId, EClass)] -> [EClassId])
-> (EGraph -> [(EClassId, EClass)]) -> EGraph -> [EClassId]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((EClassId, EClass) -> Bool)
-> [(EClassId, EClass)] -> [(EClassId, EClass)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(EClassId
ecId, EClass
ec) -> EClass -> Bool
p EClass
ec) ([(EClassId, EClass)] -> [(EClassId, EClass)])
-> (EGraph -> [(EClassId, EClass)])
-> EGraph
-> [(EClassId, EClass)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap EClass -> [(EClassId, EClass)]
forall a. IntMap a -> [(EClassId, a)]
IntMap.toList (IntMap EClass -> [(EClassId, EClass)])
-> (EGraph -> IntMap EClass) -> EGraph -> [(EClassId, EClass)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EGraph -> IntMap EClass
_eClass)
    --go ecs
        where
            go :: Monad m => [EClassId] -> EGraphST m [EClassId]
            go :: forall (m :: * -> *).
Monad m =>
[EClassId] -> EGraphST m [EClassId]
go [] = [EClassId] -> StateT EGraph m [EClassId]
forall a. a -> StateT EGraph m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
            go (EClassId
ecId:[EClassId]
ecs) = do ec <- (EGraph -> Bool) -> StateT EGraph m Bool
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (EClass -> Bool
p (EClass -> Bool) -> (EGraph -> EClass) -> EGraph -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IntMap EClass -> EClassId -> EClass
forall a. IntMap a -> EClassId -> a
IntMap.! EClassId
ecId) (IntMap EClass -> EClass)
-> (EGraph -> IntMap EClass) -> EGraph -> EClass
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EGraph -> IntMap EClass
_eClass)
                               ecs' <- go ecs
                               if ec
                                  then pure (ecId:ecs')
                                  else pure ecs'

updateFitness :: Monad m => Double -> EClassId -> EGraphST m ()
updateFitness :: forall (m :: * -> *).
Monad m =>
Double -> EClassId -> EGraphST m ()
updateFitness Double
f EClassId
ecId = do
   ec   <- (EGraph -> EClass) -> StateT EGraph m EClass
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((IntMap EClass -> EClassId -> EClass
forall a. IntMap a -> EClassId -> a
IntMap.! EClassId
ecId) (IntMap EClass -> EClass)
-> (EGraph -> IntMap EClass) -> EGraph -> EClass
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EGraph -> IntMap EClass
_eClass)
   let info = EClass -> EClassData
_info EClass
ec
   modify' $ over eClass (IntMap.insert ecId ec{_info=info{_fitness = Just f}})

-- | returns all the root e-classes (e-class without parents)
findRootClasses :: Monad m => EGraphST m [EClassId]
findRootClasses :: forall (m :: * -> *). Monad m => EGraphST m [EClassId]
findRootClasses = (EGraph -> [EClassId]) -> StateT EGraph m [EClassId]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (((EClassId, EClass) -> EClassId)
-> [(EClassId, EClass)] -> [EClassId]
forall a b. (a -> b) -> [a] -> [b]
Prelude.map (EClassId, EClass) -> EClassId
forall a b. (a, b) -> a
fst ([(EClassId, EClass)] -> [EClassId])
-> (EGraph -> [(EClassId, EClass)]) -> EGraph -> [EClassId]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((EClassId, EClass) -> Bool)
-> [(EClassId, EClass)] -> [(EClassId, EClass)]
forall a. (a -> Bool) -> [a] -> [a]
Prelude.filter (EClassId, EClass) -> Bool
isParent ([(EClassId, EClass)] -> [(EClassId, EClass)])
-> (EGraph -> [(EClassId, EClass)])
-> EGraph
-> [(EClassId, EClass)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap EClass -> [(EClassId, EClass)]
forall a. IntMap a -> [(EClassId, a)]
IntMap.toList (IntMap EClass -> [(EClassId, EClass)])
-> (EGraph -> IntMap EClass) -> EGraph -> [(EClassId, EClass)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EGraph -> IntMap EClass
_eClass)
  where
    isParent :: (EClassId, EClass) -> Bool
isParent (EClassId
k, EClass
v) = HashSet (EClassId, ENode) -> Bool
forall a. HashSet a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
Prelude.null (EClass -> HashSet (EClassId, ENode)
_parents EClass
v) Bool -> Bool -> Bool
||  (EClassId
k EClassId -> HashSet EClassId -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` (((EClassId, ENode) -> EClassId)
-> HashSet (EClassId, ENode) -> HashSet EClassId
forall b a.
(Hashable b, Eq b) =>
(a -> b) -> HashSet a -> HashSet b
Set.map (EClassId, ENode) -> EClassId
forall a b. (a, b) -> a
fst (EClass -> HashSet (EClassId, ENode)
_parents EClass
v)))

-- | returns the e-class id with the best fitness that
-- is true to a predicate
getTopECLassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId]
getTopECLassThat :: forall (m :: * -> *).
Monad m =>
EClassId -> (EClass -> Bool) -> EGraphST m [EClassId]
getTopECLassThat EClassId
n EClass -> Bool
p = do
  (EGraph -> RangeTree Double) -> StateT EGraph m (RangeTree Double)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (EGraphDB -> RangeTree Double
_fitRangeDB (EGraphDB -> RangeTree Double)
-> (EGraph -> EGraphDB) -> EGraph -> RangeTree Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EGraph -> EGraphDB
_eDB)
    StateT EGraph m (RangeTree Double)
-> (RangeTree Double -> StateT EGraph m [EClassId])
-> StateT EGraph m [EClassId]
forall a b.
StateT EGraph m a -> (a -> StateT EGraph m b) -> StateT EGraph m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= EClassId
-> [EClassId] -> RangeTree Double -> StateT EGraph m [EClassId]
forall (m :: * -> *).
Monad m =>
EClassId -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]
go EClassId
n []
  where
    go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]
    go :: forall (m :: * -> *).
Monad m =>
EClassId -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]
go EClassId
0 [EClassId]
bests RangeTree Double
rt = [EClassId] -> StateT EGraph m [EClassId]
forall a. a -> StateT EGraph m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [EClassId]
bests
    go EClassId
m [EClassId]
bests RangeTree Double
rt = case RangeTree Double
rt of
                       RangeTree Double
Empty   -> [EClassId] -> StateT EGraph m [EClassId]
forall a. a -> StateT EGraph m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [EClassId]
bests
                       RangeTree Double
t :|> (Double, EClassId)
y -> do let x :: EClassId
x = (Double, EClassId) -> EClassId
forall a b. (a, b) -> b
snd (Double, EClassId)
y
                                     ecId <- EClassId -> EGraphST m EClassId
forall (m :: * -> *). Monad m => EClassId -> EGraphST m EClassId
canonical EClassId
x
                                     ec <- gets ((IntMap.! ecId) . _eClass)
                                     if (isInfinite . fromJust . _fitness . _info $ ec)
                                       then pure bests
                                       else if p ec
                                              then go (m-1) (x:bests) t
                                              else go m bests t
getTopECLassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId]
getTopECLassWithSize :: forall (m :: * -> *).
Monad m =>
EClassId -> EClassId -> EGraphST m [EClassId]
getTopECLassWithSize EClassId
sz EClassId
n = do
  (EGraph -> Maybe (RangeTree Double))
-> StateT EGraph m (Maybe (RangeTree Double))
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ((IntMap (RangeTree Double) -> EClassId -> Maybe (RangeTree Double)
forall a. IntMap a -> EClassId -> Maybe a
IntMap.!? EClassId
sz) (IntMap (RangeTree Double) -> Maybe (RangeTree Double))
-> (EGraph -> IntMap (RangeTree Double))
-> EGraph
-> Maybe (RangeTree Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EGraphDB -> IntMap (RangeTree Double)
_sizeFitDB (EGraphDB -> IntMap (RangeTree Double))
-> (EGraph -> EGraphDB) -> EGraph -> IntMap (RangeTree Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EGraph -> EGraphDB
_eDB)
    StateT EGraph m (Maybe (RangeTree Double))
-> (Maybe (RangeTree Double) -> StateT EGraph m [EClassId])
-> StateT EGraph m [EClassId]
forall a b.
StateT EGraph m a -> (a -> StateT EGraph m b) -> StateT EGraph m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= EClassId
-> [EClassId]
-> Maybe (RangeTree Double)
-> StateT EGraph m [EClassId]
forall (m :: * -> *).
Monad m =>
EClassId
-> [EClassId] -> Maybe (RangeTree Double) -> EGraphST m [EClassId]
go EClassId
n []
  where
    go :: Monad m => Int -> [EClassId] -> Maybe (RangeTree Double) -> EGraphST m [EClassId]
    go :: forall (m :: * -> *).
Monad m =>
EClassId
-> [EClassId] -> Maybe (RangeTree Double) -> EGraphST m [EClassId]
go EClassId
_ [EClassId]
bests Maybe (RangeTree Double)
Nothing   = [EClassId] -> StateT EGraph m [EClassId]
forall a. a -> StateT EGraph m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    go EClassId
0 [EClassId]
bests (Just RangeTree Double
rt) = [EClassId] -> StateT EGraph m [EClassId]
forall a. a -> StateT EGraph m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [EClassId]
bests
    go EClassId
m [EClassId]
bests (Just RangeTree Double
rt) = case RangeTree Double
rt of
                             RangeTree Double
Empty   -> [EClassId] -> StateT EGraph m [EClassId]
forall a. a -> StateT EGraph m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [EClassId]
bests
                             RangeTree Double
t :|> (Double, EClassId)
y -> do let x :: EClassId
x = (Double, EClassId) -> EClassId
forall a b. (a, b) -> b
snd (Double, EClassId)
y
                                           ecId <- EClassId -> EGraphST m EClassId
forall (m :: * -> *). Monad m => EClassId -> EGraphST m EClassId
canonical EClassId
x
                                           ec <- gets ((IntMap.! ecId) . _eClass)
                                           if (isInfinite . fromJust . _fitness . _info $ ec)
                                             then pure bests
                                             else go (m-1) (x:bests) (Just t)