Add an empty merge_block_builders with a test. (#7952)

Author: liorgold2

Cargo.lock

@@ -18,6 +18,7 @@ cairo-lang-syntax = { path = "../cairo-lang-syntax", version = "~2.11.4" }
 cairo-lang-utils = { path = "../cairo-lang-utils", version = "~2.11.4" }
 assert_matches.workspace = true
 id-arena.workspace = true
+indent.workspace = true
 itertools = { workspace = true, default-features = true }
 log.workspace = true
 num-bigint = { workspace = true, default-features = true }

crates/cairo-lang-lowering/Cargo.toml

@@ -102,7 +102,7 @@ impl DebugWithDb<LoweredFormatter<'_>> for BlockEnd {
             BlockEnd::Goto(block_id, remapping) => {
                 return write!(f, "  Goto({:?}, {:?})", block_id.debug(ctx), remapping.debug(ctx));
             }
-            BlockEnd::NotSet => unreachable!(),
+            BlockEnd::NotSet => return write!(f, "  Not set"),
             BlockEnd::Match { info } => {
                 return write!(f, "  Match({:?})", info.debug(ctx));
             }

crates/cairo-lang-lowering/src/fmt.rs

@@ -1,6 +1,8 @@
+use cairo_lang_debug::DebugWithDb;
 use cairo_lang_defs::ids::{MemberId, NamedLanguageElementId};
 use cairo_lang_diagnostics::Maybe;
 use cairo_lang_semantic as semantic;
+use cairo_lang_semantic::expr::fmt::ExprFormatter;
 use cairo_lang_semantic::types::{peel_snapshots, wrap_in_snapshots};
 use cairo_lang_semantic::usage::{MemberPath, Usage};
 use cairo_lang_syntax::node::TypedStablePtr;
@@ -406,6 +408,22 @@ impl BlockBuilder {
     }
 }
 
+impl<'a> DebugWithDb<ExprFormatter<'a>> for BlockBuilder {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>, db: &ExprFormatter<'a>) -> std::fmt::Result {
+        writeln!(f, "block_id: {:?}", self.block_id)?;
+        if !self.statements.statements.is_empty() {
+            writeln!(f, "statements:")?;
+            for statement in &self.statements.statements {
+                writeln!(f, "  {statement:?}")?;
+            }
+        }
+        writeln!(f, "semantics:")?;
+        write!(f, "{}", indent::indent_all_with("  ", format!("{:?}", self.semantics.debug(db))))?;
+
+        Ok(())
+    }
+}
+
 /// Gets the type of a semantic variable.
 fn get_ty(ctx: &LoweringContext<'_, '_>, member_path: &MemberPath) -> semantic::TypeId {
     match member_path {
@@ -537,3 +555,28 @@ impl StructRecomposer for BlockStructRecomposer<'_, '_, '_> {
         self.ctx.db
     }
 }
+
+/// Given a list of block builders, creates a new single block builder and finalizes all
+/// the block builders with a [BlockEnd::Goto] to the new block.
+///
+/// The mapping from semantic variables to lowered variables in the new block follows these rules:
+///
+/// * Variables mapped to the same lowered variable across all input blocks are kept as-is.
+/// * Local variables that appear in only a subset of the blocks are removed.
+/// * Variables with different mappings across blocks are remapped to a new lowered variable.
+///
+/// If only one parent builder is given, returns it without creating a new block.
+// TODO(lior): Remove `allow(dead_code)` once the function is used.
+#[allow(dead_code)]
+pub fn merge_block_builders(
+    _ctx: &mut LoweringContext<'_, '_>,
+    parent_builders: Vec<BlockBuilder>,
+    _location: LocationId,
+) -> BlockBuilder {
+    // If there is only one parent builder, return it.
+    if parent_builders.len() == 1 {
+        return parent_builders.into_iter().next().unwrap();
+    }
+
+    todo!("Merging multiple block builders is not supported yet.");
+}

crates/cairo-lang-lowering/src/lower/block_builder.rs

