checking CFT info when doing optimization pruning

scripts/base/protocols/krb/main.zeek

@@ -190,8 +190,7 @@ event krb_as_response(c: connection, msg: KDC_Response) &priority=-5
 
 event krb_ap_request(c: connection, ticket: KRB::Ticket, opts: KRB::AP_Options) &priority=5
 	{
-	if ( set_session(c) )
-		return;
+	set_session(c);
 	}
 
 event krb_tgs_request(c: connection, msg: KDC_Request) &priority=5

src/script_opt/ZAM/AM-Opt.cc

@@ -941,17 +941,55 @@ void ZAMCompiler::KillInst(zeek_uint_t i) {
         }
     }
 
-    if ( num_labels == 0 )
-        // No labels to propagate.
-        return;
+    ZInstI* succ = nullptr;
 
+    if ( num_labels > 0 ) {
         for ( auto j = i + 1; j < insts1.size(); ++j ) {
-        auto succ = insts1[j];
-        if ( succ->live ) {
-            succ->num_labels += num_labels;
+            if ( insts1[j]->live ) {
+                succ = insts1[j];
                 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_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 ( 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);
+            }
+        }
+
+        // If's can be killed because their bodies become empty,
+        // and break's because they just lead to their following instruction.
+        // However, loop's and next's should not be killed.
+        ASSERT(cft.count(CFT_LOOP) == 0);
+        ASSERT(cft.count(CFT_LOOP_COND) == 0);
+        ASSERT(cft.count(CFT_NEXT) == 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) {
-    for ( auto& gi : gotos.back() )
+void ZAMCompiler::ResolveGoTos(GoToSets& gotos, const InstLabel l, ControlFlowType cft) {
+    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

@@ -167,14 +167,8 @@ private:
     void PushFallThroughs() { PushGoTos(fallthroughs); }
     void PushCatchReturns() { PushGoTos(catches); }
 
-    void ResolveNexts(const InstLabel l) {
-        ResolveGoTos(nexts, l);
-        AddCFT(l, CFT_NEXT);
-    }
-    void ResolveBreaks(const InstLabel l) {
-        ResolveGoTos(breaks, l);
-        AddCFT(l, CFT_BREAK);
-    }
+    void ResolveNexts(const InstLabel l) { ResolveGoTos(nexts, l, CFT_NEXT); }
+    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); }
 
@@ -280,7 +274,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();

src/script_opt/ZAM/Low-Level.cc

@@ -27,7 +27,14 @@ const ZAMStmt ZAMCompiler::LastInst() { return ZAMStmt(insts1.size() - 1); }
 void ZAMCompiler::AddCFT(ZInstI* inst, ControlFlowType cft) {
     if ( ! inst->aux )
         inst->aux = new ZInstAux(0);
-    inst->aux->cft.insert(cft);
+
+    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_entry->second;
+    }
 }
 
 OpaqueVals* ZAMCompiler::BuildVals(const ListExprPtr& l) { return new OpaqueVals(InternalBuildVals(l.get())); }

src/script_opt/ZAM/ZInst.h

@@ -385,6 +385,8 @@ enum ControlFlowType {
     CFT_LOOP_COND,
     CFT_NEXT,
     CFT_BREAK,
+
+    CFT_NONE,
 };
 
 // Auxiliary information, used when the fixed ZInst layout lacks
@@ -506,7 +508,7 @@ public:
     bool is_BiF_call = false;
 
     // Associated control flow information.
-    std::set<ControlFlowType> cft;
+    std::map<ControlFlowType, int> cft;
 
     // Used for referring to events.
     EventHandler* event_handler = nullptr;