ZAM control-flow tracking now explicitly includes the ends of loops

src/script_opt/ZAM/AM-Opt.cc

@@ -1031,6 +1031,23 @@ void ZAMCompiler::KillInst(zeek_uint_t i) {
             insts1[j]->aux->cft[CFT_BLOCK_END] += be_cnt;
         }
 
+        if ( cft.count(CFT_LOOP_END) > 0 ) {
+            // Propagate block-ends backwards.
+            int j = i;
+            while ( --j >= 0 )
+                if ( insts1[j]->live )
+                    break;
+
+            ASSERT(j >= 0);
+
+            // Make sure the CFT entry is created.
+            AddCFT(insts1[j], CFT_LOOP_END);
+
+            auto le_cnt = cft[CFT_LOOP_END];
+            --le_cnt; // we already did one above
+            insts1[j]->aux->cft[CFT_LOOP_END] += le_cnt;
+        }
+
         // If's can be killed because their bodies become empty,
         // break's because they just lead to their following instruction,
         // and next's if they become dead code.

src/script_opt/ZAM/Low-Level.cc

@@ -35,7 +35,7 @@ void ZAMCompiler::AddCFT(ZInstI* inst, ControlFlowType cft) {
     if ( cft_entry == inst->aux->cft.end() )
         inst->aux->cft[cft] = 1;
     else {
-        ASSERT(cft == CFT_BLOCK_END || cft == CFT_BREAK);
+        ASSERT(cft == CFT_BLOCK_END || cft == CFT_LOOP_END || cft == CFT_BREAK);
         ++cft_entry->second;
     }
 }

src/script_opt/ZAM/Stmt.cc

@@ -650,7 +650,7 @@ const ZAMStmt ZAMCompiler::While(const Stmt* cond_stmt, const Expr* cond, const
     if ( body && body->Tag() != STMT_NULL )
         (void)CompileStmt(body);
 
-    AddCFT(insts1.back(), CFT_BLOCK_END);
+    AddCFT(insts1.back(), CFT_LOOP_END);
 
     auto tail = GoTo(GoToTarget(head));
 
@@ -849,7 +849,7 @@ const ZAMStmt ZAMCompiler::Loop(const Stmt* body) {
     auto b = CompileStmt(body);
 
     AddCFT(insts1[head.stmt_num], CFT_LOOP);
-    AddCFT(insts1[b.stmt_num], CFT_BLOCK_END);
+    AddCFT(insts1[b.stmt_num], CFT_LOOP_END);
 
     auto tail = GoTo(GoToTarget(head));
 
@@ -863,14 +863,15 @@ const ZAMStmt ZAMCompiler::FinishLoop(const ZAMStmt iter_head, ZInstI& iter_stmt
                                       bool is_table) {
     auto loop_iter = AddInst(iter_stmt);
     auto body_end = CompileStmt(body);
-    AddCFT(insts1[body_end.stmt_num], CFT_BLOCK_END);
+
+    auto loop_end = GoTo(GoToTarget(iter_head));
+    AddCFT(insts1[loop_end.stmt_num], CFT_LOOP_END);
 
     // We only need cleanup for looping over tables, but for now we
     // need some sort of placeholder instruction (until the optimizer
     // can elide it) to resolve loop exits.
     ZOp op = is_table ? OP_END_TABLE_LOOP_f : OP_NOP;
 
-    auto loop_end = GoTo(GoToTarget(iter_head));
     auto z = ZInstI(op, iter_slot);
     z.op_type = is_table ? OP_V_I1 : OP_X;
     auto final_stmt = AddInst(z);

src/script_opt/ZAM/ZInst.h

@@ -390,6 +390,7 @@ enum ControlFlowType {
     CFT_ELSE,
     CFT_LOOP,
     CFT_LOOP_COND,
+    CFT_LOOP_END,
     CFT_NEXT,
     CFT_BREAK,
     CFT_DEFAULT,