diff --git a/crates/cairo-lang-lowering/src/optimizations/const_folding.rs b/crates/cairo-lang-lowering/src/optimizations/const_folding.rs index 46980a110e2..e7e8025caaf 100644 --- a/crates/cairo-lang-lowering/src/optimizations/const_folding.rs +++ b/crates/cairo-lang-lowering/src/optimizations/const_folding.rs @@ -1325,18 +1325,35 @@ impl<'db, 'mt> ConstFoldingContext<'db, 'mt> { let var_info = self.var_info.get(&info.inputs[0].var_id)?; let desnapped = try_extract_matches!(var_info.as_ref(), VarInfo::Snapshot)?; let element_var_infos = try_extract_matches!(desnapped.as_ref(), VarInfo::Array)?; - // TODO(orizi): Propagate success values as well. - if element_var_infos.is_empty() { + let wrap_snapshot_array = |var_infos: Vec>>>| { + VarInfo::Snapshot(VarInfo::Array(var_infos).into()) + }; + if let Some((element_var_info, remaining_var_infos)) = if id + == self.array_snapshot_pop_front + { + element_var_infos.split_first().map(|(first, rest)| (first.clone(), rest.to_vec())) + } else { + element_var_infos.split_last().map(|(last, rest)| (last.clone(), rest.to_vec())) + } { + let arm = &info.arms[0]; + self.var_info + .insert(arm.var_ids[0], wrap_snapshot_array(remaining_var_infos).into()); + if let Some(element_var_info) = element_var_info { + self.var_info.insert( + arm.var_ids[1], + VarInfo::Box(VarInfo::Snapshot(element_var_info).into()).into(), + ); + } + None + } else { let arm = &info.arms[1]; - self.var_info.insert(arm.var_ids[0], VarInfo::Array(vec![]).into()); + self.var_info.insert(arm.var_ids[0], wrap_snapshot_array(vec![]).into()); Some(BlockEnd::Goto( arm.block_id, VarRemapping { remapping: FromIterator::from_iter([(arm.var_ids[0], info.inputs[0])]), }, )) - } else { - None } } else { None diff --git a/crates/cairo-lang-lowering/src/optimizations/test_data/const_folding b/crates/cairo-lang-lowering/src/optimizations/test_data/const_folding index 40cf443fb0d..fdd12a6f6c9 100644 --- a/crates/cairo-lang-lowering/src/optimizations/test_data/const_folding +++ b/crates/cairo-lang-lowering/src/optimizations/test_data/const_folding @@ -6075,6 +6075,248 @@ End: //! > ========================================================================== +//! > Snapshot pop_front propagates success values. + +//! > test_runner_name +test_match_optimizer + +//! > function_code +fn foo() -> felt252 { + let mut snapshot = array![1, 2].span().snapshot; + let Some(v) = array_snapshot_pop_front(ref snapshot) else { + return 0; + }; + let Some(w) = array_snapshot_pop_front(ref snapshot) else { + return 100; + }; + *v.unbox() + *w.unbox() +} + +//! > function_name +foo + +//! > module_code +use core::array::array_snapshot_pop_front; + +//! > semantic_diagnostics + +//! > before +Parameters: +blk0 (root): +Statements: + (v0: core::array::Array::) <- core::array::array_new::() + (v1: core::felt252) <- 1 + (v2: core::array::Array::) <- core::array::array_append::(v0, v1) + (v3: core::felt252) <- 2 + (v4: core::array::Array::) <- core::array::array_append::(v2, v3) + (v5: core::array::Array::, v6: @core::array::Array::) <- snapshot(v4) +End: + Match(match core::array::array_snapshot_pop_front::(v6) { + Option::Some(v7, v8) => blk1, + Option::None(v9) => blk4, + }) + +blk1: +Statements: +End: + Match(match core::array::array_snapshot_pop_front::(v7) { + Option::Some(v10, v11) => blk2, + Option::None(v12) => blk3, + }) + +blk2: +Statements: + (v13: @core::felt252) <- unbox(v8) + (v14: core::felt252) <- desnap(v13) + (v15: @core::felt252) <- unbox(v11) + (v16: core::felt252) <- desnap(v15) + (v17: core::felt252) <- core::felt252_add(v14, v16) +End: + Return(v17) + +blk3: +Statements: + (v18: core::felt252) <- 100 +End: + Return(v18) + +blk4: +Statements: + (v19: core::felt252) <- 0 +End: + Return(v19) + +//! > after +Parameters: +blk0 (root): +Statements: + (v0: core::array::Array::) <- core::array::array_new::() + (v1: core::felt252) <- 1 + (v2: core::array::Array::) <- core::array::array_append::(v0, v1) + (v3: core::felt252) <- 2 + (v4: core::array::Array::) <- core::array::array_append::(v2, v3) + (v5: core::array::Array::, v6: @core::array::Array::) <- snapshot(v4) +End: + Match(match core::array::array_snapshot_pop_front::(v6) { + Option::Some(v7, v8) => blk1, + Option::None(v9) => blk4, + }) + +blk1: +Statements: +End: + Match(match core::array::array_snapshot_pop_front::(v7) { + Option::Some(v10, v11) => blk2, + Option::None(v12) => blk3, + }) + +blk2: +Statements: + (v13: @core::felt252) <- unbox(v8) + (v14: core::felt252) <- desnap(v13) + (v15: @core::felt252) <- unbox(v11) + (v16: core::felt252) <- desnap(v15) + (v17: core::felt252) <- 3 +End: + Return(v17) + +blk3: +Statements: + (v18: core::felt252) <- 100 +End: + Return(v18) + +blk4: +Statements: + (v19: core::felt252) <- 0 +End: + Return(v19) + +//! > lowering_diagnostics + +//! > ========================================================================== + +//! > Snapshot pop_back propagates success values. + +//! > test_runner_name +test_match_optimizer + +//! > function_code +fn foo() -> felt252 { + let mut snapshot = array![1, 2].span().snapshot; + let Some(v) = array_snapshot_pop_back(ref snapshot) else { + return 0; + }; + let Some(w) = array_snapshot_pop_back(ref snapshot) else { + return 100; + }; + *v.unbox() + *w.unbox() +} + +//! > function_name +foo + +//! > module_code +use core::array::array_snapshot_pop_back; + +//! > semantic_diagnostics + +//! > before +Parameters: +blk0 (root): +Statements: + (v0: core::array::Array::) <- core::array::array_new::() + (v1: core::felt252) <- 1 + (v2: core::array::Array::) <- core::array::array_append::(v0, v1) + (v3: core::felt252) <- 2 + (v4: core::array::Array::) <- core::array::array_append::(v2, v3) + (v5: core::array::Array::, v6: @core::array::Array::) <- snapshot(v4) +End: + Match(match core::array::array_snapshot_pop_back::(v6) { + Option::Some(v7, v8) => blk1, + Option::None(v9) => blk4, + }) + +blk1: +Statements: +End: + Match(match core::array::array_snapshot_pop_back::(v7) { + Option::Some(v10, v11) => blk2, + Option::None(v12) => blk3, + }) + +blk2: +Statements: + (v13: @core::felt252) <- unbox(v8) + (v14: core::felt252) <- desnap(v13) + (v15: @core::felt252) <- unbox(v11) + (v16: core::felt252) <- desnap(v15) + (v17: core::felt252) <- core::felt252_add(v14, v16) +End: + Return(v17) + +blk3: +Statements: + (v18: core::felt252) <- 100 +End: + Return(v18) + +blk4: +Statements: + (v19: core::felt252) <- 0 +End: + Return(v19) + +//! > after +Parameters: +blk0 (root): +Statements: + (v0: core::array::Array::) <- core::array::array_new::() + (v1: core::felt252) <- 1 + (v2: core::array::Array::) <- core::array::array_append::(v0, v1) + (v3: core::felt252) <- 2 + (v4: core::array::Array::) <- core::array::array_append::(v2, v3) + (v5: core::array::Array::, v6: @core::array::Array::) <- snapshot(v4) +End: + Match(match core::array::array_snapshot_pop_back::(v6) { + Option::Some(v7, v8) => blk1, + Option::None(v9) => blk4, + }) + +blk1: +Statements: +End: + Match(match core::array::array_snapshot_pop_back::(v7) { + Option::Some(v10, v11) => blk2, + Option::None(v12) => blk3, + }) + +blk2: +Statements: + (v13: @core::felt252) <- unbox(v8) + (v14: core::felt252) <- desnap(v13) + (v15: @core::felt252) <- unbox(v11) + (v16: core::felt252) <- desnap(v15) + (v17: core::felt252) <- 3 +End: + Return(v17) + +blk3: +Statements: + (v18: core::felt252) <- 100 +End: + Return(v18) + +blk4: +Statements: + (v19: core::felt252) <- 0 +End: + Return(v19) + +//! > lowering_diagnostics + +//! > ========================================================================== + //! > Panic with byte array value. //! > test_runner_name