Expr-depth-iterator: support iterating over mutated expressions

src/util/expr_iterator.h

@@ -42,22 +42,18 @@ class const_unique_depth_iteratort;
 struct depth_iterator_expr_statet final
 {
   typedef exprt::operandst::const_iterator operands_iteratort;
-  inline depth_iterator_expr_statet(
-    const exprt &expr,
-    operands_iteratort it,
-    operands_iteratort end):
-    expr(expr), it(it), end(end) { }
+  explicit depth_iterator_expr_statet(const exprt &expr) : expr(expr), op_idx(0)
+  {
+  }
   std::reference_wrapper<const exprt> expr;
-  operands_iteratort it;
-  operands_iteratort end;
+  std::size_t op_idx;
 };
 
 inline bool operator==(
   const depth_iterator_expr_statet &left,
   const depth_iterator_expr_statet &right)
 {
-  return distance(left.it, left.end) == distance(right.it, right.end) &&
-         left.expr.get() == right.expr.get();
+  return left.op_idx == right.op_idx && left.expr.get() == right.expr.get();
 }
 
 /// Depth first search iterator base - iterates over supplied expression
@@ -100,16 +96,19 @@ class depth_iterator_baset
     PRECONDITION(!m_stack.empty());
     while(true)
     {
-      if(m_stack.back().it==m_stack.back().end)
+      if(m_stack.back().op_idx == m_stack.back().expr.get().operands().size())
       {
         m_stack.pop_back();
         if(m_stack.empty())
           break;
       }
       // Check eg. if we haven't seen this node before
-      else if(this->downcast().push_expr(*m_stack.back().it))
+      else if(this->downcast().push_expr(
+                m_stack.back().expr.get().operands()[m_stack.back().op_idx]))
+      {
         break;
-      m_stack.back().it++;
+      }
+      ++m_stack.back().op_idx;
     }
     return this->downcast();
   }
@@ -120,7 +119,7 @@ class depth_iterator_baset
     m_stack.pop_back();
     if(!m_stack.empty())
     {
-      ++m_stack.back().it;
+      ++m_stack.back().op_idx;
       return ++(*this);
     }
     return this->downcast();
@@ -185,18 +184,11 @@ class depth_iterator_baset
     for(auto &state : m_stack)
     {
       // This deliberately breaks sharing as expr is now non-const
-      auto &operands = expr->operands();
-      // Get iterators into the operands of the new expr corresponding to the
-      // ones into the operands of the old expr
-      const auto it = operands.end() - (state.end - state.it);
+      (void)expr->write();
       state.expr = *expr;
-      state.it=it;
-      state.end=operands.end();
       // Get the expr for the next level down to use in the next iteration
-      if(!(state==m_stack.back()))
-      {
-        expr = &*it;
-      }
+      if(!(state == m_stack.back()))
+        expr = &expr->operands()[state.op_idx];
     }
     return *expr;
   }
@@ -208,7 +200,7 @@ class depth_iterator_baset
   ///   false otherwise. If returning false, child will not be iterated over.
   bool push_expr(const exprt &expr)
   {
-    m_stack.emplace_back(expr, expr.operands().begin(), expr.operands().end());
+    m_stack.emplace_back(expr);
     return true;
   }
 

unit/Makefile

@@ -50,6 +50,7 @@ SRC += analyses/ai/ai.cpp \
        util/cmdline.cpp \
        util/expr_cast/expr_cast.cpp \
        util/expr.cpp \
+       util/expr_iterator.cpp \
        util/file_util.cpp \
        util/format_number_range.cpp \
        util/get_base_name.cpp \

unit/util/expr_iterator.cpp

@@ -0,0 +1,103 @@
+/*******************************************************************\
+
+Module: Unit tests for expr_iterator
+
+Author: Diffblue Ltd
+
+\*******************************************************************/
+
+/// \file
+/// expr_iterator unit tests
+
+#include <testing-utils/use_catch.h>
+
+#include <util/expr_iterator.h>
+#include <util/std_expr.h>
+
+SCENARIO("expr_iterator", "[core][utils][expr_iterator]")
+{
+  GIVEN("An expression of depth 3")
+  {
+    symbol_exprt symbol("whatever", bool_typet());
+    notequal_exprt middle1(symbol, symbol);
+    equal_exprt middle2(symbol, symbol);
+    implies_exprt top(middle1, middle2);
+
+    WHEN("Visiting the expressions with a depth iterator")
+    {
+      std::vector<irep_idt> ids;
+      for(auto it = top.depth_begin(), itend = top.depth_end(); it != itend;
+          ++it)
+      {
+        ids.push_back(it->id());
+      }
+
+      THEN("We expect to see parents before children")
+      {
+        REQUIRE(
+          ids == std::vector<irep_idt>{ID_implies,
+                                       ID_notequal,
+                                       ID_symbol,
+                                       ID_symbol,
+                                       ID_equal,
+                                       ID_symbol,
+                                       ID_symbol});
+      }
+    }
+
+    WHEN("Replacing one of the middle nodes mid-walk")
+    {
+      std::vector<irep_idt> ids;
+      for(auto it = top.depth_begin(), itend = top.depth_end(); it != itend;
+          ++it)
+      {
+        if(it->id() == ID_notequal)
+          it.mutate() = not_exprt(equal_exprt(symbol, symbol));
+
+        ids.push_back(it->id());
+      }
+
+      THEN("We expect to see the newly added child nodes")
+      {
+        REQUIRE(
+          ids == std::vector<irep_idt>{ID_implies,
+                                       ID_not,
+                                       ID_equal,
+                                       ID_symbol,
+                                       ID_symbol,
+                                       ID_equal,
+                                       ID_symbol,
+                                       ID_symbol});
+      }
+    }
+
+    WHEN(
+      "Replacing one of the middle nodes mid-walk and skipping the new "
+      "children")
+    {
+      std::vector<irep_idt> ids;
+      for(auto it = top.depth_begin(), itend = top.depth_end(); it != itend;
+          /* no ++it */)
+      {
+        bool replace_here = it->id() == ID_notequal;
+
+        if(replace_here)
+          it.mutate() = not_exprt(equal_exprt(symbol, symbol));
+
+        ids.push_back(it->id());
+
+        if(replace_here)
+          it.next_sibling_or_parent();
+        else
+          ++it;
+      }
+
+      THEN("We expect to skip the newly added child nodes")
+      {
+        REQUIRE(
+          ids == std::vector<irep_idt>{
+                   ID_implies, ID_not, ID_equal, ID_symbol, ID_symbol});
+      }
+    }
+  }
+}