add tracking of control flow information

src/script_opt/ZAM/AM-Opt.cc

@@ -962,17 +962,76 @@ void ZAMCompiler::KillInst(zeek_uint_t i) {
         }
     }
 
-    if ( num_labels == 0 )
+    ZInstI* succ = nullptr;
-        // No labels to propagate.
-        return;
 
+    if ( num_labels > 0 ) {
         for ( auto j = i + 1; j < insts1.size(); ++j ) {
-        auto succ = insts1[j];
+            if ( insts1[j]->live ) {
-        if ( succ->live ) {
+                succ = insts1[j];
-            succ->num_labels += num_labels;
                 break;
             }
         }
+        if ( succ )
+            succ->num_labels += num_labels;
+    }
+
+    // Look into propagating control flow info.
+    if ( inst->aux && ! inst->aux->cft.empty() ) {
+        auto& cft = inst->aux->cft;
+
+        if ( cft.count(CFT_ELSE) > 0 ) {
+            // Push forward unless this was the end of the block.
+            if ( cft.count(CFT_BLOCK_END) == 0 ) {
+                ASSERT(succ);
+                AddCFT(succ, CFT_ELSE);
+            }
+            else
+                // But if it *was* the end of the block, remove that block.
+                --cft[CFT_BLOCK_END];
+        }
+
+        if ( cft.count(CFT_BREAK) > 0 ) {
+            // ### Factor this with the following
+            // Propagate breaks 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_BREAK);
+
+            auto be_cnt = cft[CFT_BREAK];
+            --be_cnt; // we already did one above
+            insts1[j]->aux->cft[CFT_BREAK] += be_cnt;
+        }
+
+        if ( cft.count(CFT_BLOCK_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_BLOCK_END);
+
+            auto be_cnt = cft[CFT_BLOCK_END];
+            --be_cnt; // we already did one above
+            insts1[j]->aux->cft[CFT_BLOCK_END] += be_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.
+        // However, loop's and next's should not be killed.
+        ASSERT(cft.count(CFT_LOOP) == 0);
+        ASSERT(cft.count(CFT_LOOP_COND) == 0);
+    }
 }
 
 void ZAMCompiler::KillInsts(zeek_uint_t i) {

src/script_opt/ZAM/Branches.cc

@@ -10,9 +10,12 @@ namespace zeek::detail {
 
 void ZAMCompiler::PushGoTos(GoToSets& gotos) { gotos.emplace_back(); }
 
-void ZAMCompiler::ResolveGoTos(GoToSets& gotos, const InstLabel l) {
+void ZAMCompiler::ResolveGoTos(GoToSets& gotos, const InstLabel l, ControlFlowType cft) {
-    for ( auto& gi : gotos.back() )
+    for ( auto& gi : gotos.back() ) {
         SetGoTo(gi, l);
+        if ( cft != CFT_NONE )
+            AddCFT(insts1[gi.stmt_num], cft);
+    }
 
     gotos.pop_back();
 }

src/script_opt/ZAM/Compile.h

@@ -159,16 +159,17 @@ private:
 
     const ZAMStmt ValueSwitch(const SwitchStmt* sw, const NameExpr* v, const ConstExpr* c);
     const ZAMStmt TypeSwitch(const SwitchStmt* sw, const NameExpr* v, const ConstExpr* c);
+    const ZAMStmt GenSwitch(const SwitchStmt* sw, int slot, InternalTypeTag it);
 
     void PushNexts() { PushGoTos(nexts); }
     void PushBreaks() { PushGoTos(breaks); }
     void PushFallThroughs() { PushGoTos(fallthroughs); }
     void PushCatchReturns() { PushGoTos(catches); }
 
-    void ResolveNexts(const InstLabel l) { ResolveGoTos(nexts, l); }
+    void ResolveNexts(const InstLabel l) { ResolveGoTos(nexts, l, CFT_NEXT); }
-    void ResolveBreaks(const InstLabel l) { ResolveGoTos(breaks, l); }
+    void ResolveBreaks(const InstLabel l) { ResolveGoTos(breaks, l, CFT_BREAK); }
     void ResolveFallThroughs(const InstLabel l) { ResolveGoTos(fallthroughs, l); }
-    void ResolveCatchReturns(const InstLabel l) { ResolveGoTos(catches, l); }
+    void ResolveCatchReturns(const InstLabel l) { ResolveGoTos(catches, l, CFT_INLINED_RETURN); }
 
     const ZAMStmt LoopOverTable(const ForStmt* f, const NameExpr* val);
     const ZAMStmt LoopOverVector(const ForStmt* f, const NameExpr* val);
@@ -229,8 +230,8 @@ private:
     const ZAMStmt CompileIndex(const NameExpr* n1, int n2_slot, const TypePtr& n2_type, const ListExpr* l,
                                bool in_when);
 
-    const ZAMStmt BuildLambda(const NameExpr* n, ExprPtr le); // marker
+    const ZAMStmt BuildLambda(const NameExpr* n, ExprPtr le);
-    const ZAMStmt BuildLambda(int n_slot, ExprPtr le);        // marker
+    const ZAMStmt BuildLambda(int n_slot, ExprPtr le);
 
     // Second argument is which instruction slot holds the branch target.
     const ZAMStmt GenCond(const Expr* e, int& branch_v);
@@ -277,7 +278,7 @@ private:
     using GoToSets = std::vector<GoToSet>;
 
     void PushGoTos(GoToSets& gotos);
-    void ResolveGoTos(GoToSets& gotos, const InstLabel l);
+    void ResolveGoTos(GoToSets& gotos, const InstLabel l, ControlFlowType cft = CFT_NONE);
 
     ZAMStmt GenGoTo(GoToSet& v);
     ZAMStmt GoToStub();
@@ -322,6 +323,9 @@ private:
     const ZAMStmt ErrorStmt();
     const ZAMStmt LastInst();
 
+    // Adds control flow information to an instruction.
+    void AddCFT(ZInstI* inst, ControlFlowType cft);
+
     // Returns a handle to state associated with building
     // up a list of values.
     std::unique_ptr<OpaqueVals> BuildVals(const ListExprPtr&);

src/script_opt/ZAM/Low-Level.cc

@@ -20,9 +20,25 @@ bool ZAMCompiler::NullStmtOK() const {
 
 const ZAMStmt ZAMCompiler::EmptyStmt() { return ZAMStmt(insts1.size() - 1); }
 
+const ZAMStmt ZAMCompiler::ErrorStmt() { return ZAMStmt(0); }
+
 const ZAMStmt ZAMCompiler::LastInst() { return ZAMStmt(insts1.size() - 1); }
 
-const ZAMStmt ZAMCompiler::ErrorStmt() { return ZAMStmt(0); }
+void ZAMCompiler::AddCFT(ZInstI* inst, ControlFlowType cft) {
+    if ( cft == CFT_NONE )
+        return;
+
+    if ( ! inst->aux )
+        inst->aux = new ZInstAux(0);
+
+    auto cft_entry = inst->aux->cft.find(cft);
+    if ( cft_entry == inst->aux->cft.end() )
+        inst->aux->cft[cft] = 1;
+    else {
+        ASSERT(cft == CFT_BLOCK_END || cft == CFT_BREAK);
+        ++cft_entry->second;
+    }
+}
 
 std::unique_ptr<OpaqueVals> ZAMCompiler::BuildVals(const ListExprPtr& l) {
     return std::make_unique<OpaqueVals>(InternalBuildVals(l.get()));

src/script_opt/ZAM/Stmt.cc

@@ -141,13 +141,21 @@ const ZAMStmt ZAMCompiler::IfElse(const Expr* e, const Stmt* s1, const Stmt* s2)
     else
         cond_stmt = GenCond(e, branch_v);
 
+    AddCFT(insts1.back(), CFT_IF);
+
     if ( s1 ) {
         auto s1_end = CompileStmt(s1);
+        AddCFT(insts1.back(), CFT_BLOCK_END);
+
         if ( s2 ) {
             auto branch_after_s1 = GoToStub();
+            auto else_start = insts1.size();
             auto s2_end = CompileStmt(s2);
+
             SetV(cond_stmt, GoToTargetBeyond(branch_after_s1), branch_v);
             SetGoTo(branch_after_s1, GoToTargetBeyond(s2_end));
+            AddCFT(insts1[else_start], CFT_ELSE);
+            AddCFT(insts1.back(), CFT_BLOCK_END);
 
             return s2_end;
         }
@@ -462,6 +470,7 @@ const ZAMStmt ZAMCompiler::GenSwitch(const SwitchStmt* sw, int slot, InternalTyp
     // Generate each of the cases.
     auto cases = sw->Cases();
     std::vector<InstLabel> case_start;
+    int case_index = 0;
 
     PushFallThroughs();
     for ( auto sw_case : *cases ) {
@@ -477,8 +486,11 @@ const ZAMStmt ZAMCompiler::GenSwitch(const SwitchStmt* sw, int slot, InternalTyp
     ResolveBreaks(sw_end);
 
     int def_ind = sw->DefaultCaseIndex();
-    if ( def_ind >= 0 )
+    if ( def_ind >= 0 ) {
-        SetV3(sw_head, case_start[def_ind]);
+        auto def = case_start[def_ind];
+        SetV3(sw_head, def);
+        AddCFT(def, CFT_DEFAULT);
+    }
     else
         SetV3(sw_head, sw_end);
 
@@ -651,18 +663,23 @@ const ZAMStmt ZAMCompiler::While(const Stmt* cond_stmt, const Expr* cond, const
 
     if ( cond->Tag() == EXPR_NAME ) {
         auto n = cond->AsNameExpr();
-        cond_IF = AddInst(ZInstI(OP_IF_VV, FrameSlot(n), 0));
+        cond_IF = AddInst(ZInstI(OP_IF_Vb, FrameSlot(n), 0));
         branch_v = 2;
     }
     else
         cond_IF = GenCond(cond, branch_v);
 
+    AddCFT(insts1[head.stmt_num], CFT_LOOP);
+    AddCFT(insts1[cond_IF.stmt_num], CFT_LOOP_COND);
+
     PushNexts();
     PushBreaks();
 
     if ( body && body->Tag() != STMT_NULL )
         (void)CompileStmt(body);
 
+    AddCFT(insts1.back(), CFT_BLOCK_END);
+
     auto tail = GoTo(GoToTarget(head));
 
     auto beyond_tail = GoToTargetBeyond(tail);
@@ -682,17 +699,21 @@ const ZAMStmt ZAMCompiler::CompileFor(const ForStmt* f) {
     PushNexts();
     PushBreaks();
 
+    ZAMStmt z;
+
     if ( et == TYPE_TABLE )
-        return LoopOverTable(f, val);
+        z = LoopOverTable(f, val);
 
     else if ( et == TYPE_VECTOR )
-        return LoopOverVector(f, val);
+        z = LoopOverVector(f, val);
 
     else if ( et == TYPE_STRING )
-        return LoopOverString(f, e);
+        z = LoopOverString(f, e);
 
     else
         reporter->InternalError("bad \"for\" loop-over value when compiling");
+
+    return z;
 }
 
 const ZAMStmt ZAMCompiler::LoopOverTable(const ForStmt* f, const NameExpr* val) {

src/script_opt/ZAM/ZInst.h

@@ -385,6 +385,20 @@ private:
     bool is_managed = false;
 };
 
+enum ControlFlowType {
+    CFT_IF,
+    CFT_BLOCK_END,
+    CFT_ELSE,
+    CFT_LOOP,
+    CFT_LOOP_COND,
+    CFT_NEXT,
+    CFT_BREAK,
+    CFT_DEFAULT,
+    CFT_INLINED_RETURN,
+
+    CFT_NONE,
+};
+
 // Auxiliary information, used when the fixed ZInst layout lacks
 // sufficient expressiveness to represent all of the elements that
 // an instruction needs.
@@ -500,11 +514,14 @@ public:
     // Whether we know that we're calling a BiF.
     bool is_BiF_call = false;
 
+    // Associated control flow information.
+    std::map<ControlFlowType, int> cft;
+
     // Used for referring to events.
     EventHandler* event_handler = nullptr;
 
     // Used for things like constructors.
-    AttributesPtr attrs = nullptr;
+    AttributesPtr attrs;
 
     // Whether the instruction can lead to globals/captures changing.
     // Currently only needed by the optimizer, but convenient to