Skip to content

Commit

Permalink
Added multi-pop for span.
Browse files Browse the repository at this point in the history
commit-id:813e40ab
  • Loading branch information
orizi committed May 19, 2024
1 parent 2361876 commit 3960647
Show file tree
Hide file tree
Showing 9 changed files with 503 additions and 16 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,10 @@ pub fn core_libfunc_ap_change<InfoProvider: InvocationApChangeInfoProvider>(
| ArrayConcreteLibfunc::SnapshotPopBack(_) => {
vec![ApChange::Known(1), ApChange::Known(1)]
}
ArrayConcreteLibfunc::SnapshotMultiPopFront(_)
| ArrayConcreteLibfunc::SnapshotMultiPopBack(_) => {
vec![ApChange::Known(3), ApChange::Known(3)]
}
ArrayConcreteLibfunc::Get(libfunc) => {
if info_provider.type_size(&libfunc.ty) == 1 { [4, 3] } else { [5, 4] }
.map(ApChange::Known)
Expand Down
7 changes: 7 additions & 0 deletions crates/cairo-lang-sierra-gas/src/core_libfunc_cost_base.rs
Original file line number Diff line number Diff line change
Expand Up @@ -258,6 +258,13 @@ pub fn core_libfunc_cost(
| ArrayConcreteLibfunc::SnapshotPopBack(_) => {
vec![ConstCost::steps(2).into(), ConstCost::steps(3).into()]
}
ArrayConcreteLibfunc::SnapshotMultiPopFront(_)
| ArrayConcreteLibfunc::SnapshotMultiPopBack(_) => {
vec![
(ConstCost::steps(4) + ConstCost::range_checks(1)).into(),
(ConstCost::steps(5) + ConstCost::range_checks(1)).into(),
]
}
ArrayConcreteLibfunc::Get(libfunc) => {
if info_provider.type_size(&libfunc.ty) == 1 {
vec![
Expand Down
127 changes: 125 additions & 2 deletions crates/cairo-lang-sierra-to-casm/src/invocations/array.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use cairo_lang_casm::builder::CasmBuilder;
use cairo_lang_casm::builder::{CasmBuilder, Var};
use cairo_lang_casm::casm_build_extend;
use cairo_lang_sierra::extensions::array::ArrayConcreteLibfunc;
use cairo_lang_sierra::extensions::array::{ArrayConcreteLibfunc, ConcreteMultiPopLibfunc};
use cairo_lang_sierra::ids::ConcreteTypeId;

use super::{CompiledInvocation, CompiledInvocationBuilder, InvocationError};
Expand Down Expand Up @@ -28,6 +28,12 @@ pub fn build(
ArrayConcreteLibfunc::SnapshotPopBack(libfunc) => {
build_pop_back(&libfunc.ty, builder, false)
}
ArrayConcreteLibfunc::SnapshotMultiPopFront(libfunc) => {
build_multi_pop_front(libfunc, builder)
}
ArrayConcreteLibfunc::SnapshotMultiPopBack(libfunc) => {
build_multi_pop_back(libfunc, builder)
}
ArrayConcreteLibfunc::Get(libfunc) => build_array_get(&libfunc.ty, builder),
ArrayConcreteLibfunc::Slice(libfunc) => build_array_slice(&libfunc.ty, builder),
ArrayConcreteLibfunc::Len(libfunc) => build_array_len(&libfunc.ty, builder),
Expand Down Expand Up @@ -365,3 +371,120 @@ fn build_array_len(
Default::default(),
))
}

/// Handles a Sierra statement for popping elements from the beginning of an array.
fn build_multi_pop_front(
libfunc: &ConcreteMultiPopLibfunc,
builder: CompiledInvocationBuilder<'_>,
) -> Result<CompiledInvocation, InvocationError> {
let [expr_range_check, expr_arr] = builder.try_get_refs()?;
let range_check = expr_range_check.try_unpack_single()?;
let [arr_start, arr_end] = expr_arr.try_unpack()?;
let element_size = builder.program_info.type_sizes[&libfunc.ty];

let mut casm_builder = CasmBuilder::default();
add_input_variables! {casm_builder,
buffer(1) range_check;
deref arr_start;
deref arr_end;
};
let popped_size: i16 = libfunc.popped_values * element_size;
casm_build_extend!(casm_builder, let orig_range_check = range_check;);
extend_multi_pop_failure_checks(
&mut casm_builder,
range_check,
arr_start,
arr_end,
popped_size,
);
casm_build_extend! {casm_builder,
const popped_size = popped_size;
tempvar rc;
// Calculating the new start, as it is required for calculating the range checked value.
tempvar new_start = arr_start + popped_size;
assert rc = arr_end - new_start;
assert rc = *(range_check++);
};
let failure_handle = get_non_fallthrough_statement_id(&builder);
Ok(builder.build_from_casm_builder(
casm_builder,
[
("Fallthrough", &[&[range_check], &[new_start, arr_end], &[arr_start]], None),
("Failure", &[&[range_check], &[arr_start, arr_end]], Some(failure_handle)),
],
CostValidationInfo {
range_check_info: Some((orig_range_check, range_check)),
extra_costs: None,
},
))
}

/// Handles a Sierra statement for popping elements from the end of an array.
fn build_multi_pop_back(
libfunc: &ConcreteMultiPopLibfunc,
builder: CompiledInvocationBuilder<'_>,
) -> Result<CompiledInvocation, InvocationError> {
let [expr_range_check, expr_arr] = builder.try_get_refs()?;
let range_check = expr_range_check.try_unpack_single()?;
let [arr_start, arr_end] = expr_arr.try_unpack()?;
let element_size = builder.program_info.type_sizes[&libfunc.ty];

let mut casm_builder = CasmBuilder::default();
add_input_variables! {casm_builder,
buffer(1) range_check;
deref arr_start;
deref arr_end;
};
let popped_size: i16 = libfunc.popped_values * element_size;
casm_build_extend!(casm_builder, let orig_range_check = range_check;);
extend_multi_pop_failure_checks(
&mut casm_builder,
range_check,
arr_start,
arr_end,
popped_size,
);
casm_build_extend! {casm_builder,
const popped_size = popped_size;
tempvar rc;
// Calculating the new end, as it is required for calculating the range checked value.
tempvar new_end = arr_end - popped_size;
assert rc = new_end - arr_start;
assert rc = *(range_check++);
};
let failure_handle = get_non_fallthrough_statement_id(&builder);
Ok(builder.build_from_casm_builder(
casm_builder,
[
("Fallthrough", &[&[range_check], &[arr_start, new_end], &[new_end]], None),
("Failure", &[&[range_check], &[arr_start, arr_end]], Some(failure_handle)),
],
CostValidationInfo {
range_check_info: Some((orig_range_check, range_check)),
extra_costs: None,
},
))
}

/// Extends the casm builder with the common part of the multi-pop front and multi-pop back.
fn extend_multi_pop_failure_checks(
casm_builder: &mut CasmBuilder,
range_check: Var,
arr_start: Var,
arr_end: Var,
popped_size: i16,
) {
casm_build_extend! {casm_builder,
const popped_size_plus_1 = popped_size + 1;
const popped_size = popped_size;
let arr_start_popped = arr_start + popped_size;
tempvar has_enough_elements;
hint TestLessThanOrEqual {lhs: arr_start_popped, rhs: arr_end} into {dst: has_enough_elements};
jump HasEnoughElements if has_enough_elements != 0;
tempvar minus_length = arr_start - arr_end;
tempvar rc = minus_length + popped_size_plus_1;
assert rc = *(range_check++);
jump Failure;
HasEnoughElements:
};
}
11 changes: 11 additions & 0 deletions crates/cairo-lang-sierra/src/extensions/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,17 @@ fn args_as_two_types(
}
}

/// Helper for extracting a type and a value from the template arguments.
fn args_as_type_and_value(
args: &[GenericArg],
) -> Result<(ConcreteTypeId, BigInt), SpecializationError> {
match args {
[GenericArg::Type(ty), GenericArg::Value(value)] => Ok((ty.clone(), value.clone())),
[_, _] => Err(SpecializationError::UnsupportedGenericArg),
_ => Err(SpecializationError::WrongNumberOfGenericArgs),
}
}

/// Helper for extracting the type from the template arguments.
fn args_as_single_user_func(args: &[GenericArg]) -> Result<FunctionId, SpecializationError> {
match args {
Expand Down
184 changes: 182 additions & 2 deletions crates/cairo-lang-sierra/src/extensions/modules/array.rs
Original file line number Diff line number Diff line change
@@ -1,19 +1,23 @@
use num_traits::ToPrimitive;

use super::boxing::box_ty;
use super::range_check::RangeCheckType;
use super::snapshot::snapshot_ty;
use super::structure::StructConcreteType;
use super::utils::fixed_array_ty;
use crate::define_libfunc_hierarchy;
use crate::extensions::lib_func::{
BranchSignature, DeferredOutputKind, LibfuncSignature, OutputVarInfo, ParamSignature,
SierraApChange, SignatureAndTypeGenericLibfunc, SignatureOnlyGenericLibfunc,
SignatureSpecializationContext, WrapSignatureAndTypeGenericLibfunc,
SignatureSpecializationContext, SpecializationContext, WrapSignatureAndTypeGenericLibfunc,
};
use crate::extensions::type_specialization_context::TypeSpecializationContext;
use crate::extensions::types::{
GenericTypeArgGenericType, GenericTypeArgGenericTypeWrapper, TypeInfo,
};
use crate::extensions::{
args_as_single_type, NamedType, OutputVarReferenceInfo, SpecializationError,
args_as_single_type, args_as_type_and_value, NamedLibfunc, NamedType, OutputVarReferenceInfo,
SignatureBasedConcreteLibfunc, SpecializationError,
};
use crate::ids::{ConcreteTypeId, GenericTypeId};
use crate::program::GenericArg;
Expand Down Expand Up @@ -60,6 +64,8 @@ define_libfunc_hierarchy! {
Len(ArrayLenLibfunc),
SnapshotPopFront(ArraySnapshotPopFrontLibfunc),
SnapshotPopBack(ArraySnapshotPopBackLibfunc),
SnapshotMultiPopFront(ArraySnapshotMultiPopFrontLibfunc),
SnapshotMultiPopBack(ArraySnapshotMultiPopBackLibfunc),
}, ArrayConcreteLibfunc
}

Expand Down Expand Up @@ -501,3 +507,177 @@ impl SignatureAndTypeGenericLibfunc for ArraySnapshotPopBackLibfuncWrapped {
}
pub type ArraySnapshotPopBackLibfunc =
WrapSignatureAndTypeGenericLibfunc<ArraySnapshotPopBackLibfuncWrapped>;

/// Libfunc for popping multiple first values from the beginning of an array snapshot.
#[derive(Default)]
pub struct ArraySnapshotMultiPopFrontLibfunc {}
impl NamedLibfunc for ArraySnapshotMultiPopFrontLibfunc {
const STR_ID: &'static str = "array_snapshot_multi_pop_front";

type Concrete = ConcreteMultiPopLibfunc;

fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
args: &[GenericArg],
) -> Result<LibfuncSignature, SpecializationError> {
let (ty, count) = args_as_type_and_value(args)?;
let arr_ty = context.get_wrapped_concrete_type(ArrayType::id(), ty.clone())?;
let arr_snapshot_ty = snapshot_ty(context, arr_ty)?;
let range_check_ty = context.get_concrete_type(RangeCheckType::id(), &[])?;
Ok(LibfuncSignature {
param_signatures: vec![
ParamSignature::new(range_check_ty.clone()).with_allow_add_const(),
ParamSignature::new(arr_snapshot_ty.clone()),
],
branch_signatures: vec![
BranchSignature {
vars: vec![
OutputVarInfo::new_builtin(range_check_ty.clone(), 0),
OutputVarInfo {
ty: arr_snapshot_ty.clone(),
ref_info: OutputVarReferenceInfo::Deferred(
DeferredOutputKind::AddConst { param_idx: 1 },
),
},
OutputVarInfo {
ty: snapshot_ty(
context,
box_ty(
context,
fixed_array_ty(
context,
ty,
count
.to_usize()
.ok_or(SpecializationError::UnsupportedGenericArg)?,
)?,
)?,
)?,
ref_info: OutputVarReferenceInfo::PartialParam { param_idx: 1 },
},
],
ap_change: SierraApChange::Known { new_vars_only: false },
},
BranchSignature {
vars: vec![
OutputVarInfo::new_builtin(range_check_ty, 0),
OutputVarInfo {
ty: arr_snapshot_ty,
ref_info: OutputVarReferenceInfo::SameAsParam { param_idx: 1 },
},
],
ap_change: SierraApChange::Known { new_vars_only: false },
},
],
fallthrough: Some(0),
})
}

fn specialize(
&self,
context: &dyn SpecializationContext,
args: &[GenericArg],
) -> Result<Self::Concrete, SpecializationError> {
let (ty, count) = args_as_type_and_value(args)?;
Ok(ConcreteMultiPopLibfunc {
ty,
signature: self.specialize_signature(context.upcast(), args)?,
// Early failing on size, although the type itself would fail on definition.
popped_values: count.to_i16().ok_or(SpecializationError::UnsupportedGenericArg)?,
})
}
}

/// Libfunc for popping the last value from the end of an array snapshot.
#[derive(Default)]
pub struct ArraySnapshotMultiPopBackLibfunc {}
impl NamedLibfunc for ArraySnapshotMultiPopBackLibfunc {
const STR_ID: &'static str = "array_snapshot_multi_pop_back";

type Concrete = ConcreteMultiPopLibfunc;

fn specialize_signature(
&self,
context: &dyn SignatureSpecializationContext,
args: &[GenericArg],
) -> Result<LibfuncSignature, SpecializationError> {
let (ty, count) = args_as_type_and_value(args)?;
let arr_ty = context.get_wrapped_concrete_type(ArrayType::id(), ty.clone())?;
let arr_snapshot_ty = snapshot_ty(context, arr_ty)?;
let range_check_ty = context.get_concrete_type(RangeCheckType::id(), &[])?;
Ok(LibfuncSignature {
param_signatures: vec![
ParamSignature::new(range_check_ty.clone()).with_allow_add_const(),
ParamSignature::new(arr_snapshot_ty.clone()),
],
branch_signatures: vec![
BranchSignature {
vars: vec![
OutputVarInfo::new_builtin(range_check_ty.clone(), 0),
OutputVarInfo {
ty: arr_snapshot_ty.clone(),
ref_info: OutputVarReferenceInfo::Deferred(
DeferredOutputKind::AddConst { param_idx: 1 },
),
},
OutputVarInfo {
ty: snapshot_ty(
context,
box_ty(
context,
fixed_array_ty(
context,
ty,
count
.to_usize()
.ok_or(SpecializationError::UnsupportedGenericArg)?,
)?,
)?,
)?,
ref_info: OutputVarReferenceInfo::Deferred(DeferredOutputKind::Generic),
},
],
ap_change: SierraApChange::Known { new_vars_only: false },
},
BranchSignature {
vars: vec![
OutputVarInfo::new_builtin(range_check_ty, 0),
OutputVarInfo {
ty: arr_snapshot_ty,
ref_info: OutputVarReferenceInfo::SameAsParam { param_idx: 1 },
},
],
ap_change: SierraApChange::Known { new_vars_only: false },
},
],
fallthrough: Some(0),
})
}

fn specialize(
&self,
context: &dyn SpecializationContext,
args: &[GenericArg],
) -> Result<Self::Concrete, SpecializationError> {
let (ty, count) = args_as_type_and_value(args)?;
Ok(ConcreteMultiPopLibfunc {
ty,
signature: self.specialize_signature(context.upcast(), args)?,
// Early failing on size, although the type itself would fail on definition.
popped_values: count.to_i16().ok_or(SpecializationError::UnsupportedGenericArg)?,
})
}
}

/// Struct the data for a multi pop action.
pub struct ConcreteMultiPopLibfunc {
pub ty: ConcreteTypeId,
pub popped_values: i16,
pub signature: LibfuncSignature,
}
impl SignatureBasedConcreteLibfunc for ConcreteMultiPopLibfunc {
fn signature(&self) -> &LibfuncSignature {
&self.signature
}
}
Loading

0 comments on commit 3960647

Please sign in to comment.