From f73e6cd9bb19680409c924056692280c86c47257 Mon Sep 17 00:00:00 2001
From: Rodrigo Mesquita <rodrigo.m.mesquita@gmail.com>
Date: Thu, 6 Jul 2023 00:26:03 +0100
Subject: [PATCH] Add operation to create mapping in union find

---
 src/Data/Equality/Graph.hs               | 34 ++++++++++++++++++++++--
 src/Data/Equality/Graph/ReprUnionFind.hs |  2 +-
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/Data/Equality/Graph.hs b/src/Data/Equality/Graph.hs
index ce47158..e07f449 100644
--- a/src/Data/Equality/Graph.hs
+++ b/src/Data/Equality/Graph.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE MagicHash #-}
 -- {-# LANGUAGE ApplicativeDo #-}
 {-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE FlexibleContexts #-}
@@ -24,7 +25,10 @@ module Data.Equality.Graph
     , find, canonicalize
 
       -- * Functions on e-graphs
-    , emptyEGraph, newEClass
+    , emptyEGraph
+
+      -- ** Low-level operations
+    , newEClass, newPointerToClassId
 
       -- * Re-exports
     , module Data.Equality.Graph.Classes
@@ -32,6 +36,8 @@ module Data.Equality.Graph
     , module Data.Equality.Language
     ) where
 
+import GHC.Exts (Int(..), (+#), (<#), isTrue#)
+
 -- ROMES:TODO: Is the E-Graph a Monad if the analysis data were the type arg? i.e. Monad (EGraph language)?
 
 -- import GHC.Conc
@@ -52,6 +58,7 @@ import Data.Equality.Utils.SizedList
 
 import Data.Equality.Graph.Internal
 import Data.Equality.Graph.ReprUnionFind
+import qualified Data.Equality.Utils.IntToIntMap as IIM
 import Data.Equality.Graph.Classes
 import Data.Equality.Graph.Nodes
 import Data.Equality.Analysis
@@ -334,7 +341,30 @@ newEClass adata egr =
             , classes   = IM.insert new_eclass_id new_eclass (classes egr)
             }
       )
-{-# INLINE newEClass #-}
+{-# INLINEABLE newEClass #-}
+
+-- | Create a mapping from some class-id that does not exist in the e-graph to
+-- the given e-class id target. In practice, this basically creates an
+-- alias from the a given class-id to the e-class id of the target
+--
+-- If, instead, you want to create a mapping from an existing class-id to another one, use 'merge'.
+--
+-- Under the hood, this operation will bump the union find counter for next-ids
+-- to the given id+1 and add an entry to the union find from given id to the
+-- given target id.
+--
+-- INVARIANT: The given e-class pointer does not exist in the e-graph
+newPointerToClassId :: ClassId -- ^ Given Id (pointer) that will point to the target
+                    -> ClassId -- ^ The target id
+                    -> EGraph a l -> EGraph a l
+newPointerToClassId (I# pointer) (I# target) egr =
+   egr { unionFind = case unionFind egr of
+                       RUF im _size ->
+                         if isTrue# (pointer <# _size)
+                            then error $ "newPointerToClassId: given pointer id (" ++ show (I# pointer) ++ ") already exists in the e-graph"
+                            else RUF (IIM.insert pointer target im) (pointer +# 1#)
+       }
+{-# INLINEABLE newPointerToClassId #-}
 
 -- | Represent an expression (in fix-point form) and merge it with the e-class with the given id
 representAndMerge :: (Analysis a l, Language l) => ClassId -> Fix l -> EGraph a l -> EGraph a l
diff --git a/src/Data/Equality/Graph/ReprUnionFind.hs b/src/Data/Equality/Graph/ReprUnionFind.hs
index b1a8f54..5dcf97b 100644
--- a/src/Data/Equality/Graph/ReprUnionFind.hs
+++ b/src/Data/Equality/Graph/ReprUnionFind.hs
@@ -8,7 +8,7 @@ Union-find-like data structure that defines equivalence classes of e-class ids.
 
 -}
 module Data.Equality.Graph.ReprUnionFind
-  ( ReprUnionFind
+  ( ReprUnionFind(..)
   , emptyUF
   , makeNewSet
   , unionSets