From d0eb40ac927ef5d84a873d02db5efa3e759496bc Mon Sep 17 00:00:00 2001
From: Vern Paxson <vern@corelight.com>
Date: Mon, 4 Dec 2023 16:50:59 -0800
Subject: [PATCH] improved AST optimizer's analysis of variable usage in
 inlined functions

---
 src/script_opt/GenIDDefs.cc | 40 ++++++++++++++++++++++++++++++++++++-
 src/script_opt/GenIDDefs.h  |  7 ++++---
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/script_opt/GenIDDefs.cc b/src/script_opt/GenIDDefs.cc
index 209176b3d7..836e99586b 100644
--- a/src/script_opt/GenIDDefs.cc
+++ b/src/script_opt/GenIDDefs.cc
@@ -62,7 +62,45 @@ TraversalCode GenIDDefs::PreStmt(const Stmt* s) {
             auto block = cr->Block();
 
             cr_active.push_back(confluence_blocks.size());
-            block->Traverse(this);
+
+            // Confluence for the bodies of catch-return's is a bit complex.
+            // We would like any expressions computed at the outermost level
+            // of the body to be available for script optimization *outside*
+            // the catch-return; this in particular is helpful in optimizing
+            // coalesced event handlers, but has other benefits as well.
+            //
+            // However, if one of the outermost statements executes a "return",
+            // then any outermost expressions computed after it might not
+            // be available. Put another way, the potentially-returning
+            // statement starts a confluence region that runs through the end
+            // of the body.
+            //
+            // To deal with this, we start off without a new confluence block,
+            // but create one upon encountering a statement that could return.
+
+            bool did_confluence = false;
+
+            if ( block->Tag() == STMT_LIST ) {
+                auto prev_stmt = s;
+                auto& stmts = block->AsStmtList()->Stmts();
+                for ( auto& st : stmts ) {
+                    if ( ! did_confluence && st->CouldReturn(false) ) {
+                        StartConfluenceBlock(prev_stmt);
+                        did_confluence = true;
+                    }
+
+                    st->Traverse(this);
+                }
+            }
+            else
+                // If there's just a single statement then there are no
+                // expressions computed subsequent to it that we need to
+                // worry about, so just do ordinary traversal.
+                block->Traverse(this);
+
+            if ( did_confluence )
+                EndConfluenceBlock();
+
             cr_active.pop_back();
 
             auto retvar = cr->RetVar();
diff --git a/src/script_opt/GenIDDefs.h b/src/script_opt/GenIDDefs.h
index bc1c2785bd..a8b05864cb 100644
--- a/src/script_opt/GenIDDefs.h
+++ b/src/script_opt/GenIDDefs.h
@@ -93,9 +93,10 @@ private:
 
     // Stack of confluence blocks corresponding to activate catch-return
     // statements.  We used to stop identifier definitions at these
-    // boundaries, but given there's no confluence (i.e., the body of the
-    // catch-return *will* execute), we can do broader optimization if we
-    // don't treat them as their own (new) confluence blocks.
+    // boundaries, but given there's limited confluence (i.e., the body of
+    // the catch-return *will* execute, up through its first return), we
+    // can do broader optimization if we don't treat them as their own
+    // (new) confluence blocks.
     std::vector<zeek_uint_t> cr_active;
 
     // The following is parallel to confluence_blocks except