1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

use {
    super::{Node, Snapshot},
    crate::{
        errors::ClockworkError,
        objects::{NodeAccount, SnapshotAccount, SnapshotStatus},
    },
    anchor_lang::{prelude::*, AnchorDeserialize},
    anchor_spl::token::TokenAccount,
    std::convert::TryFrom,
};

const SEED_REGISTRY: &[u8] = b"registry";

/**
 * Registry
 */

#[account]
#[derive(Debug)]
pub struct Registry {
    pub is_locked: bool,
    pub node_count: u64,
    pub snapshot_count: u64,
}

impl Registry {
    pub fn pubkey() -> Pubkey {
        Pubkey::find_program_address(&[SEED_REGISTRY], &crate::ID).0
    }
}

impl TryFrom<Vec<u8>> for Registry {
    type Error = Error;
    fn try_from(data: Vec<u8>) -> std::result::Result<Self, Self::Error> {
        Registry::try_deserialize(&mut data.as_slice())
    }
}

/**
 * RegistryAccount
 */

pub trait RegistryAccount {
    fn init(&mut self) -> Result<()>;

    fn new_node(
        &mut self,
        authority: &mut Signer,
        node: &mut Account<Node>,
        stake: &mut Account<TokenAccount>,
        worker: &Signer,
    ) -> Result<()>;

    fn new_snapshot(&mut self, snapshot: &mut Account<Snapshot>) -> Result<()>;

    fn rotate_snapshot(
        &mut self,
        current_snapshot: Option<&mut Account<Snapshot>>,
        next_snapshot: &mut Account<Snapshot>,
    ) -> Result<()>;

    fn lock(&mut self) -> Result<()>;

    fn unlock(&mut self) -> Result<()>;
}

impl RegistryAccount for Account<'_, Registry> {
    fn init(&mut self) -> Result<()> {
        self.is_locked = false;
        self.node_count = 0;
        self.snapshot_count = 0;
        Ok(())
    }

    fn new_node(
        &mut self,
        authority: &mut Signer,
        node: &mut Account<Node>,
        stake: &mut Account<TokenAccount>,
        worker: &Signer,
    ) -> Result<()> {
        require!(!self.is_locked, ClockworkError::RegistryLocked);
        node.init(authority, self.node_count, stake, worker)?;
        self.node_count = self.node_count.checked_add(1).unwrap();
        Ok(())
    }

    fn new_snapshot(&mut self, snapshot: &mut Account<Snapshot>) -> Result<()> {
        require!(!self.is_locked, ClockworkError::RegistryLocked);
        self.lock()?;
        snapshot.init(self.snapshot_count)?;
        Ok(())
    }

    fn rotate_snapshot(
        &mut self,
        current_snapshot: Option<&mut Account<Snapshot>>,
        next_snapshot: &mut Account<Snapshot>,
    ) -> Result<()> {
        // Require the registry is locked
        require!(self.is_locked, ClockworkError::RegistryMustBeLocked);

        // Validate the next snapshot is in progress
        require!(
            next_snapshot.status == SnapshotStatus::InProgress,
            ClockworkError::SnapshotNotInProgress
        );

        // Validate the snapshot has captured the entire registry
        require!(
            next_snapshot.node_count == self.node_count,
            ClockworkError::SnapshotIncomplete
        );

        // Archive the current snapshot
        match current_snapshot {
            Some(current_snapshot) => {
                // Validate the snapshot is current
                require!(
                    current_snapshot.status == SnapshotStatus::Current,
                    ClockworkError::SnapshotNotCurrent
                );

                // Mark the current snapshot as archived
                current_snapshot.status = SnapshotStatus::Archived;
            }
            None => require!(self.snapshot_count == 0, ClockworkError::SnapshotNotCurrent),
        }

        // Mark the next snapshot as current
        next_snapshot.status = SnapshotStatus::Current;

        // Increment snapshot counter
        self.snapshot_count = self.snapshot_count.checked_add(1).unwrap();

        // Unlock the registry
        self.unlock()?;

        Ok(())
    }

    fn lock(&mut self) -> Result<()> {
        self.is_locked = true;
        Ok(())
    }

    fn unlock(&mut self) -> Result<()> {
        self.is_locked = false;
        Ok(())
    }
}