@@ -0,0 +1,202 @@
+use cairo_lang_debug::DebugWithDb;
+use cairo_lang_semantic::corelib::unit_ty;
+use cairo_lang_semantic::expr::fmt::ExprFormatter;
+use cairo_lang_semantic::test_utils::{TestFunction, setup_test_function};
+use cairo_lang_semantic::usage::MemberPath;
+use cairo_lang_semantic::{self as semantic, Expr, Statement, StatementId};
+use cairo_lang_syntax::node::TypedStablePtr;
+use cairo_lang_test_utils::parse_test_file::TestRunnerResult;
+use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
+use cairo_lang_utils::unordered_hash_set::UnorderedHashSet;
+use cairo_lang_utils::{Upcast, extract_matches};
+use itertools::Itertools;
+
+use super::block_builder::{BlockBuilder, merge_block_builders};
+use super::context::{LoweringContext, VarRequest};
+use super::test_utils::{create_encapsulating_ctx, create_lowering_context};
+use crate::VariableId;
+use crate::fmt::LoweredFormatter;
+use crate::test_utils::LoweringDatabaseForTesting;
+
+const N_LOWERING_VARS: usize = 100;
+
+cairo_lang_test_utils::test_file_test!(
+    test_merge_block_builders,
+    "src/lower/test_data",
+    {
+        merge_block_builders: "merge_block_builders",
+    },
+    test_merge_block_builders
+);
+
+/// Tests the [merge_block_builders] function.
+///
+/// Each test case has the following input sections:
+/// - `variables`: A comma-separated list of "var_name: type".
+/// - `block_definitions`: Defines the semantic to lowering map of each input block.
+///
+///   For example:
+///   ```ignore
+///      ((x, 0),);
+///      ((x.a, 1), (y, 2));
+///   ```
+///   represents two blocks, where:
+///   * The first maps `x` to lowered variable 0.
+///   * The second maps `x.a` to lowered variable 1, and `y` to lowered variable 2.
+///
+///   Note that `x` and `x.*` should not be specified together in one block.
+///
+/// - `module_code`: Additional code for defining structs and helper functions.
+fn test_merge_block_builders(
+    inputs: &OrderedHashMap<String, String>,
+    _args: &OrderedHashMap<String, String>,
+) -> TestRunnerResult {
+    let db = LoweringDatabaseForTesting::default();
+    // Create a function with the given variables as parameters and the given block definitions
+    // as a dummy function body.
+    // Note that the function body is not lowered at any point, and it is only used to parse the
+    // semantic to lowering map.
+    let test_function = setup_test_function(
+        &db,
+        &format!("fn foo ({}) {{ {} }}", &inputs["variables"], &inputs["block_definitions"]),
+        "foo",
+        inputs.get("module_code").unwrap_or(&"".into()),
+    )
+    .unwrap();
+
+    let mut encapsulating_ctx =
+        create_encapsulating_ctx(&db, test_function.function_id, &test_function.signature);
+
+    let mut ctx = create_lowering_context(
+        &db,
+        test_function.function_id,
+        &test_function.signature,
+        &mut encapsulating_ctx,
+    );
+
+    // Create dummy lowering variables.
+    let dummy_location = ctx.get_location(test_function.signature.stable_ptr.untyped());
+    let lowering_vars: Vec<VariableId> = (0..N_LOWERING_VARS)
+        .map(|_| ctx.new_var(VarRequest { ty: unit_ty(ctx.db), location: dummy_location }))
+        .collect();
+
+    let expr_formatter = ExprFormatter { db: db.upcast(), function_id: test_function.function_id };
+
+    let input_blocks = create_block_builders(&mut ctx, &test_function, &lowering_vars);
+    let input_blocks_str =
+        input_blocks.iter().map(|b| format!("{:?}", b.debug(&expr_formatter))).join("\n");
+
+    // Invoke [merge_block_builders] on the input blocks.
+    let merged_block = merge_block_builders(&mut ctx, input_blocks, dummy_location);
+
+    let lowered_formatter = LoweredFormatter::new(db.upcast(), &ctx.variables.variables);
+    let lowered_blocks = ctx.blocks.build().unwrap();
+    let lowered_str = lowered_blocks
+        .iter()
+        .map(|(block_id, block)| {
+            format!(
+                "{:?}:\n{:?}\n",
+                block_id.debug(&lowered_formatter),
+                block.debug(&lowered_formatter)
+            )
+        })
+        .join("");
+
+    TestRunnerResult {
+        outputs: OrderedHashMap::from([
+            ("input_blocks".into(), input_blocks_str),
+            ("merged_block_builder".into(), format!("{:?}", merged_block.debug(&expr_formatter))),
+            ("lowered".into(), lowered_str),
+        ]),
+        error: None,
+    }
+}
+
+/// Creates a block builder for each semantic "statement" in the function body.
+///
+/// See [create_block_builder] for more details.
+fn create_block_builders(
+    ctx: &mut LoweringContext<'_, '_>,
+    test_function: &TestFunction,
+    lowering_vars: &[VariableId],
+) -> Vec<BlockBuilder> {
+    let expr = ctx.function_body.arenas.exprs[test_function.body].clone();
+    let block_expr = extract_matches!(expr, Expr::Block);
+
+    block_expr
+        .statements
+        .iter()
+        .map(|statement_id| create_block_builder(ctx, *statement_id, lowering_vars))
+        .collect()
+}
+
+/// Given a semantic "statement" of the form:
+///    `((member_path, lower_var_idx), ...)`
+/// creates a block builder with a semantic mapping that maps each member path to the corresponding
+/// given lowered variable.
+///
+/// Assumption: if a certain semantic variable is mapped, all its children should not be mapped.
+///
+/// Note that the statement is not a real statement - it is not lowered, and it is only used to
+/// define the semantic mapping.
+fn create_block_builder(
+    ctx: &mut LoweringContext<'_, '_>,
+    statement_id: StatementId,
+    lowering_vars: &[VariableId],
+) -> BlockBuilder {
+    let block_id = ctx.blocks.alloc_empty();
+    let mut block_builder = BlockBuilder::root(block_id);
+    let mut visited_vars: UnorderedHashSet<semantic::VarId> = Default::default();
+
+    let statement_expr =
+        extract_matches!(&ctx.function_body.arenas.statements[statement_id], Statement::Expr);
+    let external_tuple =
+        extract_matches!(&ctx.function_body.arenas.exprs[statement_expr.expr], Expr::Tuple);
+
+    let expr_ids = external_tuple.items.clone();
+    for expr_id in expr_ids {
+        let inner_tuple = extract_matches!(&ctx.function_body.arenas.exprs[expr_id], Expr::Tuple);
+        let lower_var_idx: usize = (&extract_matches!(
+            &ctx.function_body.arenas.exprs[inner_tuple.items[1]],
+            Expr::Literal
+        )
+        .value)
+            .try_into()
+            .unwrap();
+
+        match &ctx.function_body.arenas.exprs[inner_tuple.items[0]] {
+            Expr::MemberAccess(member_access) => {
+                let member_path: MemberPath = (member_access.member_path.as_ref().unwrap()).into();
+                let mut var = &member_path;
+                while let MemberPath::Member { parent: v, .. } = var {
+                    var = v;
+                }
+                let var_id = extract_matches!(var, MemberPath::Var);
+
+                if visited_vars.insert(*var_id) {
+                    block_builder.put_semantic(*var_id, lowering_vars[lower_var_idx]);
+                }
+
+                let location = ctx.get_location(member_access.stable_ptr.untyped());
+                block_builder.update_ref_raw(
+                    ctx,
+                    member_path,
+                    lowering_vars[lower_var_idx],
+                    location,
+                );
+                // Remove the statements that were created as part of the `update_ref_raw` call.
+                block_builder.statements.statements.clear();
+            }
+            Expr::Var(var) => {
+                if visited_vars.insert(var.var) {
+                    block_builder.put_semantic(var.var, lowering_vars[lower_var_idx]);
+                }
+            }
+            expr => {
+                panic!("Unexpected expression: {expr:?}");
+            }
+        }
+    }
+
+    block_builder
+}

