diff --git a/python/dune/perftool/ufl/visitor.py b/python/dune/perftool/ufl/visitor.py index 8773efff42e76ca8030e09810f143dec03e2ee21..2ac22c3e4076fdb9b9bcb0bb7357c48318ebb6a5 100644 --- a/python/dune/perftool/ufl/visitor.py +++ b/python/dune/perftool/ufl/visitor.py @@ -325,7 +325,15 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker): # def conditional(self, o): - return prim.If(*tuple(self.call(op) for op in o.ufl_operands)) + condition = self.call(o.ufl_operands[0]) + indices = self.indices + + op1 = self.call(o.ufl_operands[1]) + # Restore indexing information for the second branch + self.indices = indices + op2 = self.call(o.ufl_operands[2]) + + return prim.If(condition, op1, op2) def eq(self, o): return prim.Comparison(self.call(o.ufl_operands[0]),