update docs and add some missing expressions to EqualityMapper

Author: alexfikl

pymbolic/mapper/equality.py

@@ -27,51 +27,93 @@
 
 
 class EqualityMapper(Mapper):
+    __slots__ = ["_ids_to_result"]
+
     def __init__(self) -> None:
         self._ids_to_result: Dict[Tuple[int, int], bool] = {}
 
     def rec(self, expr: Any, other: Any) -> bool:
         key = (id(expr), id(other))
-
-        try:
-            result = self._ids_to_result[key]
-        except KeyError:
-            if expr is other:
-                result = True
-            elif expr.__class__ != other.__class__:
-                result = False
-            elif hash(expr) != hash(other):
-                result = False
-            else:
-                try:
-                    method = getattr(self, expr.mapper_method)
-                except AttributeError:
-                    if isinstance(expr, Expression):
-                        result = self.handle_unsupported_expression(expr, other)
-                    else:
-                        result = self.map_foreign(expr, other)
+        if key in self._ids_to_result:
+            return self._ids_to_result[key]
+
+        if expr is other:
+            result = True
+        elif expr.__class__ != other.__class__:
+            result = False
+        else:
+            try:
+                method = getattr(self, expr.mapper_method)
+            except AttributeError:
+                if isinstance(expr, Expression):
+                    result = self.handle_unsupported_expression(expr, other)
                 else:
-                    result = method(expr, other)
-
-            self._ids_to_result[key] = result
+                    result = self.map_foreign(expr, other)
+            else:
+                result = method(expr, other)
 
+        self._ids_to_result[key] = result
         return result
 
     def __call__(self, expr: Any, other: Any) -> bool:
         return self.rec(expr, other)
 
+    # {{{ handle_unsupported_expression
+
     def handle_unsupported_expression(self, expr, other) -> bool:
         eq = expr.make_equality_mapper()
         if type(self) == type(eq):
             raise UnsupportedExpressionError(
                     "'{}' cannot handle expressions of type '{}'".format(
                         type(self).__name__, type(expr).__name__))
 
+        # NOTE: this may look fishy, but we want to preserve the cache as we
+        # go through the expression tree, so that it does not do
+        # unnecessary checks when we change the mapper for some subclass
         eq._ids_to_result = self._ids_to_result
+
         return eq(expr, other)
 
-    def map_constant(self, expr, other) -> bool:
-        return expr == other
+    # }}}
+
+    # {{{ foreign
+
+    def map_tuple(self, expr, other) -> bool:
+        return (
+                len(expr) == len(other)
+                and all(self.rec(el, other_el)
+                    for el, other_el in zip(expr, other)))
+
+    def map_foreign(self, expr, other) -> bool:
+        from pymbolic.primitives import VALID_CONSTANT_CLASSES
+
+        if isinstance(expr, VALID_CONSTANT_CLASSES):
+            return expr == other
+        elif isinstance(expr, tuple):
+            return self.map_tuple(expr, other)
+        else:
+            raise ValueError(
+                    f"{type(self).__name__} encountered invalid foreign object: "
+                    f"{expr!r}")
+
+    # }}}
+
+    # {{{
+
+    # NOTE: `type(expr) == type(other)` is checked in `__call__`, so the
+    # checks below can assume that the two operands always have the same type
+
+    # NOTE: as much as possible, these should try to put the "cheap" checks
+    # first so that the shortcircuiting removes the need to to extra work
+
+    # NOTE: `all` is also shortcircuiting, so should be better to use a
+    # generator there to avoid extra work
+
+    def map_wildcard(self, expr, other) -> bool:
+        return True
+
+    def map_function_symbol(self, expr, other) -> bool:
+        return True
 
     def map_variable(self, expr, other) -> bool:
         return expr.name == other.name
@@ -112,6 +154,7 @@ def map_sum(self, expr, other) -> bool:
                     for child, other_child in zip(expr.children, other.children))
                 )
 
+    map_slice = map_sum
     map_product = map_sum
     map_min = map_sum
     map_max = map_sum
@@ -143,8 +186,8 @@ def map_power(self, expr, other) -> bool:
 
     def map_left_shift(self, expr, other) -> bool:
         return (
-                self.rec(expr.shiftee, other.shiftee)
-                and self.rec(expr.shift, other.shift))
+                self.rec(expr.shift, other.shift)
+                and self.rec(expr.shiftee, other.shiftee))
 
     map_right_shift = map_left_shift
 
@@ -160,6 +203,12 @@ def map_if(self, expr, other) -> bool:
                 and self.rec(expr.then, other.then)
                 and self.rec(expr.else_, other.else_))
 
+    def map_if_positive(self, expr, other) -> bool:
+        return (
+                self.rec(expr.criterion, other.criterion)
+                and self.rec(expr.then, other.then)
+                and self.rec(expr.else_, other.else_))
+
     def map_common_subexpression(self, expr, other) -> bool:
         return (
                 expr.prefix == other.prefix
@@ -188,12 +237,4 @@ def map_polynomial(self, expr, other) -> bool:
                 self.rec(expr.Base, other.Data)
                 and self.rec(expr.Data, other.Data))
 
-    # {{{ foreign
-
-    def map_tuple(self, expr, other) -> bool:
-        return (
-                len(expr) == len(other)
-                and all(self.rec(el, other_el)
-                    for el, other_el in zip(expr, other)))
-
     # }}}

pymbolic/primitives.py

@@ -193,7 +193,7 @@ class Expression:
     .. automethod:: make_stringifier
 
     .. automethod:: __eq__
-    .. automethod:: is_equal
+    .. automethod:: make_equality_mapper
     .. automethod:: __hash__
     .. automethod:: get_hash
     .. automethod:: __str__
@@ -499,11 +499,10 @@ def __repr__(self):
     # {{{ hash/equality interface
 
     def __eq__(self, other):
-        """Provides equality testing with quick positive and negative paths
-        based on :func:`id` and :meth:`__hash__`.
+        """Provides equality testing with quick positive and negative paths.
 
         Subclasses should generally not override this method, but instead
-        provide an implementation of :meth:`is_equal`.
+        provide an implementation of :meth:`make_equality_mapper`.
         """
         return self.make_equality_mapper()(self, other)