From 1d43760862de04b51695a0379aa58bc2bfd99b0b Mon Sep 17 00:00:00 2001
From: Vern Paxson <vern@corelight.com>
Date: Tue, 8 Jul 2025 18:18:27 -0700
Subject: [PATCH] ZAM optimizer fix for += / -= set operations

---
 src/script_opt/Expr.cc                        |  6 ++++++
 .../Baseline/opt.regress-set-op-opt/output    |  4 ++++
 testing/btest/opt/regress-set-op-opt.zeek     | 20 +++++++++++++++++++
 3 files changed, 30 insertions(+)
 create mode 100644 testing/btest/Baseline/opt.regress-set-op-opt/output
 create mode 100644 testing/btest/opt/regress-set-op-opt.zeek

diff --git a/src/script_opt/Expr.cc b/src/script_opt/Expr.cc
index 159d0107e7..cb2927c6de 100644
--- a/src/script_opt/Expr.cc
+++ b/src/script_opt/Expr.cc
@@ -767,6 +767,9 @@ bool AddToExpr::IsReduced(Reducer* c) const {
 }
 
 ExprPtr AddToExpr::Reduce(Reducer* c, StmtPtr& red_stmt) {
+    if ( c->Optimizing() )
+        op2 = c->UpdateExpr(op2);
+
     auto tag = op1->GetType()->Tag();
 
     switch ( tag ) {
@@ -869,6 +872,9 @@ bool RemoveFromExpr::IsReduced(Reducer* c) const {
 }
 
 ExprPtr RemoveFromExpr::Reduce(Reducer* c, StmtPtr& red_stmt) {
+    if ( c->Optimizing() )
+        op2 = c->UpdateExpr(op2);
+
     if ( op1->GetType()->Tag() == TYPE_TABLE ) {
         StmtPtr red_stmt1;
         StmtPtr red_stmt2;
diff --git a/testing/btest/Baseline/opt.regress-set-op-opt/output b/testing/btest/Baseline/opt.regress-set-op-opt/output
new file mode 100644
index 0000000000..6102e32795
--- /dev/null
+++ b/testing/btest/Baseline/opt.regress-set-op-opt/output
@@ -0,0 +1,4 @@
+### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
+[my_set={
+
+}]
diff --git a/testing/btest/opt/regress-set-op-opt.zeek b/testing/btest/opt/regress-set-op-opt.zeek
new file mode 100644
index 0000000000..8e11e0725c
--- /dev/null
+++ b/testing/btest/opt/regress-set-op-opt.zeek
@@ -0,0 +1,20 @@
+# @TEST-DOC: Regression test for ZAM optimizer mis-transforming set +=/-= ops
+# @TEST-REQUIRES: test "${ZEEK_USE_CPP}" != "1"
+# @TEST-EXEC: zeek -b -O ZAM %INPUT >output
+# @TEST-EXEC: btest-diff output
+
+type R: record {
+	my_set: set[string];
+};
+
+event zeek_init()
+	{
+	local r1 = R();
+	local r2 = R();
+
+	if ( |r1$my_set| > 0 )
+		print T;
+
+	r2$my_set += r1$my_set;
+	print r2;
+	}