Lazy evaluation of ite operator

cksystemsteaching · Jan 15, 2025 · 0eea271 · 0eea271
1 parent b5ef7cd
commit 0eea271
Showing 1 changed file with 58 additions and 23 deletions.
diff --git a/tools/bitme.py b/tools/bitme.py
@@ -436,15 +436,15 @@ def get_expression(self):
         assert len(self.values) > 0
         if isinstance(self.sid_line, Bool):
             assert len(self.values) <= 2
-            true_line = Constant.false
             false_line = Constant.false
+            true_line = Constant.false
             for value in self.values:
                 constraint_line = self.values[value]
-                if value == 1:
-                    true_line = constraint_line
-                else:
-                    assert value == 0
+                if value == 0:
                     false_line = Values.NOT(constraint_line)
+                else:
+                    assert value == 1
+                    true_line = constraint_line
             exp_line = Values.OR(false_line, true_line)
         else:
             exp_line = None
@@ -1614,25 +1614,60 @@ def get_mapped_array_expression_for(self, index):
             self.ite_cache[index] = self.copy(arg1_line, arg2_line, arg3_line)
         return self.ite_cache[index]
 
-    def propagate(self, arg1_value, arg2_value, arg3_value):
-        assert len(arg1_value.values) <= 2
-        true_line = Constant.false
-        false_line = Constant.false
-        for value in arg1_value.values:
-            constraint_line = arg1_value.values[value]
-            if value == 1:
-                true_line = constraint_line
+    def get_values(self, step):
+        if step not in self.cache_values:
+            arg1_value = self.arg1_line.get_values(step)
+            if isinstance(arg1_value, Values):
+                assert len(arg1_value.values) <= 2
+                false_line = Constant.false
+                true_line = Constant.false
+                for value in arg1_value.values:
+                    constraint_line = arg1_value.values[value]
+                    if value == 0:
+                        false_line = constraint_line
+                    else:
+                        assert value == 1
+                        true_line = constraint_line
+                if false_line == Constant.false:
+                    arg2_value = self.arg2_line.get_values(step)
+                    if isinstance(arg2_value, Values):
+                        self.cache_values[step] = arg2_value.constrain(true_line)
+                        return self.cache_values[step]
+                    elif true_line == Constant.true:
+                        # true case holds unconditionally
+                        self.cache_values[step] = arg2_value.get_expression()
+                        return self.cache_values[step]
+                    else:
+                        # lazy evaluation of false case
+                        arg3_value = self.arg3_line.get_values(step)
+                elif true_line == Constant.false:
+                    arg3_value = self.arg3_line.get_values(step)
+                    if isinstance(arg3_value, Values):
+                        self.cache_values[step] = arg3_value.constrain(false_line)
+                        return self.cache_values[step]
+                    elif false_line == Constant.true:
+                        # false case holds unconditionally
+                        self.cache_values[step] = arg3_value.get_expression()
+                        return self.cache_values[step]
+                    else:
+                        # lazy evaluation of true case
+                        arg2_value = self.arg2_line.get_values(step)
+                else:
+                    arg2_value = self.arg2_line.get_values(step)
+                    arg3_value = self.arg3_line.get_values(step)
+                    if isinstance(arg2_value, Values) and isinstance(arg3_value, Values):
+                        arg2_value = arg2_value.constrain(true_line)
+                        arg3_value = arg3_value.constrain(false_line)
+                        self.cache_values[step] = arg2_value.merge(arg3_value)
+                        return self.cache_values[step]
             else:
-                assert value == 0
-                false_line = constraint_line
-        if true_line == Constant.false:
-            return self.arg3_value.constrain(false_line)
-        elif false_line == Constant.false:
-            return self.arg2_value.constrain(true_line)
-        else:
-            results2 = self.arg2_value.constrain(true_line)
-            results3 = self.arg3_value.constrain(false_line)
-            return results2.merge(results3)
+                arg2_value = self.arg2_line.get_values(step)
+                arg3_value = self.arg3_line.get_values(step)
+            arg1_value = arg1_value.get_expression()
+            arg2_value = arg2_value.get_expression()
+            arg3_value = arg3_value.get_expression()
+            self.cache_values[step] = self.copy(arg1_value, arg2_value, arg3_value)
+        return self.cache_values[step]
 
     def get_z3(self):
         if self.z3 is None: