Clear stack map data structures upon FunctionBuilderContext reuse (b…

…ytecodealliance#8962) We were previously failing to clear `self.stack_map_vars` and `self.stack_map_values`. Clearing `self.dfs` is not strictly necessary, since calling `self.dfs.iter()` internally clears its internal data before iteration begins, but its nice to tidy up anyways. Additionally, match on all the fields of `self` so that if we add more, we get compiler errors to remind us to clear the new fields as well.
ahaoboy · Jul 16, 2024 · 51948ef · 51948ef
1 parent 158a8f5
commit 51948ef
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 5 deletions.
diff --git a/cranelift/codegen/src/traversals.rs b/cranelift/codegen/src/traversals.rs
@@ -52,8 +52,7 @@ impl Dfs {
     /// This iterator can be used to perform either pre- or post-order
     /// traversals, or a combination of the two.
     pub fn iter<'a>(&'a mut self, func: &'a ir::Function) -> DfsIter<'a> {
-        self.seen.clear();
-        self.stack.clear();
+        self.clear();
         if let Some(e) = func.layout.entry_block() {
             self.stack.push((Event::Enter, e));
         }
@@ -73,6 +72,13 @@ impl Dfs {
     pub fn post_order_iter<'a>(&'a mut self, func: &'a ir::Function) -> DfsPostOrderIter<'a> {
         DfsPostOrderIter(self.iter(func))
     }
+
+    /// Clear this DFS, but keep its allocations for future reuse.
+    pub fn clear(&mut self) {
+        let Dfs { stack, seen } = self;
+        stack.clear();
+        seen.clear();
+    }
 }
 
 /// An iterator that yields pairs of `(Event, ir::Block)` items as it performs a

diff --git a/cranelift/frontend/src/frontend.rs b/cranelift/frontend/src/frontend.rs
@@ -66,9 +66,20 @@ impl FunctionBuilderContext {
     }
 
     fn clear(&mut self) {
-        self.ssa.clear();
-        self.status.clear();
-        self.types.clear();
+        let FunctionBuilderContext {
+            ssa,
+            status,
+            types,
+            stack_map_vars,
+            stack_map_values,
+            dfs,
+        } = self;
+        ssa.clear();
+        status.clear();
+        types.clear();
+        stack_map_values.clear();
+        stack_map_vars.clear();
+        dfs.clear();
     }
 
     fn is_empty(&self) -> bool {