WIP:

commit-id:8cdbe243

Author: liorgold2

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

@@ -1,7 +1,8 @@
 use cairo_lang_semantic::{self as semantic, Condition, PatternId};
 use cairo_lang_syntax::node::TypedStablePtr;
+use filtered_patterns::IndexAndBindings;
 use itertools::Itertools;
-use patterns::{create_node_for_patterns, get_pattern};
+use patterns::{Cache, create_node_for_patterns, get_pattern};
 
 use super::graph::{
     ArmExpr, BooleanIf, EvaluateExpr, FlowControlGraph, FlowControlGraphBuilder, FlowControlNode,
@@ -61,13 +62,25 @@ 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(
                     ctx,
                     &mut graph,
                     expr_var,
                     &patterns.iter().map(|pattern| Some(get_pattern(ctx, *pattern))).collect_vec(),
-                    &|_graph, pattern_indices| {
-                        if pattern_indices.first().is_some() { current_node } else { false_branch }
+                    &|graph, pattern_indices| {
+                        if let Some(index_and_bindings) = pattern_indices.first() {
+                            cache.get_or_compute(
+                                &|graph, index_and_bindings: IndexAndBindings| {
+                                    index_and_bindings.wrap_node(graph, current_node)
+                                },
+                                graph,
+                                index_and_bindings,
+                            )
+                        } else {
+                            false_branch
+                        }
                     },
                     expr_location,
                 );
@@ -110,6 +123,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(
         ctx,
@@ -119,10 +134,17 @@ pub fn create_graph_expr_match<'db>(
             .iter()
             .map(|(pattern, _)| Some(get_pattern(ctx, *pattern)))
             .collect_vec(),
-        &|_graph, pattern_indices| {
+        &|graph, pattern_indices| {
             // TODO(lior): add diagnostics if pattern_indices is empty (instead of `unwrap`).
-            let index = pattern_indices.first().unwrap();
-            pattern_and_nodes[index].1
+            let index_and_bindings = pattern_indices.first().unwrap();
+            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,
+            )
         },
         matched_expr_location,
     );

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

@@ -1,5 +1,10 @@
+use cairo_lang_semantic::PatternVariable;
 use itertools::Itertools;
 
+use crate::lower::flow_control::graph::{
+    BindVar, FlowControlGraphBuilder, FlowControlNode, FlowControlVar, NodeId,
+};
+
 /// The pattern-matching function below take a list of patterns, and depending on the item at
 /// question, construct a filtered list of patterns that are relevant to the item.
 /// This struct represents the indices of those filtered patterns.
@@ -17,23 +22,37 @@ use itertools::Itertools;
 /// For the latter, it is important to return both `0` and `1`, because the arm that will be chosen
 /// depends on the value of `y` (which will be handled by the calling pattern matching function).
 #[derive(Clone, Hash, Eq, PartialEq)]
-pub struct FilteredPatterns {
-    /// The indices of the patterns that are accepted by the filter.
-    filter: Vec<usize>,
+pub struct FilteredPatterns<'db> {
+    /// The indices of the patterns that are accepted by the filter, together with binding
+    /// information.
+    filter: Vec<IndexAndBindings<'db>>,
 }
 
-impl FilteredPatterns {
+impl<'db> FilteredPatterns<'db> {
     pub fn empty() -> Self {
         Self { filter: vec![] }
     }
 
-    pub fn add(&mut self, idx: usize) {
-        self.filter.push(idx);
-    }
-
     /// Returns a [FilteredPatterns] that accepts all patterns (no filtering).
     pub fn all(n_patterns: usize) -> Self {
-        Self { filter: (0..n_patterns).collect_vec() }
+        Self {
+            filter: (0..n_patterns)
+                .map(|idx| IndexAndBindings { index: idx, bindings: Bindings::default() })
+                .collect_vec(),
+        }
+    }
+
+    pub fn all_with_bindings(bindings_vec: impl Iterator<Item = Bindings<'db>>) -> Self {
+        Self {
+            filter: bindings_vec
+                .enumerate()
+                .map(|(index, bindings)| IndexAndBindings { index, bindings })
+                .collect_vec(),
+        }
+    }
+
+    pub fn add(&mut self, idx: usize, bindings: Bindings<'db>) {
+        self.filter.push(IndexAndBindings { index: idx, bindings });
     }
 
     /// Assuming `self` is a [FilteredPatterns] that applies to a *subset* of a list of patterns
@@ -45,14 +64,67 @@ impl FilteredPatterns {
     /// Suppose that `bar` filters this list to only `C`.
     /// `bar` returns the filter `[1]` since it uses its own indexing.
     /// `foo` needs to lift it to `[2]` to return to its caller using `foo`'s indexing.
-    pub fn lift(self, outer_filter: &FilteredPatterns) -> Self {
+    pub fn lift(self, outer_filter: &FilteredPatterns<'db>) -> Self {
         Self {
-            filter: self.filter.into_iter().map(|index| outer_filter.filter[index]).collect_vec(),
+            filter: self
+                .filter
+                .into_iter()
+                .map(|index_and_bindings| index_and_bindings.lift(outer_filter))
+                .collect_vec(),
         }
     }
 
-    /// Returns the first pattern accepted by the filter.
-    pub fn first(self) -> Option<usize> {
+    /// Returns the first index of the filtered patterns.
+    // TODO: rename. Fix doc.
+    pub fn first(self) -> Option<IndexAndBindings<'db>> {
         self.filter.into_iter().next()
     }
+
+    /// Returns an iterator over the indices of the patterns accepted by the filter.
+    pub fn indices<'a>(&'a self) -> impl Iterator<Item = usize> + 'a {
+        self.filter.iter().map(|index_and_bindings| index_and_bindings.index)
+    }
+}
+
+#[derive(Clone, Hash, Eq, PartialEq)]
+pub struct IndexAndBindings<'db> {
+    /// The index of the pattern in the list of patterns.
+    pub index: usize,
+    /// The bindings that should be applied if the pattern is chosen.
+    bindings: Bindings<'db>,
+}
+impl<'db> IndexAndBindings<'db> {
+    /// Lifts the index of the pattern to the outer level.
+    /// See [FilteredPatterns::lift] for more details.
+    fn lift(self, outer_filter: &FilteredPatterns<'db>) -> IndexAndBindings<'db> {
+        IndexAndBindings {
+            index: outer_filter.filter[self.index].index,
+            bindings: self.bindings.union(&outer_filter.filter[self.index].bindings),
+        }
+    }
+
+    // TODO: doc.
+    pub fn wrap_node(self, graph: &mut FlowControlGraphBuilder<'db>, mut node: NodeId) -> NodeId {
+        for (input, output) in self.bindings.bindings {
+            node = graph.add_node(FlowControlNode::BindVar(BindVar { input, output, next: node }));
+        }
+        node
+    }
+}
+
+#[derive(Clone, Default, Hash, Eq, PartialEq)]
+pub struct Bindings<'db> {
+    /// The bindings that should be applied if the pattern is chosen.
+    bindings: Vec<(FlowControlVar, PatternVariable<'db>)>,
+}
+impl<'db> Bindings<'db> {
+    pub fn single(input: FlowControlVar, output: PatternVariable<'db>) -> Self {
+        Self { bindings: vec![(input, output)] }
+    }
+
+    pub fn union(&self, bindings: &Self) -> Self {
+        Self {
+            bindings: self.bindings.iter().chain(bindings.bindings.iter()).cloned().collect_vec(),
+        }
+    }
 }