flow_control: Support otherwise pattern.

commit-id:cf815f94

Author: liorgold2

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

@@ -1,10 +1,11 @@
 use cairo_lang_semantic::{self as semantic, Condition, PatternId};
 use cairo_lang_syntax::node::TypedStablePtr;
 use itertools::Itertools;
-use patterns::create_node_for_patterns;
+use patterns::{create_node_for_patterns, get_pattern};
 
 use super::graph::{
-    ArmExpr, BooleanIf, EvaluateExpr, FlowControlGraph, FlowControlGraphBuilder, FlowControlNode, NodeId,
+    ArmExpr, BooleanIf, EvaluateExpr, FlowControlGraph, FlowControlGraphBuilder, FlowControlNode,
+    NodeId,
 };
 use crate::lower::context::LoweringContext;
 
@@ -64,10 +65,7 @@ pub fn create_graph_expr_if<'db>(
                     ctx,
                     &mut graph,
                     expr_var,
-                    &patterns
-                        .iter()
-                        .map(|pattern| &ctx.function_body.arenas.patterns[*pattern])
-                        .collect_vec(),
+                    &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 }
                     },
@@ -119,7 +117,7 @@ pub fn create_graph_expr_match<'db>(
         matched_var,
         &pattern_and_nodes
             .iter()
