derive lookup_data

starkware-libs · Dec 18, 2024 · b65cb18 · b65cb18
1 parent f1cbcef
commit b65cb18
Show file tree

Hide file tree

Showing 26 changed files with 98 additions and 455 deletions.
diff --git a/stwo_cairo_prover/crates/air_structs_derive/src/lib.rs b/stwo_cairo_prover/crates/air_structs_derive/src/lib.rs
@@ -1,13 +1,15 @@
 use itertools::Itertools;
 use proc_macro2::TokenStream;
 use quote::quote;
-use syn::{parse_macro_input, Data, DeriveInput, Fields, Path, Type};
+use syn::{parse_macro_input, Data, DataStruct, DeriveInput, Fields, Path, Type};
+
+// TODO(Ohad): consider moving this crate to 'Stwo'.
 
 #[proc_macro_derive(SubComponentInputs)]
 pub fn derive_sub_component_inputs(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let input = parse_macro_input!(input as DeriveInput);
-    assert_is_struct(&input);
     let name = &input.ident;
+    let input = expect_struct(&input);
     let vec_array_fields = extract_vec_array_fields(&input);
 
     // TODO(Ohad): deprecate with_capacity.
@@ -26,40 +28,61 @@ pub fn derive_sub_component_inputs(input: proc_macro::TokenStream) -> proc_macro
     proc_macro::TokenStream::from(expanded)
 }
 
+#[proc_macro_derive(LookupData)]
+pub fn derive_lookup_data(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let input = parse_macro_input!(input as DeriveInput);
+    let name = &input.ident;
+    let input = expect_struct(&input);
+    let vec_array_fields = extract_vec_array_fields(&input);
+
+    // TODO(Ohad): deprecate with_capacity.
+    let with_capacity_method = generate_with_capacity_method(&vec_array_fields);
+    let uninitialized_method = generate_uninitialized_method(&vec_array_fields);
+
+    let expanded = quote! {
+        impl #name {
+            #with_capacity_method
+            #uninitialized_method
+        }
+    };
+
+    proc_macro::TokenStream::from(expanded)
+}
+
 #[derive(Clone)]
 struct VecArrayField {
     name: syn::Ident,
     array_length: usize,
 }
 
