1+ use blockifier:: bouncer:: BouncerWeights ;
12use chrono:: { DateTime , SubsecRound , Utc } ;
23#[ cfg( feature = "with_mongodb" ) ]
34use mongodb:: bson:: serde_helpers:: { chrono_datetime_as_bson_datetime, uuid_1_as_binary} ;
@@ -66,6 +67,12 @@ pub struct SnosBatchUpdates {
6667pub status : Option < SnosBatchStatus > ,
6768}
6869
70+ #[ derive( Debug , Clone , Serialize , Deserialize , PartialEq , Eq , Default ) ]
71+ pub struct AggregatorBatchWeights {
72+ pub l1_gas : usize ,
73+ pub message_segment_length : usize ,
74+ }
75+
6976/// Aggregator Batch
7077///
7178/// Represents a high-level batch that contains multiple SNOS batches and manages
@@ -139,13 +146,42 @@ pub struct AggregatorBatch {
139146/// Used to track the batch in the proving system
140147pub bucket_id : String ,
141148
149+ /// Builtin weights for the batch. We decide when to close a batch based on this.
150+ pub builtin_weights : AggregatorBatchWeights ,
151+
142152 /// Current status of the aggregator batch
143153pub status : AggregatorBatchStatus ,
144154 /// Starknet protocol version for all blocks in this batch
145155/// All blocks in a batch must have the same Starknet version for prover compatibility
146156pub starknet_version : String ,
147157}
148158
159+ impl AggregatorBatchWeights {
160+ pub fn new ( l1_gas : usize , message_segment_length : usize ) -> Self {
161+ Self { l1_gas, message_segment_length }
162+ }
163+
164+ pub fn checked_add ( & self , other : & AggregatorBatchWeights ) -> Option < AggregatorBatchWeights > {
165+ Some ( Self {
166+ l1_gas : self . l1_gas . checked_add ( other. l1_gas ) ?,
167+ message_segment_length : self . message_segment_length . checked_add ( other. message_segment_length ) ?,
168+ } )
169+ }
170+
171+ pub fn checked_sub ( & self , other : & AggregatorBatchWeights ) -> Option < AggregatorBatchWeights > {
172+ Some ( Self {
173+ l1_gas : self . l1_gas . checked_sub ( other. l1_gas ) ?,
174+ message_segment_length : self . message_segment_length . checked_sub ( other. message_segment_length ) ?,
175+ } )
176+ }
177+ }
178+
179+ impl From < & BouncerWeights > for AggregatorBatchWeights {
180+ fn from ( weights : & BouncerWeights ) -> Self {
181+ Self { l1_gas : weights. l1_gas , message_segment_length : weights. message_segment_length }
182+ }
183+ }
184+
149185impl AggregatorBatch {
150186 /// Creates a new aggregator batch
151187///
@@ -158,13 +194,15 @@ impl AggregatorBatch {
158194///
159195/// # Returns
160196/// A new `AggregatorBatch` instance with status `Open` and single block
197+ #[ allow( clippy:: too_many_arguments) ]
161198 pub fn new (
162199 index : u64 ,
163200 start_snos_batch : u64 ,
164201 start_block : u64 ,
165202 squashed_state_updates_path : String ,
166203 blob_path : String ,
167204 bucket_id : String ,
205+ builtin_weights : AggregatorBatchWeights ,
168206 starknet_version : String ,
169207 ) -> Self {
170208 Self {
@@ -184,6 +222,7 @@ impl AggregatorBatch {
184222 bucket_id,
185223 starknet_version,
186224 status : AggregatorBatchStatus :: Open ,
225+ builtin_weights,
187226 }
188227 }
189228}
@@ -303,3 +342,143 @@ impl SnosBatch {
303342 Ok ( ( ) )
304343 }
305344}
345+
346+ #[ cfg( test) ]
347+ mod tests {
348+ use super :: * ;
349+
350+ mod aggregator_batch_weights_tests {
351+ use super :: * ;
352+
353+ #[ test]
354+ fn test_new ( ) {
355+ let weights = AggregatorBatchWeights :: new ( 1000 , 500 ) ;
356+ assert_eq ! ( weights. l1_gas, 1000 ) ;
357+ assert_eq ! ( weights. message_segment_length, 500 ) ;
358+ }
359+
360+ #[ test]
361+ fn test_checked_add_success ( ) {
362+ let weights1 = AggregatorBatchWeights :: new ( 1000 , 500 ) ;
363+ let weights2 = AggregatorBatchWeights :: new ( 2000 , 300 ) ;
364+
365+ let result = weights1. checked_add ( & weights2) ;
366+ assert ! ( result. is_some( ) ) ;
367+
368+ let sum = result. unwrap ( ) ;
369+ assert_eq ! ( sum. l1_gas, 3000 ) ;
370+ assert_eq ! ( sum. message_segment_length, 800 ) ;
371+ }
372+
373+ #[ test]
374+ fn test_checked_add_overflow_l1_gas ( ) {
375+ let weights1 = AggregatorBatchWeights :: new ( usize:: MAX , 100 ) ;
376+ let weights2 = AggregatorBatchWeights :: new ( 1 , 100 ) ;
377+
378+ let result = weights1. checked_add ( & weights2) ;
379+ assert ! ( result. is_none( ) ) ;
380+ }
381+
382+ #[ test]
383+ fn test_checked_add_overflow_message_segment_length ( ) {
384+ let weights1 = AggregatorBatchWeights :: new ( 100 , usize:: MAX ) ;
385+ let weights2 = AggregatorBatchWeights :: new ( 100 , 1 ) ;
386+
387+ let result = weights1. checked_add ( & weights2) ;
388+ assert ! ( result. is_none( ) ) ;
389+ }
390+
391+ #[ test]
392+ fn test_checked_add_max_values ( ) {
393+ let weights1 = AggregatorBatchWeights :: new ( usize:: MAX / 2 , usize:: MAX / 2 ) ;
394+ let weights2 = AggregatorBatchWeights :: new ( usize:: MAX / 2 , usize:: MAX / 2 ) ;
395+
396+ let result = weights1. checked_add ( & weights2) ;
397+ assert ! ( result. is_some( ) ) ;
398+
399+ let sum = result. unwrap ( ) ;
400+ assert_eq ! ( sum. l1_gas, usize :: MAX - 1 ) ;
401+ assert_eq ! ( sum. message_segment_length, usize :: MAX - 1 ) ;
402+ }
403+
404+ #[ test]
405+ fn test_checked_sub_success ( ) {
406+ let weights1 = AggregatorBatchWeights :: new ( 2000 , 500 ) ;
407+ let weights2 = AggregatorBatchWeights :: new ( 1000 , 300 ) ;
408+
409+ let result = weights1. checked_sub ( & weights2) ;
410+ assert ! ( result. is_some( ) ) ;
411+
412+ let diff = result. unwrap ( ) ;
413+ assert_eq ! ( diff. l1_gas, 1000 ) ;
414+ assert_eq ! ( diff. message_segment_length, 200 ) ;
415+ }
416+
417+ #[ test]
418+ fn test_checked_sub_with_zero ( ) {
419+ let weights1 = AggregatorBatchWeights :: new ( 1000 , 500 ) ;
420+ let weights2 = AggregatorBatchWeights :: new ( 0 , 0 ) ;
421+
422+ let result = weights1. checked_sub ( & weights2) ;
423+ assert ! ( result. is_some( ) ) ;
424+
425+ let diff = result. unwrap ( ) ;
426+ assert_eq ! ( diff. l1_gas, 1000 ) ;
427+ assert_eq ! ( diff. message_segment_length, 500 ) ;
428+ }
429+
430+ #[ test]
431+ fn test_checked_sub_equal_values ( ) {
432+ let weights1 = AggregatorBatchWeights :: new ( 1000 , 500 ) ;
433+ let weights2 = AggregatorBatchWeights :: new ( 1000 , 500 ) ;
434+
435+ let result = weights1. checked_sub ( & weights2) ;
436+ assert ! ( result. is_some( ) ) ;
437+
438+ let diff = result. unwrap ( ) ;
439+ assert_eq ! ( diff. l1_gas, 0 ) ;
440+ assert_eq ! ( diff. message_segment_length, 0 ) ;
441+ }
442+
443+ #[ test]
444+ fn test_checked_sub_underflow_l1_gas ( ) {
445+ let weights1 = AggregatorBatchWeights :: new ( 100 , 500 ) ;
446+ let weights2 = AggregatorBatchWeights :: new ( 200 , 300 ) ;
447+
448+ let result = weights1. checked_sub ( & weights2) ;
449+ assert ! ( result. is_none( ) ) ;
450+ }
451+
452+ #[ test]
453+ fn test_checked_sub_underflow_message_segment_length ( ) {
454+ let weights1 = AggregatorBatchWeights :: new ( 500 , 100 ) ;
455+ let weights2 = AggregatorBatchWeights :: new ( 300 , 200 ) ;
456+
457+ let result = weights1. checked_sub ( & weights2) ;
458+ assert ! ( result. is_none( ) ) ;
459+ }
460+
461+ #[ test]
462+ fn test_checked_sub_from_max ( ) {
463+ let weights1 = AggregatorBatchWeights :: new ( usize:: MAX , usize:: MAX ) ;
464+ let weights2 = AggregatorBatchWeights :: new ( 1 , 1 ) ;
465+
466+ let result = weights1. checked_sub ( & weights2) ;
467+ assert ! ( result. is_some( ) ) ;
468+
469+ let diff = result. unwrap ( ) ;
470+ assert_eq ! ( diff. l1_gas, usize :: MAX - 1 ) ;
471+ assert_eq ! ( diff. message_segment_length, usize :: MAX - 1 ) ;
472+ }
473+
474+ #[ test]
475+ fn test_from_bouncer_weights ( ) {
476+ let bouncer_weights =
477+ BouncerWeights { l1_gas : 1234 , message_segment_length : usize:: MAX , ..Default :: default ( ) } ;
478+
479+ let agg_weights = AggregatorBatchWeights :: from ( & bouncer_weights) ;
480+ assert_eq ! ( agg_weights. l1_gas, 1234 ) ;
481+ assert_eq ! ( agg_weights. message_segment_length, usize :: MAX ) ;
482+ }
483+ }
484+ }
0 commit comments