propagate variable hits to parent branches in push_missing_vars (#2183)

src/codegen.rs

@@ -61,6 +61,7 @@ impl BranchCodeStack {
         marker: &mut DebrayAllocator,
     ) -> SubsumedBranchHits {
         let mut subsumed_hits = SubsumedBranchHits::with_hasher(FxBuildHasher::default());
+        let mut propagated_var_nums = IndexSet::with_hasher(FxBuildHasher::default());
 
         for idx in (self.stack.len() - depth..self.stack.len()).rev() {
             let branch = &mut marker.branch_stack[idx];
@@ -84,11 +85,19 @@ impl BranchCodeStack {
 
                             self.stack[idx][branch_idx].push_back(instr!("put_variable", r, 0));
                         }
+
+                        if idx > self.stack.len() - depth {
+                            propagated_var_nums.insert(var_num);
+                        }
                     }
 
                     subsumed_hits.insert(var_num);
                 }
             }
+
+            for var_num in propagated_var_nums.drain(..) {
+                marker.branch_stack[idx - 1].add_branch_occurrence(var_num);
+            }
         }
 
         subsumed_hits

src/debray_allocator.rs

@@ -39,6 +39,19 @@ impl BranchOccurrences {
             subsumed_hits: SubsumedBranchHits::with_hasher(FxBuildHasher::default()),
         }
     }
+
+    pub(crate) fn add_branch_occurrence(&mut self, var_num: usize) {
+        debug_assert!(self.current_branch < self.num_branches);
+        let num_branches = self.num_branches;
+
+        let entry = self
+            .hits
+            .entry(var_num)
+            .or_insert_with(|| BitVec::repeat(false, num_branches));
+
+        entry.set(self.current_branch, true);
+        self.subsumed_hits.insert(var_num);
+    }
 }
 
 #[derive(Debug)]
@@ -92,17 +105,7 @@ impl BranchStack {
 
     pub(crate) fn add_branch_occurrence(&mut self, var_num: usize) {
         if let Some(occurrences) = self.last_mut() {
-            debug_assert!(occurrences.current_branch < occurrences.num_branches);
+            occurrences.add_branch_occurrence(var_num);
-
-            let num_branches = occurrences.num_branches;
-
-            let entry = occurrences
-                .hits
-                .entry(var_num)
-                .or_insert_with(|| BitVec::repeat(false, num_branches));
-
-            entry.set(occurrences.current_branch, true);
-            occurrences.subsumed_hits.insert(var_num);
         }
     }