test(starknet_os): add small fft regression test

crates/starknet_os/src/hint_processor/snos_hint_processor.rs

@@ -51,7 +51,6 @@ impl<S: StateReader> SnosHintProcessor<S> {
 
     /// Stores the data-availabilty segment, to be used for computing the KZG commitment in blob
     /// mode.
-    #[allow(dead_code)]
     pub(crate) fn set_da_segment(&mut self, da_segment: Vec<Felt>) -> Result<(), OsHintError> {
         if self.da_segment.is_some() {
             return Err(OsHintError::AssertionFailed {

crates/starknet_os/src/hints/error.rs

@@ -15,6 +15,7 @@ use starknet_api::StarknetApiError;
 use starknet_types_core::felt::Felt;
 
 use crate::hints::enum_definition::AllHints;
+use crate::hints::hint_implementation::kzg::utils::FftError;
 use crate::hints::vars::{Const, Ids};
 
 #[derive(Debug, thiserror::Error)]
@@ -34,11 +35,15 @@ pub enum OsHintError {
     #[error("{id:?} value {felt} is not a bit.")]
     ExpectedBit { id: Ids, felt: Felt },
     #[error(transparent)]
+    Fft(#[from] FftError),
+    #[error(transparent)]
     FromUtf8(#[from] FromUtf8Error),
     #[error("The identifier {0:?} has no full name.")]
     IdentifierHasNoFullName(Box<Identifier>),
     #[error("The identifier {0:?} has no members.")]
     IdentifierHasNoMembers(Box<Identifier>),
+    #[error("Failed to convert {variant:?} felt value {felt:?} to type {ty}: {reason:?}.")]
+    IdsConversion { variant: Ids, felt: Felt, ty: String, reason: String },
     #[error(
         "Inconsistent block numbers: {actual}, {expected}. The constant STORED_BLOCK_HASH_BUFFER \
          is probably out of sync."

crates/starknet_os/src/hints/hint_implementation/kzg/implementation.rs

@@ -1,8 +1,76 @@
 use blockifier::state::state_api::StateReader;
+use cairo_vm::hint_processor::builtin_hint_processor::hint_utils::{
+    get_integer_from_var_name,
+    get_ptr_from_var_name,
+    insert_value_from_var_name,
+};
+use cairo_vm::types::relocatable::MaybeRelocatable;
+use num_bigint::BigInt;
+use starknet_types_core::felt::Felt;
 
-use crate::hints::error::OsHintResult;
+use crate::hints::error::{OsHintError, OsHintResult};
+use crate::hints::hint_implementation::kzg::utils::polynomial_coefficients_to_kzg_commitment;
 use crate::hints::types::HintArgs;
+use crate::hints::vars::{Const, Ids};
 
-pub(crate) fn store_da_segment<S: StateReader>(HintArgs { .. }: HintArgs<'_, S>) -> OsHintResult {
-    todo!()
+pub(crate) fn store_da_segment<S: StateReader>(
+    HintArgs { hint_processor, vm, ids_data, ap_tracking, constants, .. }: HintArgs<'_, S>,
+) -> OsHintResult {
+    log::debug!("Executing store_da_segment hint.");
+    let state_updates_start =
+        get_ptr_from_var_name(Ids::StateUpdatesStart.into(), vm, ids_data, ap_tracking)?;
+    let da_size_felt = get_integer_from_var_name(Ids::DaSize.into(), vm, ids_data, ap_tracking)?;
+    let da_size =
+        usize::try_from(da_size_felt.to_biguint()).map_err(|error| OsHintError::IdsConversion {
+            variant: Ids::DaSize,
+            felt: da_size_felt,
+            ty: "usize".to_string(),
+            reason: error.to_string(),
+        })?;
+
+    let da_segment: Vec<Felt> =
+        vm.get_integer_range(state_updates_start, da_size)?.into_iter().map(|s| *s).collect();
+
+    let blob_length_felt = Const::BlobLength.fetch(constants)?;
+    let blob_length = usize::try_from(blob_length_felt.to_biguint()).map_err(|error| {
+        OsHintError::ConstConversion {
+            variant: Const::BlobLength,
+            felt: *blob_length_felt,
+            ty: "usize".to_string(),
+            reason: error.to_string(),
+        }
+    })?;
+
+    let kzg_commitments: Vec<(Felt, Felt)> = da_segment
+        .chunks(blob_length)
+        .enumerate()
+        .map(|(chunk_id, chunk)| {
+            let coefficients: Vec<BigInt> = chunk.iter().map(|f| f.to_bigint()).collect();
+            log::debug!("Computing KZG commitment on chunk {chunk_id}...");
+            polynomial_coefficients_to_kzg_commitment(coefficients)
+        })
+        .collect::<Result<_, _>>()?;
+    log::debug!("Done computing KZG commitments.");
+
+    hint_processor.set_da_segment(da_segment)?;
+
+    let n_blobs = kzg_commitments.len();
+    let kzg_commitments_segment = vm.add_temporary_segment();
+    let evals_segment = vm.add_temporary_segment();
+
+    insert_value_from_var_name(Ids::NBlobs.into(), n_blobs, vm, ids_data, ap_tracking)?;
+    insert_value_from_var_name(
+        Ids::KzgCommitments.into(),
+        kzg_commitments_segment,
+        vm,
+        ids_data,
+        ap_tracking,
+    )?;
+    insert_value_from_var_name(Ids::Evals.into(), evals_segment, vm, ids_data, ap_tracking)?;
+
+    let kzg_commitments_flattened: Vec<MaybeRelocatable> =
+        kzg_commitments.into_iter().flat_map(|c| [c.0.into(), c.1.into()]).collect();
+    vm.write_arg(kzg_commitments_segment, &kzg_commitments_flattened)?;
+
+    Ok(())
 }

crates/starknet_os/src/hints/hint_implementation/kzg/test.rs

@@ -4,6 +4,7 @@ use c_kzg::KzgCommitment;
 use num_bigint::BigInt;
 use num_traits::{Num, One, Zero};
 use rstest::rstest;
+use starknet_types_core::felt::Felt;
 
 use crate::hints::hint_implementation::kzg::utils::{fft, split_commitment, BLS_PRIME};
 
@@ -22,6 +23,19 @@ fn generate(generator: &BigInt) -> Vec<BigInt> {
     array
 }
 
+#[rstest]
+fn test_small_fft_regression(#[values(true, false)] bit_reversed: bool) {
+    let prime = BigInt::from(17);
+    let generator = BigInt::from(3);
+    let coeffs: Vec<BigInt> = [0, 1, 2, 3].into_iter().map(BigInt::from).collect();
+    let expected_eval: Vec<BigInt> = (if bit_reversed { [6, 15, 9, 4] } else { [6, 9, 15, 4] })
+        .into_iter()
+        .map(BigInt::from)
+        .collect();
+    let actual_eval = fft(&coeffs, &generator, &prime, bit_reversed).unwrap();
+    assert_eq!(actual_eval, expected_eval);
+}
+
 #[rstest]
 fn test_fft(#[values(true, false)] bit_reversed: bool) {
     let prime = BigInt::from_str_radix(BLS_PRIME, 10).unwrap();
@@ -80,8 +94,8 @@ fn test_fft(#[values(true, false)] bit_reversed: bool) {
         16,
     ).unwrap(),
     (
-        BigInt::from_str_radix("590b91a225b1bf7fa5877ec546d090e0059f019c74675362", 16).unwrap(),
-        BigInt::from_str_radix("b7a71dc9d8e15ea474a69c0531e720cf5474b189ac9afc81", 16).unwrap(),
+        Felt::from_hex_unchecked("590b91a225b1bf7fa5877ec546d090e0059f019c74675362"),
+        Felt::from_hex_unchecked("b7a71dc9d8e15ea474a69c0531e720cf5474b189ac9afc81"),
     )
 )]
 #[case(
@@ -90,13 +104,13 @@ fn test_fft(#[values(true, false)] bit_reversed: bool) {
         16,
     ).unwrap(),
     (
-        BigInt::from_str_radix("51732ecbaa75842afd3d8860d40b3e8eeea433bce18b5c8", 16).unwrap(),
-        BigInt::from_str_radix("a797c1973c99961e357246ee81bde0acbdd27e801d186ccb", 16).unwrap(),
+        Felt::from_hex_unchecked("51732ecbaa75842afd3d8860d40b3e8eeea433bce18b5c8"),
+        Felt::from_hex_unchecked("a797c1973c99961e357246ee81bde0acbdd27e801d186ccb"),
     )
 )]
 fn test_split_commitment_function(
     #[case] commitment: BigInt,
-    #[case] expected_output: (BigInt, BigInt),
+    #[case] expected_output: (Felt, Felt),
 ) {
     let commitment = KzgCommitment::from_bytes(&commitment.to_bytes_be().1).unwrap();
     assert_eq!(split_commitment(&commitment).unwrap(), expected_output);

crates/starknet_os/src/hints/hint_implementation/kzg/utils.rs

@@ -4,9 +4,10 @@ use std::sync::LazyLock;
 
 use blockifier::utils::usize_from_u32;
 use c_kzg::{Blob, KzgCommitment, KzgSettings, BYTES_PER_FIELD_ELEMENT};
-use num_bigint::{BigInt, ParseBigIntError, Sign};
+use num_bigint::{BigInt, ParseBigIntError};
 use num_traits::{Num, One, Zero};
 use starknet_infra_utils::compile_time_cargo_manifest_dir;
+use starknet_types_core::felt::Felt;
 
 const BLOB_SUBGROUP_GENERATOR: &str =
     "39033254847818212395286706435128746857159659164139250548781411570340225835782";
@@ -155,14 +156,14 @@ pub(crate) fn fft(
     Ok(values)
 }
 
-pub(crate) fn split_commitment(commitment: &KzgCommitment) -> Result<(BigInt, BigInt), FftError> {
+pub(crate) fn split_commitment(commitment: &KzgCommitment) -> Result<(Felt, Felt), FftError> {
     let commitment_bytes: [u8; COMMITMENT_BYTES_LENGTH] = *commitment.to_bytes().as_ref();
 
     // Split the number.
     let low = &commitment_bytes[COMMITMENT_BYTES_MIDPOINT..];
     let high = &commitment_bytes[..COMMITMENT_BYTES_MIDPOINT];
 
-    Ok((BigInt::from_bytes_be(Sign::Plus, low), BigInt::from_bytes_be(Sign::Plus, high)))
+    Ok((Felt::from_bytes_be_slice(low), Felt::from_bytes_be_slice(high)))
 }
 
 fn polynomial_coefficients_to_blob(coefficients: Vec<BigInt>) -> Result<Vec<u8>, FftError> {
@@ -184,10 +185,9 @@ fn polynomial_coefficients_to_blob(coefficients: Vec<BigInt>) -> Result<Vec<u8>,
     serialize_blob(&fft_result)
 }
 
-#[allow(dead_code)]
 pub(crate) fn polynomial_coefficients_to_kzg_commitment(
     coefficients: Vec<BigInt>,
-) -> Result<(BigInt, BigInt), FftError> {
+) -> Result<(Felt, Felt), FftError> {
     let blob = polynomial_coefficients_to_blob(coefficients)?;
     let commitment_bytes = blob_to_kzg_commitment(&Blob::from_bytes(&blob)?)?;
     split_commitment(&commitment_bytes)

crates/starknet_os/src/hints/vars.rs

@@ -86,12 +86,14 @@ pub enum Ids {
     CompressedStart,
     ContractAddress,
     ContractStateChanges,
+    DaSize,
     DataEnd,
     DataStart,
     DecompressedDst,
     DictPtr,
     Edge,
     ElmBound,
+    Evals,
     ExecutionContext,
     FinalRoot,
     FullOutput,
@@ -100,7 +102,9 @@ pub enum Ids {
     InitialCarriedOutputs,
     InitialRoot,
     IsLeaf,
+    KzgCommitments,
     Low,
+    NBlobs,
     NCompiledClassFacts,
     NTxs,
     NewLength,
@@ -137,12 +141,14 @@ impl From<Ids> for &'static str {
             Ids::CompressedStart => "compressed_start",
             Ids::ContractAddress => "contract_address",
             Ids::ContractStateChanges => "contract_state_changes",
+            Ids::DaSize => "da_size",
             Ids::DataEnd => "data_end",
             Ids::DataStart => "data_start",
             Ids::DecompressedDst => "decompressed_dst",
             Ids::DictPtr => "dict_ptr",
             Ids::Edge => "edge",
             Ids::ElmBound => "elm_bound",
+            Ids::Evals => "evals",
             Ids::ExecutionContext => "execution_context",
             Ids::FinalRoot => "final_root",
             Ids::FullOutput => "full_output",
@@ -151,7 +157,9 @@ impl From<Ids> for &'static str {
             Ids::InitialCarriedOutputs => "initial_carried_outputs",
             Ids::InitialRoot => "initial_root",
             Ids::IsLeaf => "is_leaf",
+            Ids::KzgCommitments => "kzg_commitments",
             Ids::Low => "low",
+            Ids::NBlobs => "n_blobs",
             Ids::NCompiledClassFacts => "n_compiled_class_facts",
             Ids::NTxs => "n_txs",
             Ids::NewLength => "new_length",
@@ -182,6 +190,7 @@ pub enum Const {
     AliasContractAddress,
     AliasCounterStorageKey,
     Base,
+    BlobLength,
     BlockHashContractAddress,
     CompiledClassVersion,
     InitialAvailableAlias,
@@ -195,6 +204,7 @@ impl From<Const> for &'static str {
             Const::AliasContractAddress => "ALIAS_CONTRACT_ADDRESS",
             Const::AliasCounterStorageKey => "ALIAS_COUNTER_STORAGE_KEY",
             Const::Base => "BASE",
+            Const::BlobLength => "BLOB_LENGTH",
             Const::BlockHashContractAddress => "BLOCK_HASH_CONTRACT_ADDRESS",
             Const::CompiledClassVersion => "COMPILED_CLASS_VERSION",
             Const::InitialAvailableAlias => "INITIAL_AVAILABLE_ALIAS",