flow-control: Add cache node optimization.

commit-id:edbf06d2

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 cache = Cache::default();
+
                 let match_node_id = create_node_for_patterns(
                     CreateNodeParams {
                         ctx,
@@ -71,7 +76,13 @@ pub fn create_graph_expr_if<'db>(
                             .collect_vec(),
                         build_node_callback: &|graph, pattern_indices| {
                             if let Some(index_and_bindings) = pattern_indices.first() {
-                                index_and_bindings.wrap_node(graph, current_node)
+                                cache.get_or_compute(
+                                    &|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 cache = Cache::default();
+
     // TODO(lior): add diagnostics if there is an unreachable arm.
     let match_node_id = create_node_for_patterns(
         CreateNodeParams {
@@ -131,8 +144,14 @@ pub fn create_graph_expr_match<'db>(
             build_node_callback: &|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(
+                    &|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,42 @@
+use std::cell::RefCell;
+
+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.
+    //
+    // The cache is wrapped in a `RefCell` to allow modifying it without holding a `&mut` to it
+    // (which would complicate its usage).
+    cache: RefCell<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>(
+        &self,
+        callback: &dyn Fn(&mut FlowControlGraphBuilder<'db>, Input) -> NodeId,
+        graph: &mut FlowControlGraphBuilder<'db>,
+        input: Input,
+    ) -> NodeId {
+        if let Some(node_id) = self.cache.borrow().get(&input) {
+            return *node_id;
+        }
+
+        let node_id = callback(graph, input.clone());
+        assert!(!self.cache.borrow().contains_key(&input));
+        self.cache.borrow_mut().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;
@@ -82,10 +83,16 @@ pub fn create_node_for_patterns<'db>(
         })
         .collect_vec();
 
-    // Wrap `build_node_callback` to add the bindings to the patterns.
+    let cache = Cache::default();
+
+    // Wrap `build_node_callback` to add the bindings to the patterns and cache the result.
     let build_node_callback = |graph: &mut FlowControlGraphBuilder<'db>,
                                pattern_indices: FilteredPatterns| {
-        build_node_callback(graph, pattern_indices.add_bindings(bindings.clone()))
+        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`.

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