From e93db75f78dd7ce13e52d2d136b35bb1fd061e9f Mon Sep 17 00:00:00 2001
From: Vern Paxson <vern@corelight.com>
Date: Mon, 5 Aug 2024 09:21:13 +0100
Subject: [PATCH] fixes for script optimization of coerce-to-any expressions

---
 src/Expr.h             |  3 +++
 src/Stmt.cc            |  3 ++-
 src/script_opt/Expr.cc | 17 +++++++++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/Expr.h b/src/Expr.h
index 73929a2114..ffead932af 100644
--- a/src/Expr.h
+++ b/src/Expr.h
@@ -1642,6 +1642,9 @@ class CoerceToAnyExpr : public UnaryExpr {
 public:
     CoerceToAnyExpr(ExprPtr op);
 
+    bool IsReduced(Reducer* c) const override;
+    ExprPtr Reduce(Reducer* c, StmtPtr& red_stmt) override;
+
 protected:
     ValPtr Fold(Val* v) const override;
 
diff --git a/src/Stmt.cc b/src/Stmt.cc
index a1edb51ee0..01e892ead0 100644
--- a/src/Stmt.cc
+++ b/src/Stmt.cc
@@ -1849,7 +1849,8 @@ void WhenInfo::Build(StmtPtr ws) {
     auto else_branch = timeout_s ? timeout_s : empty;
 
     auto do_bodies = make_intrusive<IfStmt>(two_test, s, else_branch);
-    auto dummy_return = make_intrusive<ReturnStmt>(true_const);
+    auto any_true_const = make_intrusive<CoerceToAnyExpr>(true_const);
+    auto dummy_return = make_intrusive<ReturnStmt>(any_true_const);
 
     auto shebang = make_intrusive<StmtList>(do_test, do_bodies, dummy_return);
 
diff --git a/src/script_opt/Expr.cc b/src/script_opt/Expr.cc
index 353604f1a7..5e6500ae0d 100644
--- a/src/script_opt/Expr.cc
+++ b/src/script_opt/Expr.cc
@@ -3109,6 +3109,23 @@ CoerceToAnyExpr::CoerceToAnyExpr(ExprPtr arg_op) : UnaryExpr(EXPR_TO_ANY_COERCE,
     type = base_type(TYPE_ANY);
 }
 
+bool CoerceToAnyExpr::IsReduced(Reducer* c) const { return HasReducedOps(c); }
+
+ExprPtr CoerceToAnyExpr::Reduce(Reducer* c, StmtPtr& red_stmt) {
+    if ( c->Optimizing() )
+        op = c->UpdateExpr(op);
+
+    red_stmt = nullptr;
+
+    if ( ! op->IsSingleton(c) )
+        op = op->ReduceToSingleton(c, red_stmt);
+
+    if ( c->Optimizing() )
+        return ThisPtr();
+    else
+        return AssignToTemporary(c, red_stmt);
+}
+
 ValPtr CoerceToAnyExpr::Fold(Val* v) const { return {NewRef{}, v}; }
 
 ExprPtr CoerceToAnyExpr::Duplicate() { return SetSucc(new CoerceToAnyExpr(op->Duplicate())); }