Exemplify IR::Traversal on few places in testgen

Signed-off-by: Vladimír Štill <vladimir.still@intel.com>

backends/p4tools/modules/testgen/core/small_step/abstract_stepper.cpp

@@ -14,6 +14,7 @@
 #include "ir/dump.h"
 #include "ir/id.h"
 #include "ir/indexed_vector.h"
+#include "ir/ir-traversal.h"
 #include "ir/irutils.h"
 #include "lib/cstring.h"
 #include "lib/exceptions.h"
@@ -171,10 +172,10 @@ bool AbstractStepper::stepToStructSubexpr(
     return stepToSubexpr(
         nonValueComponent, result, state,
         [nonValueComponentIdx, rebuildCmd, subexpr](const Continuation::Parameter *v) {
-            auto *result = subexpr->clone();
-            auto *namedClone = result->components[nonValueComponentIdx]->clone();
-            namedClone->expression = v->param;
-            result->components[nonValueComponentIdx] = namedClone;
+            auto *result = IR::Traversal::apply(subexpr, &IR::StructExpression::components,
+                                                IR::Traversal::Index(nonValueComponentIdx),
+                                                &IR::NamedExpression::expression,
+                                                IR::Traversal::Assign(v->param));
             return rebuildCmd(result);
         });
 }

backends/p4tools/modules/testgen/core/small_step/expr_stepper.cpp

@@ -12,6 +12,7 @@
 #include "backends/p4tools/common/lib/trace_event_types.h"
 #include "backends/p4tools/common/lib/variables.h"
 #include "ir/declaration.h"
+#include "ir/ir-traversal.h"
 #include "ir/irutils.h"
 #include "ir/node.h"
 #include "ir/solver.h"
@@ -137,23 +138,22 @@ bool ExprStepper::resolveMethodCallArguments(const IR::MethodCallExpression *cal
         // the argument is not yet symbolic, try to resolve it.
         return stepToSubexpr(
             argExpr, result, state, [call, idx, param](const Continuation::Parameter *v) {
-                // TODO: It seems expensive to copy the function every time we resolve an argument.
+                // TODO: It seems expensive to copy the function call every time we resolve an
+                // argument.
                 // We should do this all at once. But how?
                 // This is the same problem as in stepToListSubexpr
                 // Thankfully, most method calls have less than 10 arguments.
-                auto *clonedCall = call->clone();
-                auto *arguments = clonedCall->arguments->clone();
-                auto *arg = arguments->at(idx)->clone();
-                const IR::Expression *computedExpr = v->param;
-                // A parameter with direction InOut might be read and also written to.
-                // We capture this ambiguity with an InOutReference.
-                if (param->direction == IR::Direction::InOut) {
-                    auto stateVar = ToolsVariables::convertReference(arg->expression);
-                    computedExpr = new IR::InOutReference(stateVar, computedExpr);
-                }
-                arg->expression = computedExpr;
-                (*arguments)[idx] = arg;
-                clonedCall->arguments = arguments;
+                auto *clonedCall = IR::Traversal::apply(
+                    call, &IR::MethodCallExpression::arguments, IR::Traversal::Index(idx),
+                    &IR::Argument::expression, [&](IR::Expression *expr) -> const IR::Expression * {
+                        // A parameter with direction InOut might be read and also written to.
+                        // We capture this ambiguity with an InOutReference.
+                        if (param->direction == IR::Direction::InOut) {
+                            auto stateVar = ToolsVariables::convertReference(expr);
+                            return new IR::InOutReference(stateVar, v->param);
+                        }
+                        return v->param;
+                    });
                 return Continuation::Return(clonedCall);
             });
     }

backends/p4tools/modules/testgen/core/small_step/extern_stepper.cpp

@@ -14,6 +14,7 @@
 #include "backends/p4tools/common/lib/trace_event_types.h"
 #include "backends/p4tools/common/lib/variables.h"
 #include "ir/id.h"
+#include "ir/ir-traversal.h"
 #include "ir/ir.h"
 #include "ir/irutils.h"
 #include "ir/vector.h"
@@ -32,6 +33,14 @@
 
 namespace P4::P4Tools::P4Testgen {
 
+/// Replace argument of given index in a method call with given value.
+const IR::MethodCallExpression *replaceCallArg(const IR::MethodCallExpression *call, size_t argIdx,
+                                               const IR::Expression *replacement) {
+    return IR::Traversal::apply(call, &IR::MethodCallExpression::arguments,
+                                IR::Traversal::Index(argIdx), &IR::Argument::expression,
+                                IR::Traversal::Assign(replacement));
+}
+
 std::vector<std::pair<IR::StateVariable, const IR::Expression *>> ExprStepper::setFields(
     ExecutionState &nextState, const std::vector<IR::StateVariable> &flatFields,
     int varBitFieldSize) {
@@ -392,12 +401,8 @@ const ExprStepper::ExternMethodImpls<ExprStepper> ExprStepper::CORE_EXTERN_METHO
          if (!SymbolicEnv::isSymbolicValue(advanceExpr)) {
              stepToSubexpr(advanceExpr, stepper.result, stepper.state,
                            [&externInfo](const Continuation::Parameter *v) {
-                               auto *clonedCall = externInfo.originalCall.clone();
-                               auto *arguments = clonedCall->arguments->clone();
-                               auto *arg = arguments->at(0)->clone();
-                               arg->expression = v->param;
-                               (*arguments)[0] = arg;
-                               clonedCall->arguments = arguments;
+                               const auto *clonedCall =
+                                   replaceCallArg(&externInfo.originalCall, 0, v->param);
                                return Continuation::Return(clonedCall);
                            });
              return;
@@ -541,12 +546,8 @@ const ExprStepper::ExternMethodImpls<ExprStepper> ExprStepper::CORE_EXTERN_METHO
          if (!SymbolicEnv::isSymbolicValue(varbitExtractExpr)) {
              stepToSubexpr(varbitExtractExpr, stepper.result, stepper.state,
                            [&externInfo](const Continuation::Parameter *v) {
-                               auto *clonedCall = externInfo.originalCall.clone();
-                               auto *arguments = clonedCall->arguments->clone();
-                               auto *arg = arguments->at(1)->clone();
-                               arg->expression = v->param;
-                               (*arguments)[1] = arg;
-                               clonedCall->arguments = arguments;
+                               const auto *clonedCall =
+                                   replaceCallArg(&externInfo.originalCall, 1, v->param);
                                return Continuation::Return(clonedCall);
                            });
              return;