From d5078b27226a6fcb5d91f7cda1dd86d4dc2e6978 Mon Sep 17 00:00:00 2001
From: Vern Paxson <vern@corelight.com>
Date: Tue, 11 Jun 2024 14:36:12 -0700
Subject: [PATCH] ZBody/ZInst traversal

---
 src/script_opt/ZAM/ZBody.cc |  5 ++++
 src/script_opt/ZAM/ZInst.cc | 50 +++++++++++++++++++++++++++++++++++++
 src/script_opt/ZAM/ZInst.h  |  5 ++++
 3 files changed, 60 insertions(+)

diff --git a/src/script_opt/ZAM/ZBody.cc b/src/script_opt/ZAM/ZBody.cc
index 0eb283284a..03677d9d42 100644
--- a/src/script_opt/ZAM/ZBody.cc
+++ b/src/script_opt/ZAM/ZBody.cc
@@ -591,6 +591,11 @@ TraversalCode ZBody::Traverse(TraversalCallback* cb) const {
     TraversalCode tc = cb->PreStmt(this);
     HANDLE_TC_STMT_PRE(tc);
 
+    for ( size_t i = 0; i < NumInsts(); ++i ) {
+        tc = insts[i].Traverse(cb);
+        HANDLE_TC_STMT_PRE(tc);
+    }
+
     tc = cb->PostStmt(this);
     HANDLE_TC_STMT_POST(tc);
 }
diff --git a/src/script_opt/ZAM/ZInst.cc b/src/script_opt/ZAM/ZInst.cc
index 26308a6c87..d84a80259c 100644
--- a/src/script_opt/ZAM/ZInst.cc
+++ b/src/script_opt/ZAM/ZInst.cc
@@ -298,6 +298,56 @@ string ZInst::ConstDump() const {
     return d.Description();
 }
 
+TraversalCode ZInst::Traverse(TraversalCallback* cb) const {
+    TraversalCode tc;
+    if ( t ) {
+        tc = t->Traverse(cb);
+        HANDLE_TC_STMT_PRE(tc);
+        if ( t2 ) {
+            tc = t2->Traverse(cb);
+            HANDLE_TC_STMT_PRE(tc);
+        }
+    }
+
+    if ( aux ) {
+        tc = aux->Traverse(cb);
+        HANDLE_TC_STMT_POST(tc);
+    }
+
+    return TC_CONTINUE;
+}
+
+TraversalCode ZInstAux::Traverse(TraversalCallback* cb) const {
+    TraversalCode tc;
+
+    if ( id_val ) {
+        tc = id_val->Traverse(cb);
+        HANDLE_TC_STMT_PRE(tc);
+    }
+
+    if ( attrs ) {
+        tc = attrs->Traverse(cb);
+        HANDLE_TC_STMT_PRE(tc);
+    }
+
+    if ( value_var_type ) {
+        tc = value_var_type->Traverse(cb);
+        HANDLE_TC_STMT_PRE(tc);
+    }
+
+    for ( auto& lvt : loop_var_types ) {
+        tc = lvt->Traverse(cb);
+        HANDLE_TC_STMT_PRE(tc);
+    }
+
+    for ( auto& lvt : loop_var_types ) {
+        tc = lvt->Traverse(cb);
+        HANDLE_TC_STMT_PRE(tc);
+    }
+
+    return TC_CONTINUE;
+}
+
 void ZInstI::Dump(FILE* f, const FrameMap* frame_ids, const FrameReMap* remappings) const {
     int n = NumFrameSlots();
     // fprintf(f, "v%d ", n);
diff --git a/src/script_opt/ZAM/ZInst.h b/src/script_opt/ZAM/ZInst.h
index c0c9bfa002..3772fa97b5 100644
--- a/src/script_opt/ZAM/ZInst.h
+++ b/src/script_opt/ZAM/ZInst.h
@@ -6,6 +6,7 @@
 
 #include "zeek/Desc.h"
 #include "zeek/Func.h"
+#include "zeek/TraverseTypes.h"
 #include "zeek/script_opt/ZAM/BuiltInSupport.h"
 #include "zeek/script_opt/ZAM/Support.h"
 #include "zeek/script_opt/ZAM/ZOp.h"
@@ -109,6 +110,8 @@ public:
     // Returns a string describing the constant.
     std::string ConstDump() const;
 
+    TraversalCode Traverse(TraversalCallback* cb) const;
+
     ZOp op = OP_NOP;
     ZAMOpType op_type = OP_X;
 
@@ -455,6 +458,8 @@ public:
     // Same but for constants.
     void Add(int i, ValPtr c) { elems[i].SetConstant(c); }
 
+    TraversalCode Traverse(TraversalCallback* cb) const;
+
     // Member variables.  We could add accessors for manipulating
     // these (and make the variables private), but for convenience we
     // make them directly available.