@@ -11,6 +11,8 @@ use light_registry::{
1111} ;
1212use solana_sdk:: signature:: { Keypair , Signature , Signer } ;
1313
14+ use crate :: error:: ForesterUtilsError ;
15+
1416// What does the forester need to know?
1517// What are my public keys (current epoch account, last epoch account, known Merkle trees)
1618// 1. The current epoch
@@ -84,11 +86,20 @@ pub fn get_schedule_for_queue(
8486 protocol_config : & ProtocolConfig ,
8587 total_epoch_weight : u64 ,
8688 epoch : u64 ,
87- ) -> Vec < Option < ForesterSlot > > {
89+ current_phase_start_slot : u64 ,
90+ ) -> Result < Vec < Option < ForesterSlot > > , ForesterUtilsError > {
8891 let mut vec = Vec :: new ( ) ;
89- let start_slot = 0 ;
90- // TODO: enforce that active_phase_length is a multiple of slot_length
91- let end_slot = start_slot + ( protocol_config. active_phase_length / protocol_config. slot_length ) ;
92+
93+ let current_light_slot = if start_solana_slot >= current_phase_start_slot {
94+ ( start_solana_slot - current_phase_start_slot) / protocol_config. slot_length
95+ } else {
96+ return Err ( ForesterUtilsError :: InvalidSlotNumber ) ;
97+ } ;
98+
99+ let start_slot = current_light_slot;
100+ start_solana_slot =
101+ current_phase_start_slot + ( current_light_slot * protocol_config. slot_length ) ;
102+ let end_slot = protocol_config. active_phase_length / protocol_config. slot_length ;
92103
93104 for light_slot in start_slot..end_slot {
94105 let forester_index = ForesterEpochPda :: get_eligible_forester_index (
@@ -106,30 +117,31 @@ pub fn get_schedule_for_queue(
106117 } ) ) ;
107118 start_solana_slot += protocol_config. slot_length ;
108119 }
109- vec
120+ Ok ( vec)
110121}
111122
112123pub fn get_schedule_for_forester_in_queue (
113124 start_solana_slot : u64 ,
114125 queue_pubkey : & Pubkey ,
115126 total_epoch_weight : u64 ,
116127 forester_epoch_pda : & ForesterEpochPda ,
117- ) -> Vec < Option < ForesterSlot > > {
128+ ) -> Result < Vec < Option < ForesterSlot > > , ForesterUtilsError > {
118129 let mut slots = get_schedule_for_queue (
119130 start_solana_slot,
120131 queue_pubkey,
121132 & forester_epoch_pda. protocol_config ,
122133 total_epoch_weight,
123134 forester_epoch_pda. epoch ,
124- ) ;
135+ forester_epoch_pda. epoch_active_phase_start_slot ,
136+ ) ?;
125137 slots. iter_mut ( ) . for_each ( |slot_option| {
126138 if let Some ( slot) = slot_option {
127139 if !forester_epoch_pda. is_eligible ( slot. forester_index ) {
128140 * slot_option = None ;
129141 }
130142 }
131143 } ) ;
132- slots
144+ Ok ( slots)
133145}
134146
135147#[ derive( Debug , Clone , PartialEq , Eq ) ]
@@ -153,7 +165,7 @@ impl TreeForesterSchedule {
153165 solana_slot : u64 ,
154166 forester_epoch_pda : & ForesterEpochPda ,
155167 epoch_pda : & EpochPda ,
156- ) -> Self {
168+ ) -> Result < Self , ForesterUtilsError > {
157169 let mut _self = Self {
158170 tree_accounts : * tree_accounts,
159171 slots : Vec :: new ( ) ,
@@ -163,8 +175,8 @@ impl TreeForesterSchedule {
163175 & _self. tree_accounts . queue ,
164176 epoch_pda. registered_weight ,
165177 forester_epoch_pda,
166- ) ;
167- _self
178+ ) ? ;
179+ Ok ( _self)
168180 }
169181
170182 pub fn is_eligible ( & self , forester_slot : u64 ) -> bool {
@@ -211,6 +223,12 @@ pub struct Phase {
211223 pub end : u64 ,
212224}
213225
226+ impl Phase {
227+ pub fn length ( & self ) -> u64 {
228+ self . end - self . start
229+ }
230+ }
231+
214232pub fn get_epoch_phases ( protocol_config : & ProtocolConfig , epoch : u64 ) -> EpochPhases {
215233 let epoch_start_slot = protocol_config
216234 . genesis_slot
@@ -382,7 +400,11 @@ impl Epoch {
382400 if forester_epoch_pda. total_epoch_weight . is_none ( ) {
383401 forester_epoch_pda. total_epoch_weight = Some ( epoch_pda. registered_weight ) ;
384402 }
385- self . add_trees_with_schedule ( & forester_epoch_pda, & epoch_pda, trees, current_solana_slot) ;
403+ self . add_trees_with_schedule ( & forester_epoch_pda, & epoch_pda, trees, current_solana_slot)
404+ . map_err ( |e| {
405+ println ! ( "Error adding trees with schedule: {:?}" , e) ;
406+ RpcError :: AssertRpcError ( "Error adding trees with schedule" . to_string ( ) )
407+ } ) ?;
386408 Ok ( ( ) )
387409 }
388410 /// Internal function to init Epoch struct with registered account
@@ -395,7 +417,7 @@ impl Epoch {
395417 epoch_pda : & EpochPda ,
396418 trees : & [ TreeAccounts ] ,
397419 current_solana_slot : u64 ,
398- ) {
420+ ) -> Result < ( ) , ForesterUtilsError > {
399421 // let state = self.phases.get_current_epoch_state(current_solana_slot);
400422 // TODO: add epoch state to sync schedule
401423 for tree in trees {
@@ -404,9 +426,10 @@ impl Epoch {
404426 current_solana_slot,
405427 forester_epoch_pda,
406428 epoch_pda,
407- ) ;
429+ ) ? ;
408430 self . merkle_trees . push ( tree_schedule) ;
409431 }
432+ Ok ( ( ) )
410433 }
411434
412435 pub fn update_state ( & mut self , current_solana_slot : u64 ) -> EpochState {
@@ -490,14 +513,22 @@ mod test {
490513 let queue_pubkey = Pubkey :: new_unique ( ) ;
491514 let start_solana_slot = 0 ;
492515 let epoch = 0 ;
516+ let current_phase_start_slot = 0 ;
493517
494518 let schedule = get_schedule_for_queue (
495519 start_solana_slot,
496520 & queue_pubkey,
497521 & protocol_config,
498522 total_epoch_weight,
499523 epoch,
500- ) ;
524+ current_phase_start_slot,
525+ )
526+ . unwrap ( ) ;
527+
528+ // Expected number of light slots in the active phase
529+ let expected_light_slots =
530+ ( protocol_config. active_phase_length / protocol_config. slot_length ) as usize ;
531+ assert_eq ! ( schedule. len( ) , expected_light_slots) ; // Should generate 100 slots
501532
502533 assert_eq ! (
503534 schedule. len( ) ,
@@ -518,4 +549,142 @@ mod test {
518549 assert ! ( slot. forester_index < total_epoch_weight) ;
519550 }
520551 }
552+
553+ #[ test]
554+ fn test_get_schedule_for_queue_offset_phase_start ( ) {
555+ let protocol_config = ProtocolConfig {
556+ genesis_slot : 1000 , // Genesis starts later
557+ min_weight : 100 ,
558+ slot_length : 10 ,
559+ registration_phase_length : 100 ,
560+ active_phase_length : 1000 , // 100 light slots
561+ report_work_phase_length : 100 ,
562+ network_fee : 5000 ,
563+ ..Default :: default ( )
564+ } ;
565+
566+ let total_epoch_weight = 500 ;
567+ let queue_pubkey = Pubkey :: new_unique ( ) ;
568+ let epoch = 0 ;
569+
570+ // Calculate actual start of the active phase for epoch 0
571+ // Registration: 1000 to 1099
572+ // Active: 1100 to 2099
573+ let current_phase_start_slot = 1100 ;
574+
575+ // Start calculating right from the beginning of this active phase
576+ let start_solana_slot = current_phase_start_slot;
577+
578+ let schedule = get_schedule_for_queue (
579+ start_solana_slot,
580+ & queue_pubkey,
581+ & protocol_config,
582+ total_epoch_weight,
583+ epoch,
584+ current_phase_start_slot, // Pass the calculated start slot
585+ )
586+ . unwrap ( ) ;
587+
588+ let expected_light_slots =
589+ ( protocol_config. active_phase_length / protocol_config. slot_length ) as usize ;
590+ assert_eq ! ( schedule. len( ) , expected_light_slots) ; // Still 100 light slots expected
591+
592+ // Check the first slot details
593+ let first_slot = schedule[ 0 ] . as_ref ( ) . unwrap ( ) ;
594+ assert_eq ! ( first_slot. slot, 0 ) ; // First light slot index is 0
595+ // Its Solana start slot should be the phase start slot
596+ assert_eq ! ( first_slot. start_solana_slot, current_phase_start_slot) ;
597+ assert_eq ! (
598+ first_slot. end_solana_slot,
599+ current_phase_start_slot + protocol_config. slot_length
600+ ) ;
601+
602+ // Check the second slot details
603+ let second_slot = schedule[ 1 ] . as_ref ( ) . unwrap ( ) ;
604+ assert_eq ! ( second_slot. slot, 1 ) ; // Second light slot index is 1
605+ // Its Solana start slot should be offset by one slot_length
606+ assert_eq ! (
607+ second_slot. start_solana_slot,
608+ current_phase_start_slot + protocol_config. slot_length
609+ ) ;
610+ assert_eq ! (
611+ second_slot. end_solana_slot,
612+ current_phase_start_slot + 2 * protocol_config. slot_length
613+ ) ;
614+ }
615+
616+ // NEW TEST: Case where current_light_slot > 0
617+ #[ test]
618+ fn test_get_schedule_for_queue_mid_phase_start ( ) {
619+ let protocol_config = ProtocolConfig {
620+ genesis_slot : 0 ,
621+ min_weight : 100 ,
622+ slot_length : 10 ,
623+ registration_phase_length : 100 , // Reg: 0-99
624+ active_phase_length : 1000 , // Active: 100-1099 (100 light slots)
625+ report_work_phase_length : 100 ,
626+ network_fee : 5000 ,
627+ ..Default :: default ( )
628+ } ;
629+
630+ let total_epoch_weight = 500 ;
631+ let queue_pubkey = Pubkey :: new_unique ( ) ;
632+ let epoch = 0 ;
633+ let current_phase_start_slot = 100 ; // Active phase starts at slot 100
634+
635+ // Start calculating from Solana slot 155, which is within the active phase
636+ let start_solana_slot = 155 ;
637+
638+ // Calculation:
639+ // current_light_slot = floor((155 - 100) / 10) = floor(55 / 10) = 5
640+ // Effective start_solana_slot for loop = 100 + (5 * 10) = 150
641+ // End light slot = 1000 / 10 = 100
642+ // Loop runs from light_slot 5 to 99 (inclusive). Length = 100 - 5 = 95
643+
644+ let schedule = get_schedule_for_queue (
645+ start_solana_slot,
646+ & queue_pubkey,
647+ & protocol_config,
648+ total_epoch_weight,
649+ epoch,
650+ current_phase_start_slot,
651+ )
652+ . unwrap ( ) ;
653+
654+ let expected_light_slots_total =
655+ protocol_config. active_phase_length / protocol_config. slot_length ; // 100
656+ let expected_start_light_slot = 5 ;
657+ let expected_schedule_len =
658+ ( expected_light_slots_total - expected_start_light_slot) as usize ; // 100 - 5 = 95
659+
660+ assert_eq ! ( schedule. len( ) , expected_schedule_len) ; // Should generate 95 slots
661+
662+ // Check the first slot in the *returned* schedule
663+ let first_returned_slot = schedule[ 0 ] . as_ref ( ) . unwrap ( ) ;
664+ assert_eq ! ( first_returned_slot. slot, expected_start_light_slot) ; // Light slot index starts at 5
665+ // Its Solana start slot should align to the beginning of light slot 5
666+ let expected_first_solana_start =
667+ current_phase_start_slot + expected_start_light_slot * protocol_config. slot_length ; // 100 + 5 * 10 = 150
668+ assert_eq ! (
669+ first_returned_slot. start_solana_slot,
670+ expected_first_solana_start
671+ ) ;
672+ assert_eq ! (
673+ first_returned_slot. end_solana_slot,
674+ expected_first_solana_start + protocol_config. slot_length // 150 + 10 = 160
675+ ) ;
676+
677+ // Check the second slot in the *returned* schedule
678+ let second_returned_slot = schedule[ 1 ] . as_ref ( ) . unwrap ( ) ;
679+ assert_eq ! ( second_returned_slot. slot, expected_start_light_slot + 1 ) ; // Light slot index 6
680+ // Its Solana start slot should be 160
681+ assert_eq ! (
682+ second_returned_slot. start_solana_slot,
683+ expected_first_solana_start + protocol_config. slot_length
684+ ) ;
685+ assert_eq ! (
686+ second_returned_slot. end_solana_slot,
687+ expected_first_solana_start + 2 * protocol_config. slot_length // 170
688+ ) ;
689+ }
521690}
0 commit comments