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

ab_client_proof_of_time/source/
state.rs
