Merge pull request #95 from firedrakeproject/conditional-simplification

* conditional-simplification:
  Add test of simplification in Conditional
  gem: Simplification in Conditional

gem/gem.py

@@ -361,12 +361,19 @@ def __init__(self, a, b):
 class Conditional(Node):
     __slots__ = ('children', 'shape')
 
-    def __init__(self, condition, then, else_):
+    def __new__(cls, condition, then, else_):
         assert not condition.shape
         assert then.shape == else_.shape
 
+        # If both branches are the same, just return one of them.  In
+        # particular, this will help constant-fold zeros.
+        if then == else_:
+            return then
+
+        self = super(Conditional, cls).__new__(cls)
         self.children = condition, then, else_
         self.shape = then.shape
+        return self
 
 
 class IndexBase(with_metaclass(ABCMeta)):

tests/test_simplification.py

@@ -0,0 +1,33 @@
+from __future__ import absolute_import, print_function, division
+
+import pytest
+
+from gem.gem import Variable, Zero, Conditional, \
+    LogicalAnd, Index, Indexed, Product
+
+
+def test_conditional_simplification():
+    a = Variable("A", ())
+    b = Variable("B", ())
+
+    expr = Conditional(LogicalAnd(b, a), a, a)
+
+    assert expr == a
+
+
+def test_conditional_zero_folding():
+    b = Variable("B", ())
+    a = Variable("A", (3, ))
+    i = Index()
+    expr = Conditional(LogicalAnd(b, b),
+                       Product(Indexed(a, (i, )),
+                               Zero()),
+                       Zero())
+
+    assert expr == Zero()
+
+
+if __name__ == "__main__":
+    import os
+    import sys
+    pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:])