Make Expr a Merkle binary tree to allow subtree comparison

1 files changed, 12 insertions, 11 deletions

diff --git a/mdpoly/abc.py b/mdpoly/abc.py
index 1914168..87a6c13 100644
--- a/mdpoly/abc.py
+++ b/mdpoly/abc.py
@@ -30,7 +30,7 @@ class Algebra(Enum):
 
 
 class Expr(ABC):
-    """ Binary tree to represent a mathematical expression. """
+    """ Merkle binary tree to represent a mathematical expression. """
 
     # --- Properties ---
 
@@ -93,6 +93,14 @@ class Expr(ABC):
     def __iter__(self) -> Iterable[Expr]:
         yield from self.children()
 
+    def __hash__(self) -> int:
+        return hash(self.left) + hash(self.right)
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Expr):
+            return hash(self) == hash(other)
+        return False
+
     # --- Methods for polynomial expression ---
 
     @abstractmethod
@@ -143,16 +151,6 @@ class Expr(ABC):
         It can also be used for more advanced replacements since ``old`` and
         ``new`` are free to be any expression.
         """
-        # FIXME: this doesn't work if old is something like y ** 2,
-        # TODO: Define equality on expr object
-        if not isinstance(old, Leaf):
-            import warnings
-            warnings.warn(
-                "FIXME: You are trying to replace an expression with another "
-                "expression, but this does not work yet. Currently it is only "
-                "possible to replace leaves (Variable, Constant, Parameter) with "
-                "leaves / expressions.")
-
         def replace_all(node):
             if node == old:
                 return new 
@@ -254,6 +252,9 @@ class Leaf(Expr):
     def __repr__(self) -> str:
         return self.name
 
+    def __hash__(self) -> int:
+        return hash((self.__class__.__qualname__, self.algebra, self.name, self.shape))
+
     # -- Overloads --- 
 
     @property