WIP: flow_control: Support match expressions.

commit-id:df88d67e

Author: liorgold2

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

@@ -1,10 +1,10 @@
-use cairo_lang_semantic::{self as semantic, Condition};
+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 super::graph::{
-    ArmExpr, BooleanIf, EvaluateExpr, FlowControlGraph, FlowControlGraphBuilder, FlowControlNode,
+    ArmExpr, BooleanIf, EvaluateExpr, FlowControlGraph, FlowControlGraphBuilder, FlowControlNode, NodeId,
 };
 use crate::lower::context::LoweringContext;
 
@@ -86,3 +86,54 @@ pub fn create_graph_expr_if<'db>(
 
     graph.finalize(current_node)
 }
+
+/// Creates a graph node for [semantic::ExprMatch].
+#[allow(dead_code)]
+pub fn create_graph_expr_match<'db>(
+    ctx: &LoweringContext<'db, '_>,
+    expr: &semantic::ExprMatch<'db>,
+) -> FlowControlGraph<'db> {
+    let mut graph = FlowControlGraphBuilder::default();
+
+    let matched_expr = &ctx.function_body.arenas.exprs[expr.matched_expr];
+    let matched_expr_location = ctx.get_location(matched_expr.stable_ptr().untyped());
+    let matched_var = graph.new_var(matched_expr.ty(), matched_expr_location);
+
+    // Create a list of patterns and nodes.
+    let pattern_and_nodes: Vec<(PatternId, NodeId)> = expr
+        .arms
+        .iter()
+        .flat_map(|match_arm| {
+            // For each arm, create a node for the arm expression.
+            let arm_node =
+                graph.add_node(FlowControlNode::ArmExpr(ArmExpr { expr: match_arm.expression }));
+            // Then map the patterns to that node.
+            match_arm.patterns.iter().map(move |pattern| (*pattern, arm_node))
+        })
+        .collect();
+
+    // TODO(lior): add diagnostics if there is an unreachable arm.
+    let match_node_id = create_node_for_patterns(
+        ctx,
+        &mut graph,
+        matched_var,
+        &pattern_and_nodes
+            .iter()
+            .map(|(pattern, _)| &ctx.function_body.arenas.patterns[*pattern])
+            .collect_vec(),
+        &|_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
+        },
+        matched_expr_location,
+    );
+
+    let root = graph.add_node(FlowControlNode::EvaluateExpr(EvaluateExpr {
+        expr: expr.matched_expr,
+        var_id: matched_var,
+        next: match_node_id,
+    }));
+
+    graph.finalize(root)
+}

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

@@ -10,7 +10,7 @@ use cairo_lang_utils::extract_matches;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 use itertools::Itertools;
 
-use super::create_graph::create_graph_expr_if;
+use super::create_graph::{create_graph_expr_if, create_graph_expr_match};
 use super::graph::FlowControlGraph;
 use super::lower_graph::lower_graph;
 use crate::Lowered;
@@ -27,6 +27,7 @@ cairo_lang_test_utils::test_file_test!(
     "src/lower/flow_control/test_data",
     {
         if_: "if",
+        match_: "match",
     },
     test_create_graph
 );
@@ -67,6 +68,7 @@ fn test_create_graph(
 
     let graph = match &expr {
         semantic::Expr::If(expr) => create_graph_expr_if(&ctx, expr),
+        semantic::Expr::Match(expr) => create_graph_expr_match(&ctx, expr),
         _ => {
             panic!("Unsupported expression: {:?}", expr.debug(&expr_formatter));
         }

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

@@ -180,7 +180,7 @@ End:
 //! > Test if with panic in condition
 
 //! > test_runner_name
-test_create_graph(expect_diagnostics: false, skip_lowering: false)
+test_create_graph(expect_diagnostics: false)
 
 //! > function_code
 fn foo() -> felt252 {
@@ -221,7 +221,7 @@ End:
 //! > Test if with panic in body
 
 //! > test_runner_name
-test_create_graph(expect_diagnostics: false, skip_lowering: false)
+test_create_graph(expect_diagnostics: false)
 
 //! > function_code
 fn foo() -> felt252 {
@@ -283,7 +283,7 @@ End:
 //! > Test if with panic or return in both body and else
 
 //! > test_runner_name
-test_create_graph(expect_diagnostics: false, skip_lowering: false)
+test_create_graph(expect_diagnostics: false)
 
 //! > function_code
 fn foo() -> felt252 {
@@ -340,7 +340,7 @@ End:
 //! > Test if-let with nested enums.
 
 //! > test_runner_name
-test_create_graph(expect_diagnostics: false, skip_lowering: false)
+test_create_graph(expect_diagnostics: false)
 
 //! > function_code
 fn foo(x: Option<Option<felt252>>) -> felt252 {

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

@@ -0,0 +1,90 @@
+//! > Simple match
+
+//! > test_runner_name
+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,
+    }
+}
+
+//! > module_code
+enum Color {
+    Red,
+    Green,
+    Blue,
+    Black,
+    White,
+}
+
+//! > graph
+Root: 4
+0 ArmExpr { expr: ExprId(1) }
+1 ArmExpr { expr: ExprId(2) }
+2 ArmExpr { expr: ExprId(3) }
+3 EnumMatch { matched_var: FlowControlVar { idx: 0 }, variants: (NodeId(0), FlowControlVar { idx: 1 }), (NodeId(0), FlowControlVar { idx: 2 }), (NodeId(1), FlowControlVar { idx: 3 }), (NodeId(2), FlowControlVar { idx: 4 }), (NodeId(2), FlowControlVar { idx: 5 })}
+4 EvaluateExpr { expr: ExprId(0), var_id: FlowControlVar { idx: 0 }, next: NodeId(3) }
+
+//! > semantic_diagnostics
+
+//! > lowered
+Parameters: v0: test::Color
+blk0 (root):
+Statements:
+End:
+  Match(match_enum(v0) {
+    Color::Red(v1) => blk1,
+    Color::Green(v2) => blk2,
+    Color::Blue(v3) => blk3,
+    Color::Black(v4) => blk4,
+    Color::White(v5) => blk5,
+  })
+
+blk1:
+Statements:
+End:
+  Goto(blk7, {})
+
+blk2:
+Statements:
+End:
+  Goto(blk7, {})
+
+blk3:
+Statements:
+  (v7: core::felt252) <- 2
+End:
+  Goto(blk8, {v7 -> v9})
+
+blk4:
+Statements:
+End:
+  Goto(blk6, {})
+
+blk5:
+Statements:
+End:
+  Goto(blk6, {})
+
+blk6:
+Statements:
+  (v6: core::felt252) <- 3
+End:
+  Goto(blk8, {v6 -> v9})
+
+blk7:
+Statements:
+  (v8: core::felt252) <- 1
+End:
+  Goto(blk8, {v8 -> v9})
+
+blk8:
+Statements:
+End:
+  Return(v9)