-            .map(|(pattern, _)| &ctx.function_body.arenas.patterns[*pattern])
+            .map(|(pattern, _)| Some(get_pattern(ctx, *pattern)))
             .collect_vec(),
         &|_graph, pattern_indices| {
             // TODO(lior): add diagnostics if pattern_indices is empty (instead of `unwrap`).

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

@@ -47,11 +47,7 @@ impl FilteredPatterns {
     /// `foo` needs to lift it to `[2]` to return to its caller using `foo`'s indexing.
     pub fn lift(self, outer_filter: &FilteredPatterns) -> Self {
         Self {
-            filter: self
-                .filter
-                .into_iter()
-                .map(|index| outer_filter.filter[index])
-                .collect_vec(),
+            filter: self.filter.into_iter().map(|index| outer_filter.filter[index]).collect_vec(),
         }
     }
 

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

@@ -39,6 +39,9 @@ use crate::lower::context::LoweringContext;
 type BuildNodeCallback<'db, 'a> =
     &'a dyn Fn(&mut FlowControlGraphBuilder<'db>, FilteredPatterns) -> NodeId;
 
+/// A thin wrapper around [semantic::Pattern], where `None` represents the `_` pattern.
+pub type Pattern<'a, 'db> = Option<&'a semantic::Pattern<'db>>;
+
 /// Given a list of patterns and the nodes to go to if the pattern matches,
 /// returns a new graph node to handle the patterns.
 ///
@@ -50,17 +53,14 @@ pub fn create_node_for_patterns<'db>(
     ctx: &LoweringContext<'db, '_>,
     graph: &mut FlowControlGraphBuilder<'db>,
     input_var: FlowControlVar,
-    patterns: &[&semantic::Pattern<'db>],
+    patterns: &[Pattern<'_, 'db>],
     build_node_callback: BuildNodeCallback<'db, '_>,
     location: LocationId<'db>,
 ) -> NodeId {
     // 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)) {
         // Call the callback with all patterns accepted.
-        return build_node_callback(
-            graph,
-            FilteredPatterns::all(patterns.len())
-        );
+        return build_node_callback(graph, FilteredPatterns::all(patterns.len()));
     }
 
     let (n_snapshots, long_ty) = peel_snapshots(ctx.db, graph.var_ty(input_var));
@@ -87,7 +87,7 @@ fn create_node_for_enum<'db>(
     input_var: FlowControlVar,
     concrete_enum_id: ConcreteEnumId<'db>,
     n_snapshots: usize,
-    patterns: &[&semantic::Pattern<'db>],
+    patterns: &[Pattern<'_, 'db>],
     build_node_callback: BuildNodeCallback<'db, '_>,
     location: LocationId<'db>,
 ) -> NodeId {
@@ -100,18 +100,28 @@ fn create_node_for_enum<'db>(
     // Maps variant index to the list of the inner patterns.
     // For example, a pattern `A(B(x))` will add the (inner) pattern `B(x)` to the vector at the
     // index of the variant `A`.
-    let mut variant_to_inner_patterns: Vec<Vec<&semantic::Pattern>> = vec![vec![]; concrete_variants.len()];
+    let mut variant_to_inner_patterns: Vec<Vec<Pattern>> = vec![vec![]; concrete_variants.len()];
 
     for (idx, pattern) in patterns.iter().enumerate() {
         match pattern {
-            semantic::Pattern::EnumVariant(PatternEnumVariant { variant, inner_pattern, .. }) => {
+            Some(semantic::Pattern::EnumVariant(PatternEnumVariant {
+                variant,
+                inner_pattern,
+                ..
+            })) => {
                 variant_to_pattern_indices[variant.idx].add(idx);
-                // TODO(lior): Fix the unwrap below.
-                variant_to_inner_patterns[variant.idx].push(
-                    inner_pattern
-                        .map(|inner_pattern| &ctx.function_body.arenas.patterns[inner_pattern])
-                        .unwrap(),
-                );
+                variant_to_inner_patterns[variant.idx]
+                    .push(inner_pattern.map(|inner_pattern| get_pattern(ctx, inner_pattern)));
+            }
+            Some(semantic::Pattern::Otherwise(..)) | None => {
+                // Add `idx` to all the variants.
+                for pattern_indices in variant_to_pattern_indices.iter_mut() {
+                    pattern_indices.add(idx);
+                }
+                // Add the `_` pattern (represented by `None`) to all the variants.
+                for inner_patterns in variant_to_inner_patterns.iter_mut() {
+                    inner_patterns.push(None);
+                }
             }
             _ => todo!("Pattern {:?} is not supported yet.", pattern),
         }
@@ -146,16 +156,26 @@ fn create_node_for_enum<'db>(
 }
 
 /// Returns `true` if the pattern accepts any value (`_` or a variable name).
-fn pattern_is_any(pattern: &semantic::Pattern) -> bool {
+fn pattern_is_any(pattern: &Pattern) -> bool {
     match pattern {
-        semantic::Pattern::Otherwise(..)
-        | semantic::Pattern::Variable(..) => true,
-        semantic::Pattern::Literal(..) |
-        semantic::Pattern::StringLiteral(..) |
-        semantic::Pattern::Struct(..) |
-        semantic::Pattern::Tuple(..) |
-        semantic::Pattern::FixedSizeArray(..) |
-        semantic::Pattern::EnumVariant(..) |
-        semantic::Pattern::Missing(..) => false,
+        Some(semantic_pattern) => match semantic_pattern {
+            semantic::Pattern::Otherwise(..) | semantic::Pattern::Variable(..) => true,
+            semantic::Pattern::Literal(..)
+            | semantic::Pattern::StringLiteral(..)
+            | semantic::Pattern::Struct(..)
+            | semantic::Pattern::Tuple(..)
+            | semantic::Pattern::FixedSizeArray(..)
+            | semantic::Pattern::EnumVariant(..)
+            | semantic::Pattern::Missing(..) => false,
+        },
+        None => true,
     }
 }
+
+/// Returns a reference to a [semantic::Pattern] from a [semantic::PatternId].
+pub fn get_pattern<'db, 'a>(
+    ctx: &'a LoweringContext<'db, '_>,
+    semantic_pattern: semantic::PatternId,
+) -> &'a semantic::Pattern<'db> {
+    &ctx.function_body.arenas.patterns[semantic_pattern]
+}

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

@@ -6,11 +6,9 @@ test_create_graph(expect_diagnostics: false)
 //! > function_code
 fn foo(color: Color) -> felt252 {
     match color {
-        // TODO(lior): Remove some of the `(_)` once supported.
-        Color::Red(_) | Color::Green(_) => 1,
-        Color::Red(_) | Color::Blue(_) => 2,
-        // TODO(lior): Change to `_` once supported.
-        Color::Black(_) | Color::White(_) => 3,
+        Color::Red | Color::Green(_) => 1,
+        Color::Red(_) | Color::Blue => 2,
+        _ => 3,
     }
 }