ab_client_proof_of_time/source/
state.rs

1use crate::PotNextSlotInput;
2use crate::verifier::PotVerifier;
3use ab_core_primitives::pot::{PotOutput, PotParametersChange, SlotNumber};
4use parking_lot::Mutex;
5
6#[derive(Debug, Copy, Clone, Eq, PartialEq)]
7struct InnerState {
8    next_slot_input: PotNextSlotInput,
9    parameters_change: Option<PotParametersChange>,
10}
11
12impl InnerState {
13    #[expect(
14        clippy::option_option,
15        reason = "External option indicates whether the change should be done, while internal is \
16        the value itself"
17    )]
18    fn update(
19        mut self,
20        mut slot: SlotNumber,
21        mut output: PotOutput,
22        maybe_updated_parameters_change: Option<Option<PotParametersChange>>,
23        pot_verifier: &PotVerifier,
24    ) -> Self {
25        if let Some(updated_parameters_change) = maybe_updated_parameters_change {
26            self.parameters_change = updated_parameters_change;
27        }
28
29        loop {
30            self.next_slot_input = PotNextSlotInput::derive(
31                self.next_slot_input.slot_iterations,
32                slot,
33                output,
34                &self.parameters_change,
35            );
36
37            // Advance further as far as possible using previously verified proofs/checkpoints
38            if let Some(checkpoints) = pot_verifier.try_get_checkpoints(
39                self.next_slot_input.slot_iterations,
40                self.next_slot_input.seed,
41            ) {
42                slot = self.next_slot_input.slot;
43                output = checkpoints.output();
44            } else {
45                break;
46            }
47        }
48
49        self
50    }
51}
52
53/// Result of [`PotState::set_known_good_output()`] call
54#[derive(Debug)]
55pub enum PotStateSetOutcome {
56    /// Nothing has changed
57    NoChange,
58    /// PoT chain extension
59    Extension {
60        from: PotNextSlotInput,
61        to: PotNextSlotInput,
62    },
63    /// PoT chain reorg
64    Reorg {
65        from: PotNextSlotInput,
66        to: PotNextSlotInput,
67    },
68}
69
70/// Global PoT state.
71///
72/// Maintains the accurate information about PoT state and current tip.
73#[derive(Debug)]
74pub struct PotState {
75    inner_state: Mutex<InnerState>,
76    verifier: PotVerifier,
77}
78
79impl PotState {
80    /// Create a new PoT state
81    pub fn new(
82        next_slot_input: PotNextSlotInput,
83        parameters_change: Option<PotParametersChange>,
84        verifier: PotVerifier,
85    ) -> Self {
86        let inner = InnerState {
87            next_slot_input,
88            parameters_change,
89        };
90
91        Self {
92            inner_state: Mutex::new(inner),
93            verifier,
94        }
95    }
96
97    /// PoT input for the next slot
98    pub fn next_slot_input(&self) -> PotNextSlotInput {
99        self.inner_state.lock().next_slot_input
100    }
101
102    /// Extend PoT chain if it matches provided expected next slot input.
103    ///
104    /// Returns `Ok(new_next_slot_input)` if PoT chain was extended successfully and
105    /// `Err(existing_next_slot_input)` in case the state was changed in the meantime.
106    pub fn try_extend(
107        &self,
108        expected_existing_next_slot_input: PotNextSlotInput,
109        best_slot: SlotNumber,
110        best_output: PotOutput,
111        maybe_updated_parameters_change: Option<Option<PotParametersChange>>,
112    ) -> Result<PotNextSlotInput, PotNextSlotInput> {
113        let mut existing_inner_state = self.inner_state.lock();
114        if expected_existing_next_slot_input != existing_inner_state.next_slot_input {
115            return Err(existing_inner_state.next_slot_input);
116        }
117
118        *existing_inner_state = existing_inner_state.update(
119            best_slot,
120            best_output,
121            maybe_updated_parameters_change,
122            &self.verifier,
123        );
124
125        Ok(existing_inner_state.next_slot_input)
126    }
127
128    /// Set known good output for time slot, overriding PoT chain if it doesn't match the provided
129    /// output.
130    ///
131    /// This is typically called with information obtained from received block. It typically lags
132    /// behind PoT tip and is used as a correction mechanism in case PoT reorg is needed.
133    pub fn set_known_good_output(
134        &self,
135        slot: SlotNumber,
136        output: PotOutput,
137        updated_parameters_change: Option<PotParametersChange>,
138    ) -> PotStateSetOutcome {
139        let previous_best_state;
140        let new_best_state;
141        {
142            let mut inner_state = self.inner_state.lock();
143            previous_best_state = *inner_state;
144            new_best_state = previous_best_state.update(
145                slot,
146                output,
147                Some(updated_parameters_change),
148                &self.verifier,
149            );
150            *inner_state = new_best_state;
151        }
152
153        if previous_best_state.next_slot_input == new_best_state.next_slot_input {
154            return PotStateSetOutcome::NoChange;
155        }
156
157        if previous_best_state.next_slot_input.slot < new_best_state.next_slot_input.slot {
158            let mut slot_iterations = previous_best_state.next_slot_input.slot_iterations;
159            let mut seed = previous_best_state.next_slot_input.seed;
160
161            for slot in
162                previous_best_state.next_slot_input.slot..new_best_state.next_slot_input.slot
163            {
164                let Some(checkpoints) = self.verifier.try_get_checkpoints(slot_iterations, seed)
165                else {
166                    break;
167                };
168
169                let pot_input = PotNextSlotInput::derive(
170                    slot_iterations,
171                    slot,
172                    checkpoints.output(),
173                    &updated_parameters_change,
174                );
175
176                // TODO: Consider carrying of the whole `PotNextSlotInput` rather than individual
177                //  variables
178                let next_slot = slot + SlotNumber::ONE;
179                slot_iterations = pot_input.slot_iterations;
180                seed = pot_input.seed;
181
182                if next_slot == new_best_state.next_slot_input.slot
183                    && slot_iterations == new_best_state.next_slot_input.slot_iterations
184                    && seed == new_best_state.next_slot_input.seed
185                {
186                    return PotStateSetOutcome::Extension {
187                        from: previous_best_state.next_slot_input,
188                        to: new_best_state.next_slot_input,
189                    };
190                }
191            }
192        }
193
194        PotStateSetOutcome::Reorg {
195            from: previous_best_state.next_slot_input,
196            to: new_best_state.next_slot_input,
197        }
198    }
199}
ab_client_proof_of_time/source/state.rs

ab_client_proof_of_time/source/
state.rs
