flow-control: Add cache node optimization. (#8115)

Author: liorgold2

crates/cairo-lang-lowering/src/lower/flow_control/create_graph.rs

@@ -1,5 +1,7 @@
+use cache::Cache;
 use cairo_lang_semantic::{self as semantic, Condition, PatternId};
 use cairo_lang_syntax::node::TypedStablePtr;
+use filtered_patterns::IndexAndBindings;
 use itertools::Itertools;
 use patterns::{CreateNodeParams, create_node_for_patterns, get_pattern};
 
@@ -9,6 +11,7 @@ use super::graph::{
 };
 use crate::lower::context::LoweringContext;
 
+mod cache;
 mod filtered_patterns;
 mod patterns;
 
@@ -61,6 +64,8 @@ pub fn create_graph_expr_if<'db>(
                 let expr_location = ctx.get_location(expr.stable_ptr().untyped());
                 let expr_var = graph.new_var(expr.ty(), expr_location);
 
+                let mut cache = Cache::default();
+
                 let match_node_id = create_node_for_patterns(
                     CreateNodeParams {
                         ctx,
@@ -69,9 +74,15 @@ pub fn create_graph_expr_if<'db>(
                             .iter()
                             .map(|pattern| Some(get_pattern(ctx, *pattern)))
                             .collect_vec(),
-                        build_node_callback: &|graph, pattern_indices| {
+                        build_node_callback: &mut |graph, pattern_indices| {
                             if let Some(index_and_bindings) = pattern_indices.first() {
-                                index_and_bindings.wrap_node(graph, current_node)
+                                cache.get_or_compute(
+                                    &mut |graph, index_and_bindings: IndexAndBindings| {
+                                        index_and_bindings.wrap_node(graph, current_node)
+                                    },
+                                    graph,
+                                    index_and_bindings,
+                                )
                             } else {
                                 false_branch
                             }
@@ -119,6 +130,8 @@ pub fn create_graph_expr_match<'db>(
         })
         .collect();
 
+    let mut cache = Cache::default();
+
     // TODO(lior): add diagnostics if there is an unreachable arm.
     let match_node_id = create_node_for_patterns(
         CreateNodeParams {
@@ -128,11 +141,17 @@ pub fn create_graph_expr_match<'db>(
                 .iter()
                 .map(|(pattern, _)| Some(get_pattern(ctx, *pattern)))
                 .collect_vec(),
-            build_node_callback: &|graph, pattern_indices| {
+            build_node_callback: &mut |graph, pattern_indices| {
                 // TODO(lior): add diagnostics if pattern_indices is empty (instead of `unwrap`).
                 let index_and_bindings = pattern_indices.first().unwrap();
-                let index = index_and_bindings.index();
-                index_and_bindings.wrap_node(graph, pattern_and_nodes[index].1)
+                cache.get_or_compute(
+                    &mut |graph, index_and_bindings: IndexAndBindings| {
+                        let index = index_and_bindings.index();
+                        index_and_bindings.wrap_node(graph, pattern_and_nodes[index].1)
+                    },
+                    graph,
+                    index_and_bindings,
+                )
             },
             location: matched_expr_location,
         },

crates/cairo-lang-lowering/src/lower/flow_control/create_graph/cache.rs

@@ -0,0 +1,37 @@
+use cairo_lang_utils::unordered_hash_map::UnorderedHashMap;
+
+use crate::lower::flow_control::graph::{FlowControlGraphBuilder, NodeId};
+
+/// Implements a simple memoization mechanism to optimize the flow control graph that is created.
+///
+/// The cache is used before calling `BuildNodeCallback` to avoid creating two nodes that behave
+/// identically.
+pub struct Cache<Input> {
+    /// A map from input to the cached result.
+    cache: UnorderedHashMap<Input, NodeId>,
+}
+impl<Input: std::hash::Hash + Eq + Clone> Cache<Input> {
+    /// Calls the callback if this is the first time the input is seen.
+    /// Returns the previous result, otherwise.
+    pub fn get_or_compute<'db>(
+        &mut self,
+        callback: &mut dyn FnMut(&mut FlowControlGraphBuilder<'db>, Input) -> NodeId,
+        graph: &mut FlowControlGraphBuilder<'db>,
+        input: Input,
+    ) -> NodeId {
+        if let Some(node_id) = self.cache.get(&input) {
+            return *node_id;
+        }
+
+        let node_id = callback(graph, input.clone());
+        assert!(!self.cache.contains_key(&input));
+        self.cache.insert(input, node_id);
+        node_id
+    }
+}
+
+impl<Input: std::hash::Hash + Eq + Clone> std::default::Default for Cache<Input> {
+    fn default() -> Self {
+        Self { cache: Default::default() }
+    }
+}

crates/cairo-lang-lowering/src/lower/flow_control/create_graph/patterns.rs

@@ -10,6 +10,7 @@ use itertools::Itertools;
 use super::super::graph::{
     EnumMatch, FlowControlGraphBuilder, FlowControlNode, FlowControlVar, NodeId,
 };
+use super::cache::Cache;
 use super::filtered_patterns::{Bindings, FilteredPatterns};
 use crate::ids::LocationId;
 use crate::lower::context::LoweringContext;
@@ -39,7 +40,7 @@ use crate::lower::flow_control::graph::Deconstruct;
 /// Finally, the inner pattern-matching function (for `x`) will construct a [EnumMatch] node
 /// that leads to the two nodes returned by the callback.
 type BuildNodeCallback<'db, 'a> =
-    &'a dyn Fn(&mut FlowControlGraphBuilder<'db>, FilteredPatterns) -> NodeId;
+    &'a mut dyn FnMut(&mut FlowControlGraphBuilder<'db>, FilteredPatterns) -> NodeId;
 
 /// A thin wrapper around [semantic::Pattern], where `None` represents the `_` pattern.
 type PatternOption<'a, 'db> = Option<&'a semantic::Pattern<'db>>;
@@ -82,11 +83,17 @@ pub fn create_node_for_patterns<'db>(
         })
         .collect_vec();
 
-    // Wrap `build_node_callback` to add the bindings to the patterns.
-    let build_node_callback = |graph: &mut FlowControlGraphBuilder<'db>,
-                               pattern_indices: FilteredPatterns| {
-        build_node_callback(graph, pattern_indices.add_bindings(bindings.clone()))
-    };
+    let mut cache = Cache::default();
+
+    // Wrap `build_node_callback` to add the bindings to the patterns and cache the result.
+    let mut build_node_callback =
+        |graph: &mut FlowControlGraphBuilder<'db>, pattern_indices: FilteredPatterns| {
+            cache.get_or_compute(
+                build_node_callback,
+                graph,
+                pattern_indices.add_bindings(bindings.clone()),
+            )
+        };
 
     // If all the patterns are catch-all, we do not need to look into `input_var`.
     if patterns.iter().all(|pattern| pattern_is_any(pattern)) {
@@ -99,7 +106,7 @@ pub fn create_node_for_patterns<'db>(
         ctx,
         graph,
         patterns: &patterns,
-        build_node_callback: &build_node_callback,
+        build_node_callback: &mut build_node_callback,
         location,
     };
     match long_ty {
@@ -169,7 +176,7 @@ fn create_node_for_enum<'db>(
                         ctx,
                         graph,
                         patterns: &inner_patterns,
-                        build_node_callback: &|graph, pattern_indices_inner| {
+                        build_node_callback: &mut |graph, pattern_indices_inner| {
                             build_node_callback(graph, pattern_indices_inner.lift(&pattern_indices))
                         },
                         location,
@@ -262,14 +269,14 @@ fn create_node_for_tuple_inner<'db>(
             ctx,
             graph,
             patterns: &patterns_on_current_item,
-            build_node_callback: &|graph, pattern_indices| {
+            build_node_callback: &mut |graph, pattern_indices| {
                 // Call `create_node_for_tuple_inner` recursively to handle the rest of the tuple.
                 create_node_for_tuple_inner(
                     CreateNodeParams {
                         ctx,
                         graph,
                         patterns: &pattern_indices.indices().map(|idx| patterns[idx]).collect_vec(),
-                        build_node_callback: &|graph, pattern_indices_inner| {
+                        build_node_callback: &mut |graph, pattern_indices_inner| {
                             build_node_callback(graph, pattern_indices_inner.lift(&pattern_indices))
                         },
                         location,

crates/cairo-lang-lowering/src/lower/flow_control/test_data/match

@@ -193,15 +193,14 @@ enum Color {
 }
 
 //! > graph
-Root: 7
+Root: 6
 0 ArmExpr { expr: ExprId(1) }
 1 ArmExpr { expr: ExprId(2) }
 2 ArmExpr { expr: ExprId(3) }
 3 EnumMatch { matched_var: v2, variants: (NodeId(1), v8), (NodeId(2), v9), (NodeId(2), v10)}
-4 EnumMatch { matched_var: v2, variants: (NodeId(1), v12), (NodeId(2), v13), (NodeId(2), v14)}
-5 EnumMatch { matched_var: v1, variants: (NodeId(0), v3), (NodeId(3), v7), (NodeId(4), v11)}
-6 Deconstruct { input: v0, outputs: [v1, v2], next: NodeId(5) }
-7 EvaluateExpr { expr: ExprId(0), var_id: v0, next: NodeId(6) }
+4 EnumMatch { matched_var: v1, variants: (NodeId(0), v3), (NodeId(3), v7), (NodeId(3), v11)}
+5 Deconstruct { input: v0, outputs: [v1, v2], next: NodeId(4) }
+6 EvaluateExpr { expr: ExprId(0), var_id: v0, next: NodeId(5) }
 
 //! > semantic_diagnostics