crates/cairo-lang-lowering/src/lower/block_builder_test.rs

@@ -70,6 +70,12 @@ mod lower_let_else;
 mod lower_match;
 pub mod refs;
 
+#[cfg(test)]
+mod test_utils;
+
+#[cfg(test)]
+mod block_builder_test;
+
 #[cfg(test)]
 mod generated_test;
 

crates/cairo-lang-lowering/src/lower/mod.rs

@@ -5,7 +5,7 @@ use cairo_lang_semantic::usage::MemberPath;
 use cairo_lang_semantic::{self as semantic};
 use cairo_lang_utils::extract_matches;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
-use itertools::chain;
+use itertools::{Itertools, chain};
 
 use crate::VariableId;
 use crate::db::LoweringGroup;
@@ -154,6 +154,15 @@ impl SemanticLoweringMapping {
     }
 }
 
+impl<'a> cairo_lang_debug::debug::DebugWithDb<ExprFormatter<'a>> for SemanticLoweringMapping {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>, db: &ExprFormatter<'a>) -> std::fmt::Result {
+        for (member_path, value) in self.scattered.iter() {
+            writeln!(f, "{:?}: {value}", member_path.debug(db))?;
+        }
+        Ok(())
+    }
+}
+
 /// A trait for deconstructing and constructing structs.
 pub trait StructRecomposer {
     fn deconstruct(
@@ -183,11 +192,26 @@ pub trait StructRecomposer {
 enum Value {
     /// The value of member path is stored in a lowered variable.
     Var(VariableId),
-    /// The value of the member path is not stored. It should be reconstructed from the member
-    /// values.
+    /// The value of the member path is not stored. If needed, it should be reconstructed from the
+    /// member values.
     Scattered(Box<Scattered>),
 }
 
+impl std::fmt::Display for Value {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Value::Var(var) => write!(f, "v{}", var.index()),
+            Value::Scattered(scattered) => {
+                write!(
+                    f,
+                    "Scattered({})",
+                    scattered.members.values().map(|value| value.to_string()).join(", ")
+                )
+            }
+        }
+    }
+}
+
 /// A value for a non-stored member path. Recursively holds the [Value] for the members.
 #[derive(Clone, Debug, DebugWithDb)]
 #[debug_db(ExprFormatter<'a>)]

crates/cairo-lang-lowering/src/lower/refs.rs

@@ -0,0 +1,29 @@
+//! > Single block
+
+//! > test_runner_name
+test_merge_block_builders
+
+//! > variables
+x: felt252
+
+//! > block_definitions
+// A single block with a single variable `x`, mapped to lowered variable `v0`.
+(
+    (x, 0),
+);
+
+//! > input_blocks
+block_id: BlockId(0)
+semantics:
+  Var(ParamId(test::x)): v0
+
+//! > merged_block_builder
+block_id: BlockId(0)
+semantics:
+  Var(ParamId(test::x)): v0
+
+//! > lowered
+blk0:
+Statements:
+End:
+  Not set