From 4bd0a46c2936ea77c1f660a2a768b619215cefc6 Mon Sep 17 00:00:00 2001
From: Vern Paxson <vern@corelight.com>
Date: Sun, 14 Jan 2024 09:06:25 -0800
Subject: [PATCH] fixes for ZAM optimization of "switch" statements

---
 src/script_opt/ZAM/Compile.h |  5 ++++-
 src/script_opt/ZAM/Driver.cc | 26 ++++++++++++++++++++------
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/script_opt/ZAM/Compile.h b/src/script_opt/ZAM/Compile.h
index 686c92d40d..9d3744bed7 100644
--- a/src/script_opt/ZAM/Compile.h
+++ b/src/script_opt/ZAM/Compile.h
@@ -110,11 +110,14 @@ private:
     template<typename T>
     using CaseMapsI = std::vector<CaseMapI<T>>;
 
+    template<typename T>
+    void AdjustSwitchTables(CaseMapsI<T>& abstract_cases);
+
     template<typename T>
     void ConcretizeSwitchTables(const CaseMapsI<T>& abstract_cases, CaseMaps<T>& concrete_cases);
 
     template<typename T>
-    void DumpCases(const T& cases, const char* type_name) const;
+    void DumpCases(const CaseMaps<T>& cases, const char* type_name) const;
     void DumpInsts1(const FrameReMap* remappings);
 
 #include "zeek/ZAM-MethodDecls.h"
diff --git a/src/script_opt/ZAM/Driver.cc b/src/script_opt/ZAM/Driver.cc
index 5840d56a9f..e165ac2f54 100644
--- a/src/script_opt/ZAM/Driver.cc
+++ b/src/script_opt/ZAM/Driver.cc
@@ -264,6 +264,20 @@ void ZAMCompiler::AdjustBranches() {
         if ( auto t = inst->target )
             inst->target = FindLiveTarget(t);
     }
+
+    // Fix up the implicit branches in switches, too.
+    AdjustSwitchTables(int_casesI);
+    AdjustSwitchTables(uint_casesI);
+    AdjustSwitchTables(double_casesI);
+    AdjustSwitchTables(str_casesI);
+}
+
+template<typename T>
+void ZAMCompiler::AdjustSwitchTables(CaseMapsI<T>& abstract_cases) {
+    for ( auto& targs : abstract_cases ) {
+        for ( auto& targ : targs )
+            targ.second = FindLiveTarget(targ.second);
+    }
 }
 
 void ZAMCompiler::RetargetBranches() {
@@ -386,14 +400,14 @@ void ZAMCompiler::Dump() {
         inst->Dump(&frame_denizens, remappings);
     }
 
-    DumpCases(int_casesI, "int");
-    DumpCases(uint_casesI, "uint");
-    DumpCases(double_casesI, "double");
-    DumpCases(str_casesI, "str");
+    DumpCases(int_cases, "int");
+    DumpCases(uint_cases, "uint");
+    DumpCases(double_cases, "double");
+    DumpCases(str_cases, "str");
 }
 
 template<typename T>
-void ZAMCompiler::DumpCases(const T& cases, const char* type_name) const {
+void ZAMCompiler::DumpCases(const CaseMaps<T>& cases, const char* type_name) const {
     for ( auto i = 0U; i < cases.size(); ++i ) {
         printf("%s switch table #%d:", type_name, i);
         for ( auto& m : cases[i] ) {
@@ -404,7 +418,7 @@ void ZAMCompiler::DumpCases(const T& cases, const char* type_name) const {
                                 std::is_same_v<T, double> )
                 case_val = std::to_string(m.first);
 
-            printf(" %s->%d", case_val.c_str(), m.second->inst_num);
+            printf(" %s->%d", case_val.c_str(), m.second);
         }
         printf("\n");
     }