content: {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE MonoLocalBinds #-}
  {-|
     Monadic interface to e-graph stateful computations
   -}
  module Data.Equality.Graph.Monad
    (
      -- * Threading e-graphs in a stateful computation
      --
      -- | These are the same operations over e-graphs as in 'Data.Equality.Graph',
      -- but defined in the context of a 'State' monad threading around the e-graph.
      egraph
    , represent
    , add
    , merge
    , rebuild
    , EG.canonicalize
    , EG.find
    , EG.emptyEGraph
  
      -- * E-graph transformations for monadic analysis
      --
      -- | The same e-graph operations in a stateful computation threading around
      -- the e-graph, but for 'Analysis' defined monadically ('AnalysisM').
    , representM, addM, mergeM, rebuildM
  
    -- * E-graph stateful computations
    , EGraphM
    , EGraphMT
    , runEGraphM
    , runEGraphMT
  
    -- * E-graph definition re-export
    , EG.EGraph
  
    -- * 'State' monad re-exports
    , modify, get, gets
    ) where
  
  import Control.Monad ((>=>))
  import Control.Monad.Trans.State.Strict
  
  import Data.Equality.Utils (Fix, cata)
  
  import Data.Equality.Analysis
  import qualified Data.Equality.Analysis.Monadic as AM
  import Data.Equality.Graph (EGraph, ClassId, Language, ENode(..))
  import qualified Data.Equality.Graph as EG
  
  -- | E-graph stateful computation
  type EGraphM a l = State (EGraph a l)
  -- | E-graph stateful computation over an arbitrary monad
  type EGraphMT a l = StateT (EGraph a l)
  
  -- | Run EGraph computation on an empty e-graph
  --
  -- === Example
  -- @
  -- egraph $ do
  --  id1 <- represent t1
  --  id2 <- represent t2
  --  merge id1 id2
  -- @
  egraph :: Language l => EGraphM anl l a -> (a, EGraph anl l)
  egraph = runEGraphM EG.emptyEGraph
  {-# INLINE egraph #-}
  
  -- | Represent an expression (@Fix l@) in an e-graph by recursively
  -- representing sub expressions
  represent :: (Analysis anl l, Language l) => Fix l -> EGraphM anl l ClassId
  represent = cata $ sequence >=> add . Node
  {-# INLINE represent #-}
  
  -- | Add an e-node to the e-graph
  add :: (Analysis anl l, Language l) => ENode l -> EGraphM anl l ClassId
  add = StateT . fmap pure . EG.add
  {-# INLINE add #-}
  
  -- | Merge two e-classes by id
  --
  -- E-graph invariants may be broken by merging, and 'rebuild' should be used
  -- /eventually/ to restore them
  merge :: (Analysis anl l, Language l) => ClassId -> ClassId -> EGraphM anl l ClassId
  merge a b = StateT (pure <$> EG.merge a b)
  {-# INLINE merge #-}
  
  -- | Rebuild: Restore e-graph invariants
  --
  -- E-graph invariants are traditionally maintained after every merge, but we
  -- allow operations to temporarilly break the invariants (specifically, until we call
  -- 'rebuild')
  --
  -- The paper describing rebuilding in detail is https://arxiv.org/abs/2004.03082
  rebuild :: (Analysis anl l, Language l) => EGraphM anl l ()
  rebuild = StateT (pure . ((),). EG.rebuild)
  {-# INLINE rebuild #-}
  
  -- | Run 'EGraphM' computation on a given e-graph
  runEGraphM :: EGraph anl l -> EGraphM anl l a -> (a, EGraph anl l)
  runEGraphM = flip runState
  {-# INLINE runEGraphM #-}
  
  --------------------------------------------------------------------------------
  -- Monadic Analysis interface
  
  -- | Run 'EGraphM' computation on a given e-graph over a monadic analysis
  runEGraphMT :: EGraph anl l -> EGraphMT anl l m a -> m (a, EGraph anl l)
  runEGraphMT = flip runStateT
  {-# INLINE runEGraphMT #-}
  
  -- | Like 'represent', but for a monadic analysis
  representM :: (AM.AnalysisM m anl l, Language l) => Fix l -> EGraphMT anl l m ClassId
  representM = StateT . EG.representM
  {-# INLINE representM #-}
  
  -- | Like 'add', but for a monadic analysis
  addM :: (AM.AnalysisM m anl l, Language l) => ENode l -> EGraphMT anl l m ClassId
  addM = StateT . EG.addM
  {-# INLINE addM #-}
  
  -- | Like 'merge', but for a monadic analysis
  mergeM :: (AM.AnalysisM m anl l, Language l) => ClassId -> ClassId -> EGraphMT anl l m ClassId
  mergeM a b = StateT (EG.mergeM a b)
  {-# INLINE mergeM #-}
  
  -- | Like 'rebuild', but for a monadic analysis
  rebuildM :: (AM.AnalysisM m anl l, Language l) => EGraphMT anl l m ()
  rebuildM = StateT (fmap ((),) . EG.rebuildM)
  {-# INLINE rebuildM #-}
  