Optimize Interface with a bitmask to avoid re-visiting expressions

src/back/msl/writer.rs

@@ -70,6 +70,7 @@ pub struct Writer<W> {
     out: W,
     names: FastHashMap<NameKey, String>,
     named_expressions: BitSet,
+    visit_mask: BitSet,
     typifier: Typifier,
     namer: Namer,
     temp_bake_handles: Vec<Handle<crate::Expression>>,
@@ -177,6 +178,7 @@ impl<W: Write> Writer<W> {
             out,
             names: FastHashMap::default(),
             named_expressions: BitSet::new(),
+            visit_mask: BitSet::new(),
             typifier: Typifier::new(),
             namer: Namer::default(),
             temp_bake_handles: Vec::new(),
@@ -709,6 +711,7 @@ impl<W: Write> Writer<W> {
         exclude_root: bool,
     ) -> Result<(), Error> {
         // set up the search
+        self.visit_mask.clear();
         let mut interface = Interface {
             expressions: &context.expression.function.expressions,
             local_variables: &context.expression.function.local_variables,
@@ -722,6 +725,7 @@ impl<W: Write> Writer<W> {
                     None
                 },
             },
+            mask: &mut self.visit_mask,
         };
         // populate the bake handles
         interface.traverse_expr(root_handle);

src/front/global_usage.rs

@@ -2,6 +2,7 @@ use crate::{
     arena::{Arena, Handle},
     proc::{Interface, Visitor},
 };
+use bit_set::BitSet;
 
 struct GlobalUseVisitor<'a> {
     usage: &'a mut [crate::GlobalUse],
@@ -30,6 +31,7 @@ impl Visitor for GlobalUseVisitor<'_> {
 
 impl crate::Function {
     pub fn fill_global_use(&mut self, globals_num: usize, functions: &Arena<crate::Function>) {
+        let mut mask = BitSet::with_capacity(self.expressions.len());
         self.global_usage.clear();
         self.global_usage
             .resize(globals_num, crate::GlobalUse::empty());
@@ -41,6 +43,7 @@ impl crate::Function {
                 usage: &mut self.global_usage,
                 functions,
             },
+            mask: &mut mask,
         };
         io.traverse(&self.body);
     }

src/proc/call_graph.rs

@@ -3,6 +3,7 @@ use crate::{
     proc::{Interface, Visitor},
     Function,
 };
+use bit_set::BitSet;
 use petgraph::{
     graph::{DefaultIx, NodeIndex},
     Graph,
@@ -18,6 +19,7 @@ impl<'a> CallGraphBuilder<'a> {
     pub fn process(&self, func: &Function) -> CallGraph {
         let mut graph = Graph::new();
         let mut children = Vec::new();
+        let mut mask = BitSet::with_capacity(func.expressions.len());
 
         let visitor = CallGraphVisitor {
             children: &mut children,
@@ -27,29 +29,38 @@ impl<'a> CallGraphBuilder<'a> {
             expressions: &func.expressions,
             local_variables: &func.local_variables,
             visitor,
+            mask: &mut mask,
         };
 
         interface.traverse(&func.body);
 
         for handle in children {
             let id = graph.add_node(handle);
-            self.collect(handle, id, &mut graph);
+            self.collect(handle, id, &mut graph, &mut mask);
         }
 
         graph
     }
 
-    fn collect(&self, handle: Handle<Function>, id: NodeIndex<DefaultIx>, graph: &mut CallGraph) {
+    fn collect(
+        &self,
+        handle: Handle<Function>,
+        id: NodeIndex<DefaultIx>,
+        graph: &mut CallGraph,
+        mask: &mut BitSet,
+    ) {
         let mut children = Vec::new();
         let visitor = CallGraphVisitor {
             children: &mut children,
         };
         let func = &self.functions[handle];
+        mask.clear();
 
         let mut interface = Interface {
             expressions: &func.expressions,
             local_variables: &func.local_variables,
             visitor,
+            mask,
         };
 
         interface.traverse(&func.body);
@@ -58,7 +69,7 @@ impl<'a> CallGraphBuilder<'a> {
             let child_id = graph.add_node(handle);
             graph.add_edge(id, child_id, ());
 
-            self.collect(handle, child_id, graph);
+            self.collect(handle, child_id, graph, mask);
         }
     }
 }

src/proc/interface.rs

@@ -1,9 +1,11 @@
 use crate::arena::{Arena, Handle};
+use bit_set::BitSet;
 
 pub struct Interface<'a, T> {
+    pub visitor: T,
     pub expressions: &'a Arena<crate::Expression>,
     pub local_variables: &'a Arena<crate::LocalVariable>,
-    pub visitor: T,
+    pub mask: &'a mut BitSet,
 }
 
 pub trait Visitor {
@@ -12,13 +14,13 @@ pub trait Visitor {
     fn visit_fun(&mut self, _: Handle<crate::Function>) {}
 }
 
-impl<'a, T> Interface<'a, T>
-where
-    T: Visitor,
-{
+impl<'a, T: Visitor> Interface<'a, T> {
     pub fn traverse_expr(&mut self, handle: Handle<crate::Expression>) {
         use crate::Expression as E;
 
+        if !self.mask.insert(handle.index()) {
+            return;
+        }
         let expr = &self.expressions[handle];
 
         self.visitor.visit_expr(handle, expr);