From fa142438fe7912a7d55e01dfef003c838c592737 Mon Sep 17 00:00:00 2001
From: Vern Paxson <vern@corelight.com>
Date: Fri, 7 Jan 2022 14:18:42 -0800
Subject: [PATCH] logic (other than in profiling) for assignments that yield
 separate values

---
 src/Expr.h                  |  9 +++++++++
 src/script_opt/CPP/Exprs.cc | 11 ++++++++++-
 src/script_opt/Expr.cc      | 21 ++++++++++++++++++++-
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/Expr.h b/src/Expr.h
index 83c34616f5..e8863a59b3 100644
--- a/src/Expr.h
+++ b/src/Expr.h
@@ -959,6 +959,15 @@ public:
 	bool IsTemp() const { return is_temp; }
 	void SetIsTemp() { is_temp = true; }
 
+	// The following is a hack that's used in "when" expressions to support
+	// assignments to new locals, like "when ( (local l = foo()) && ...".
+	// These methods return the value to use when evaluating such
+	// assignments.  That would normally be the RHS of the assignment,
+	// but to get when's to work in a convenient fashion, for them it's
+	// instead boolean T.
+	ValPtr AssignVal() { return val; }
+	const ValPtr& AssignVal() const { return val; }
+
 protected:
 	bool TypeCheck(const AttributesPtr& attrs = nullptr);
 	bool TypeCheckArithmetics(TypeTag bt1, TypeTag bt2);
diff --git a/src/script_opt/CPP/Exprs.cc b/src/script_opt/CPP/Exprs.cc
index f65be5eb84..a9c41a8377 100644
--- a/src/script_opt/CPP/Exprs.cc
+++ b/src/script_opt/CPP/Exprs.cc
@@ -461,7 +461,16 @@ string CPPCompile::GenAssignExpr(const Expr* e, GenType gt, bool top_level)
 	if ( rhs_is_any && ! lhs_is_any && t1->Tag() != TYPE_LIST )
 		rhs_native = rhs_val_ptr = GenericValPtrToGT(rhs_val_ptr, t1, GEN_NATIVE);
 
-	return GenAssign(op1, op2, rhs_native, rhs_val_ptr, gt, top_level);
+	auto gen = GenAssign(op1, op2, rhs_native, rhs_val_ptr, gt, top_level);
+	auto av = e->AsAssignExpr()->AssignVal();
+	if ( av )
+		{
+		auto av_e = make_intrusive<ConstExpr>(av);
+		auto av_gen = GenExpr(av_e, gt, false);
+		return string("(") + gen + ", " + av_gen + ")";
+		}
+	else
+		return gen;
 	}
 
 string CPPCompile::GenAddToExpr(const Expr* e, GenType gt, bool top_level)
diff --git a/src/script_opt/Expr.cc b/src/script_opt/Expr.cc
index b69ee6d30d..b4db48a6ec 100644
--- a/src/script_opt/Expr.cc
+++ b/src/script_opt/Expr.cc
@@ -1556,6 +1556,11 @@ bool AssignExpr::IsReduced(Reducer* c) const
 		// Cascaded assignments are never reduced.
 		return false;
 
+	if ( val )
+		// Initializations of "local" variables in "when" statements
+		// are never reduced.
+		return false;
+
 	const auto& t1 = op1->GetType();
 	const auto& t2 = op2->GetType();
 
@@ -1619,6 +1624,16 @@ ExprPtr AssignExpr::Reduce(Reducer* c, StmtPtr& red_stmt)
 		// These are generated for reduced expressions.
 		return ThisPtr();
 
+	if ( val )
+		{
+		// These are reduced to the assignment followed by
+		// the assignment value.
+		auto assign_val = make_intrusive<ConstExpr>(val);
+		val = nullptr;
+		red_stmt = make_intrusive<ExprStmt>(ThisPtr());
+		return assign_val;
+		}
+
 	auto& t1 = op1->GetType();
 	auto& t2 = op2->GetType();
 
@@ -1757,7 +1772,8 @@ ExprPtr AssignExpr::Reduce(Reducer* c, StmtPtr& red_stmt)
 ExprPtr AssignExpr::ReduceToSingleton(Reducer* c, StmtPtr& red_stmt)
 	{
 	// Yields a statement performing the assignment and for the
-	// expression the LHS (but turned into an RHS).
+	// expression the LHS (but turned into an RHS), or the assignment
+	// value if present.
 	if ( op1->Tag() != EXPR_REF )
 		Internal("Confusion in AssignExpr::ReduceToSingleton");
 
@@ -1765,6 +1781,9 @@ ExprPtr AssignExpr::ReduceToSingleton(Reducer* c, StmtPtr& red_stmt)
 	auto ae_stmt = make_intrusive<ExprStmt>(assign_expr);
 	red_stmt = ae_stmt->Reduce(c);
 
+	if ( val )
+		return make_intrusive<ConstExpr>(val);
+
 	return op1->AsRefExprPtr()->GetOp1();
 	}