-fn assert_is_struct(input: &DeriveInput) {
-    if !matches!(input.data, Data::Struct(_)) {
-        panic!("Derive(SubComponentInputs) can only be applied to structs.");
+fn expect_struct(input: &DeriveInput) -> DataStruct {
+    if let Data::Struct(data_struct) = &input.data {
+        data_struct.clone()
+    } else {
+        panic!("Expected a struct");
     }
 }
 
-fn extract_vec_array_fields(input: &DeriveInput) -> Vec<VecArrayField> {
+fn extract_vec_array_fields(data_struct: &DataStruct) -> Vec<VecArrayField> {
     let mut vec_array_fields = Vec::new();
 
-    if let Data::Struct(data_struct) = &input.data {
-        if let Fields::Named(fields) = &data_struct.fields {
-            for field in &fields.named {
-                // Field is an array of Vecs.
-                if let Type::Array(type_array) = &field.ty {
-                    if let Type::Path(element_type) = &*type_array.elem {
-                        // Element is a Vec.
-                        if is_vec_type(&element_type.path) {
-                            // Get the array length
-                            if let syn::Expr::Lit(syn::ExprLit {
-                                lit: syn::Lit::Int(length_lit),
-                                ..
-                            }) = type_array.len.clone()
-                            {
-                                vec_array_fields.push(VecArrayField {
-                                    name: field.ident.clone().unwrap(),
-                                    array_length: length_lit.base10_parse().unwrap(),
-                                });
-                            }
+    if let Fields::Named(fields) = &data_struct.fields {
+        for field in &fields.named {
+            // Field is an array of Vecs.
+            if let Type::Array(type_array) = &field.ty {
+                if let Type::Path(element_type) = &*type_array.elem {
+                    // Element is a Vec.
+                    if is_vec_type(&element_type.path) {
+                        // Get the array length
+                        if let syn::Expr::Lit(syn::ExprLit {
+                            lit: syn::Lit::Int(length_lit),
+                            ..
+                        }) = type_array.len.clone()
+                        {
+                            vec_array_fields.push(VecArrayField {
+                                name: field.ident.clone().unwrap(),
+                                array_length: length_lit.base10_parse().unwrap(),
+                            });
                         }
                     }
                 }

diff --git a/..._cairo_prover/crates/prover/src/components/add_ap_opcode_is_imm_f_op1_base_fp_f/prover.rs b/..._cairo_prover/crates/prover/src/components/add_ap_opcode_is_imm_f_op1_base_fp_f/prover.rs
@@ -1,6 +1,6 @@
 #![allow(unused_parens)]
 #![allow(unused_imports)]
-use air_structs_derive::SubComponentInputs;
+use air_structs_derive::{LookupData, SubComponentInputs};
 use itertools::{chain, zip_eq, Itertools};
 use num_traits::{One, Zero};
 use prover_types::cpu::*;
@@ -301,23 +301,13 @@ pub fn write_trace_simd(
     (trace, sub_components_inputs, lookup_data)
 }
 
+#[derive(LookupData)]
 pub struct LookupData {
     pub memoryaddresstoid: [Vec<[PackedM31; 2]>; 1],
     pub memoryidtobig: [Vec<[PackedM31; 29]>; 1],
     pub opcodes: [Vec<[PackedM31; 3]>; 2],
     pub verifyinstruction: [Vec<[PackedM31; 19]>; 1],
 }
-impl LookupData {
-    #[allow(unused_variables)]
-    fn with_capacity(capacity: usize) -> Self {
-        Self {
-            memoryaddresstoid: [Vec::with_capacity(capacity)],
-            memoryidtobig: [Vec::with_capacity(capacity)],
-            opcodes: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-            verifyinstruction: [Vec::with_capacity(capacity)],
-        }
-    }
-}
 
 pub struct InteractionClaimGenerator {
     pub n_calls: usize,

diff --git a/..._cairo_prover/crates/prover/src/components/add_ap_opcode_is_imm_f_op1_base_fp_t/prover.rs b/..._cairo_prover/crates/prover/src/components/add_ap_opcode_is_imm_f_op1_base_fp_t/prover.rs
@@ -1,6 +1,6 @@
 #![allow(unused_parens)]
 #![allow(unused_imports)]
-use air_structs_derive::SubComponentInputs;
+use air_structs_derive::{LookupData, SubComponentInputs};
 use itertools::{chain, zip_eq, Itertools};
 use num_traits::{One, Zero};
 use prover_types::cpu::*;
@@ -301,23 +301,13 @@ pub fn write_trace_simd(
     (trace, sub_components_inputs, lookup_data)
 }
 
+#[derive(LookupData)]
 pub struct LookupData {
     pub memoryaddresstoid: [Vec<[PackedM31; 2]>; 1],
     pub memoryidtobig: [Vec<[PackedM31; 29]>; 1],
     pub opcodes: [Vec<[PackedM31; 3]>; 2],
     pub verifyinstruction: [Vec<[PackedM31; 19]>; 1],
 }
-impl LookupData {
-    #[allow(unused_variables)]
-    fn with_capacity(capacity: usize) -> Self {
-        Self {
-            memoryaddresstoid: [Vec::with_capacity(capacity)],
-            memoryidtobig: [Vec::with_capacity(capacity)],
-            opcodes: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-            verifyinstruction: [Vec::with_capacity(capacity)],
-        }
-    }
-}
 
 pub struct InteractionClaimGenerator {
     pub n_calls: usize,

diff --git a/..._cairo_prover/crates/prover/src/components/add_ap_opcode_is_imm_t_op1_base_fp_f/prover.rs b/..._cairo_prover/crates/prover/src/components/add_ap_opcode_is_imm_t_op1_base_fp_f/prover.rs
@@ -1,6 +1,6 @@
 #![allow(unused_parens)]
 #![allow(unused_imports)]
-use air_structs_derive::SubComponentInputs;
+use air_structs_derive::{LookupData, SubComponentInputs};
 use itertools::{chain, zip_eq, Itertools};
 use num_traits::{One, Zero};
 use prover_types::cpu::*;
@@ -285,23 +285,13 @@ pub fn write_trace_simd(
     (trace, sub_components_inputs, lookup_data)
 }
 
+#[derive(LookupData)]
 pub struct LookupData {
     pub memoryaddresstoid: [Vec<[PackedM31; 2]>; 1],
     pub memoryidtobig: [Vec<[PackedM31; 29]>; 1],
     pub opcodes: [Vec<[PackedM31; 3]>; 2],
     pub verifyinstruction: [Vec<[PackedM31; 19]>; 1],
 }
-impl LookupData {
-    #[allow(unused_variables)]
-    fn with_capacity(capacity: usize) -> Self {
-        Self {
-            memoryaddresstoid: [Vec::with_capacity(capacity)],
-            memoryidtobig: [Vec::with_capacity(capacity)],
-            opcodes: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-            verifyinstruction: [Vec::with_capacity(capacity)],
-        }
-    }
-}
 
 pub struct InteractionClaimGenerator {
     pub n_calls: usize,

diff --git a/stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_f_is_imm_f/prover.rs b/stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_f_is_imm_f/prover.rs
@@ -1,6 +1,6 @@
 #![allow(unused_parens)]
 #![allow(unused_imports)]
-use air_structs_derive::SubComponentInputs;
+use air_structs_derive::{LookupData, SubComponentInputs};
 use itertools::{chain, zip_eq, Itertools};
 use num_traits::{One, Zero};
 use prover_types::cpu::*;
@@ -658,31 +658,13 @@ pub fn write_trace_simd(
     (trace, sub_components_inputs, lookup_data)
 }
 
+#[derive(LookupData)]
 pub struct LookupData {
     pub memoryaddresstoid: [Vec<[PackedM31; 2]>; 3],
     pub memoryidtobig: [Vec<[PackedM31; 29]>; 3],
     pub opcodes: [Vec<[PackedM31; 3]>; 2],
     pub verifyinstruction: [Vec<[PackedM31; 19]>; 1],
 }
-impl LookupData {
-    #[allow(unused_variables)]
-    fn with_capacity(capacity: usize) -> Self {
-        Self {
-            memoryaddresstoid: [
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-            ],
-            memoryidtobig: [
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-            ],
-            opcodes: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-            verifyinstruction: [Vec::with_capacity(capacity)],
-        }
-    }
-}
 
 pub struct InteractionClaimGenerator {
     pub n_calls: usize,

diff --git a/stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_f_is_imm_t/prover.rs b/stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_f_is_imm_t/prover.rs
@@ -1,6 +1,6 @@
 #![allow(unused_parens)]
 #![allow(unused_imports)]
-use air_structs_derive::SubComponentInputs;
+use air_structs_derive::{LookupData, SubComponentInputs};
 use itertools::{chain, zip_eq, Itertools};
 use num_traits::{One, Zero};
 use prover_types::cpu::*;
@@ -620,31 +620,13 @@ pub fn write_trace_simd(
     (trace, sub_components_inputs, lookup_data)
 }
 
+#[derive(LookupData)]
 pub struct LookupData {
     pub memoryaddresstoid: [Vec<[PackedM31; 2]>; 3],
     pub memoryidtobig: [Vec<[PackedM31; 29]>; 3],
     pub opcodes: [Vec<[PackedM31; 3]>; 2],
     pub verifyinstruction: [Vec<[PackedM31; 19]>; 1],
 }
-impl LookupData {
-    #[allow(unused_variables)]
-    fn with_capacity(capacity: usize) -> Self {
-        Self {
-            memoryaddresstoid: [
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-            ],
-            memoryidtobig: [
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-            ],
-            opcodes: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-            verifyinstruction: [Vec::with_capacity(capacity)],
-        }
-    }
-}
 
 pub struct InteractionClaimGenerator {
     pub n_calls: usize,

diff --git a/stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_t_is_imm_f/prover.rs b/stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_t_is_imm_f/prover.rs
@@ -1,6 +1,6 @@
 #![allow(unused_parens)]
 #![allow(unused_imports)]
-use air_structs_derive::SubComponentInputs;
+use air_structs_derive::{LookupData, SubComponentInputs};
 use itertools::{chain, zip_eq, Itertools};
 use num_traits::{One, Zero};
 use prover_types::cpu::*;
@@ -532,31 +532,13 @@ pub fn write_trace_simd(
     (trace, sub_components_inputs, lookup_data)
 }
 
+#[derive(LookupData)]
 pub struct LookupData {
     pub memoryaddresstoid: [Vec<[PackedM31; 2]>; 3],
     pub memoryidtobig: [Vec<[PackedM31; 29]>; 3],
     pub opcodes: [Vec<[PackedM31; 3]>; 2],
     pub verifyinstruction: [Vec<[PackedM31; 19]>; 1],
 }
-impl LookupData {
-    #[allow(unused_variables)]
-    fn with_capacity(capacity: usize) -> Self {
-        Self {
-            memoryaddresstoid: [
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-            ],
-            memoryidtobig: [
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-            ],
-            opcodes: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-            verifyinstruction: [Vec::with_capacity(capacity)],
-        }
-    }
-}
 
 pub struct InteractionClaimGenerator {
     pub n_calls: usize,

diff --git a/stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_t_is_imm_t/prover.rs b/stwo_cairo_prover/crates/prover/src/components/add_opcode_is_small_t_is_imm_t/prover.rs
@@ -1,6 +1,6 @@
 #![allow(unused_parens)]
 #![allow(unused_imports)]
-use air_structs_derive::SubComponentInputs;
+use air_structs_derive::{LookupData, SubComponentInputs};
 use itertools::{chain, zip_eq, Itertools};
 use num_traits::{One, Zero};
 use prover_types::cpu::*;
@@ -494,31 +494,13 @@ pub fn write_trace_simd(
     (trace, sub_components_inputs, lookup_data)
 }
 
+#[derive(LookupData)]
 pub struct LookupData {
     pub memoryaddresstoid: [Vec<[PackedM31; 2]>; 3],
     pub memoryidtobig: [Vec<[PackedM31; 29]>; 3],
     pub opcodes: [Vec<[PackedM31; 3]>; 2],
     pub verifyinstruction: [Vec<[PackedM31; 19]>; 1],
 }
-impl LookupData {
-    #[allow(unused_variables)]
-    fn with_capacity(capacity: usize) -> Self {
-        Self {
-            memoryaddresstoid: [
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-            ],
-            memoryidtobig: [
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-                Vec::with_capacity(capacity),
-            ],
-            opcodes: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-            verifyinstruction: [Vec::with_capacity(capacity)],
-        }
-    }
-}
 
 pub struct InteractionClaimGenerator {
     pub n_calls: usize,

diff --git a/...prover/crates/prover/src/components/assert_eq_opcode_is_double_deref_f_is_imm_f/prover.rs b/...prover/crates/prover/src/components/assert_eq_opcode_is_double_deref_f_is_imm_f/prover.rs
@@ -1,6 +1,6 @@
 #![allow(unused_parens)]
 #![allow(unused_imports)]
-use air_structs_derive::SubComponentInputs;
+use air_structs_derive::{LookupData, SubComponentInputs};
 use itertools::{chain, zip_eq, Itertools};
 use num_traits::{One, Zero};
 use prover_types::cpu::*;
@@ -316,21 +316,12 @@ pub fn write_trace_simd(
     (trace, sub_components_inputs, lookup_data)
 }
 
+#[derive(LookupData)]
 pub struct LookupData {
     pub memoryaddresstoid: [Vec<[PackedM31; 2]>; 2],
     pub opcodes: [Vec<[PackedM31; 3]>; 2],
     pub verifyinstruction: [Vec<[PackedM31; 19]>; 1],
 }
-impl LookupData {
-    #[allow(unused_variables)]
-    fn with_capacity(capacity: usize) -> Self {
-        Self {
-            memoryaddresstoid: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-            opcodes: [Vec::with_capacity(capacity), Vec::with_capacity(capacity)],
-            verifyinstruction: [Vec::with_capacity(capacity)],
-        }
-    }
-}
 
 pub struct InteractionClaimGenerator {
     pub n_calls: usize,