tempo_precompiles/tip20_rewards_registry/
mod.rs

1// Module for tip20_rewards_registry precompile
2pub mod dispatch;
3
4use crate::{
5    TIP20_REWARDS_REGISTRY_ADDRESS,
6    error::{Result, TempoPrecompileError},
7    storage::{Handler, Mapping},
8    tip20::{TIP20Token, address_to_token_id_unchecked},
9};
10use alloy::{
11    primitives::{Address, B256, U256, keccak256},
12    sol_types::SolValue,
13};
14
15pub use tempo_contracts::precompiles::{ITIP20RewardsRegistry, TIP20RewardsRegistryError};
16use tempo_precompiles_macros::contract;
17
18/// TIPRewardsRegistry precompile that tracks stream end times
19/// Maps timestamp -> Vec of token addresses with streams ending at that time
20#[contract(addr = TIP20_REWARDS_REGISTRY_ADDRESS)]
21pub struct TIP20RewardsRegistry {
22    last_updated_timestamp: u128,
23    ending_streams: Mapping<u128, Vec<Address>>,
24    stream_index: Mapping<B256, U256>,
25}
26
27impl TIP20RewardsRegistry {
28    /// Initializes the TIP20 rewards registry contract.
29    ///
30    /// Ensures the [`TIP20RewardsRegistry`] account isn't empty and prevents state clear.
31    pub fn initialize(&mut self) -> Result<()> {
32        self.__initialize()
33    }
34
35    /// Add a token to the registry for a given stream end time
36    pub fn add_stream(&mut self, token: Address, end_time: u128) -> Result<()> {
37        let stream_key = keccak256((token, end_time).abi_encode());
38        let stream_ending_at = self.ending_streams.at(end_time);
39        let length = stream_ending_at.len()?;
40
41        self.stream_index.at(stream_key).write(U256::from(length))?;
42        stream_ending_at.push(token)
43    }
44
45    /// Remove stream before it is finalized
46    pub fn remove_stream(&mut self, token: Address, end_time: u128) -> Result<()> {
47        let stream_key = keccak256((token, end_time).abi_encode());
48        let index: usize = self.stream_index.at(stream_key).read()?.to();
49
50        let stream_ending_at = self.ending_streams.at(end_time);
51        let length = stream_ending_at.len()?;
52        let last_index = length
53            .checked_sub(1)
54            .ok_or(TempoPrecompileError::under_overflow())?;
55
56        // If removing element that's not the last, swap with last element
57        if index != last_index {
58            let last_token = stream_ending_at.at_unchecked(last_index).read()?;
59            stream_ending_at.at_unchecked(index).write(last_token)?;
60
61            // Update stream_index for the moved element
62            let last_stream_key = keccak256((last_token, end_time).abi_encode());
63            self.stream_index
64                .at(last_stream_key)
65                .write(U256::from(index))?;
66        }
67
68        // Remove last element and clear its index
69        stream_ending_at.pop()?;
70        self.stream_index.at(stream_key).delete()?;
71
72        Ok(())
73    }
74
75    /// Finalize streams for all tokens ending at the current timestamp
76    pub fn finalize_streams(&mut self, sender: Address) -> Result<()> {
77        if sender != Address::ZERO {
78            return Err(TIP20RewardsRegistryError::unauthorized().into());
79        }
80
81        let current_timestamp = self.storage.timestamp().to::<u128>();
82        let mut last_updated = self.last_updated_timestamp.read()?;
83
84        if last_updated == 0 {
85            last_updated = current_timestamp.saturating_sub(1);
86        }
87
88        if current_timestamp == last_updated {
89            return Ok(());
90        }
91
92        let mut next_timestamp = last_updated
93            .checked_add(1)
94            .ok_or(TempoPrecompileError::under_overflow())?;
95
96        while current_timestamp >= next_timestamp {
97            let tokens = self.ending_streams.at(next_timestamp).read()?;
98
99            for token in tokens {
100                let token_id = address_to_token_id_unchecked(token);
101                let mut tip20_token = TIP20Token::new(token_id);
102                tip20_token.finalize_streams(self.address, next_timestamp)?;
103
104                let stream_key = keccak256((token, next_timestamp).abi_encode());
105                self.stream_index.at(stream_key).delete()?;
106            }
107
108            // Clear all elements from the vec
109            self.ending_streams.at(next_timestamp).delete()?;
110
111            next_timestamp = next_timestamp
112                .checked_add(1)
113                .ok_or(TempoPrecompileError::under_overflow())?;
114        }
115
116        self.last_updated_timestamp.write(current_timestamp)?;
117
118        Ok(())
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use crate::{
126        error::TempoPrecompileError,
127        storage::{ContractStorage, StorageCtx, hashmap::HashMapStorageProvider},
128        test_util::TIP20Setup,
129        tip20_rewards_registry::TIP20RewardsRegistry,
130    };
131    use alloy::primitives::Address;
132    use tempo_chainspec::hardfork::TempoHardfork;
133    use tempo_contracts::precompiles::ITIP20;
134
135    #[test]
136    fn test_add_stream() -> eyre::Result<()> {
137        let mut storage = HashMapStorageProvider::new(1);
138        let token = Address::random();
139        let token2 = Address::random();
140        StorageCtx::enter(&mut storage, || {
141            StorageCtx.set_timestamp(U256::from(1000));
142
143            let mut registry = TIP20RewardsRegistry::new();
144            registry.initialize()?;
145
146            let end_time = 2000u128;
147
148            registry.add_stream(token, end_time)?;
149
150            let streams = registry.ending_streams.at(end_time).read()?;
151            assert_eq!(streams.len(), 1);
152            assert_eq!(streams[0], token);
153
154            let stream_key = keccak256((token, end_time).abi_encode());
155            let index = registry.stream_index.at(stream_key).read()?;
156            assert_eq!(index, U256::ZERO);
157
158            registry.add_stream(token2, end_time)?;
159
160            let streams = registry.ending_streams.at(end_time).read()?;
161            assert_eq!(streams.len(), 2);
162            assert!(streams.contains(&token));
163            assert!(streams.contains(&token2));
164
165            let stream_key2 = keccak256((token2, end_time).abi_encode());
166            let index2 = registry.stream_index.at(stream_key2).read()?;
167            assert_eq!(index2, U256::ONE);
168
169            Ok(())
170        })
171    }
172
173    #[test]
174    fn test_remove_stream() -> eyre::Result<()> {
175        let mut storage = HashMapStorageProvider::new(1);
176        let token1 = Address::random();
177        let token2 = Address::random();
178        let token3 = Address::random();
179        let non_existent_token = Address::random();
180        StorageCtx::enter(&mut storage, || {
181            StorageCtx.set_timestamp(U256::from(1000));
182
183            let mut registry = TIP20RewardsRegistry::new();
184            registry.initialize()?;
185
186            let end_time = 2000u128;
187
188            // Add three streams
189            registry.add_stream(token1, end_time)?;
190            registry.add_stream(token2, end_time)?;
191            registry.add_stream(token3, end_time)?;
192
193            let streams = registry.ending_streams.at(end_time).read()?;
194            assert_eq!(streams.len(), 3);
195            assert_eq!(streams[0], token1);
196            assert_eq!(streams[1], token2);
197            assert_eq!(streams[2], token3);
198
199            registry.remove_stream(token2, end_time)?;
200
201            let streams = registry.ending_streams.at(end_time).read()?;
202            assert_eq!(streams.len(), 2);
203            assert_eq!(streams[0], token1);
204            assert_eq!(streams[1], token3);
205
206            // Verify indices are updated correctly
207            let stream_key1 = keccak256((token1, end_time).abi_encode());
208            let stream_key2 = keccak256((token2, end_time).abi_encode());
209            let stream_key3 = keccak256((token3, end_time).abi_encode());
210
211            let index1 = registry.stream_index.at(stream_key1).read()?;
212            let index2 = registry.stream_index.at(stream_key2).read()?;
213            let index3 = registry.stream_index.at(stream_key3).read()?;
214
215            assert_eq!(index1, U256::ZERO);
216            assert_eq!(index2, U256::ZERO);
217            assert_eq!(index3, U256::ONE);
218
219            registry.remove_stream(token3, end_time)?;
220            let streams = registry.ending_streams.at(end_time).read()?;
221            assert_eq!(streams.len(), 1);
222            assert_eq!(streams[0], token1);
223
224            registry.remove_stream(token1, end_time)?;
225
226            let streams = registry.ending_streams.at(end_time).read()?;
227            assert_eq!(streams.len(), 0);
228
229            // Test removing non-existent stream
230            let result = registry.remove_stream(non_existent_token, end_time);
231            assert!(result.is_err());
232
233            Ok(())
234        })
235    }
236
237    #[test]
238    fn test_ending_streams() -> eyre::Result<()> {
239        let mut storage = HashMapStorageProvider::new(1);
240        let token1 = Address::random();
241        let token2 = Address::random();
242        let token3 = Address::random();
243        let token4 = Address::random();
244        StorageCtx::enter(&mut storage, || {
245            StorageCtx.set_timestamp(U256::from(1000));
246
247            let mut registry = TIP20RewardsRegistry::new();
248            registry.initialize()?;
249
250            let timestamp = 2000u128;
251
252            let empty_streams = registry.ending_streams.at(timestamp).read()?;
253            assert_eq!(empty_streams.len(), 0);
254
255            registry.add_stream(token1, timestamp)?;
256            registry.add_stream(token2, timestamp)?;
257            registry.add_stream(token3, timestamp)?;
258
259            let streams = registry.ending_streams.at(timestamp).read()?;
260            assert_eq!(streams.len(), 3);
261            assert_eq!(streams[0], token1);
262            assert_eq!(streams[1], token2);
263            assert_eq!(streams[2], token3);
264
265            let other_timestamp = 3000u128;
266            let other_streams = registry.ending_streams.at(other_timestamp).read()?;
267            assert_eq!(other_streams.len(), 0);
268
269            registry.add_stream(token4, other_timestamp)?;
270
271            let streams1 = registry.ending_streams.at(timestamp).read()?;
272            let streams2 = registry.ending_streams.at(other_timestamp).read()?;
273
274            assert_eq!(streams1.len(), 3);
275            assert_eq!(streams2.len(), 1);
276            assert_eq!(streams2[0], token4);
277
278            Ok(())
279        })
280    }
281
282    #[test]
283    fn test_finalize_streams() -> eyre::Result<()> {
284        let mut storage = HashMapStorageProvider::new(1).with_spec(TempoHardfork::Adagio);
285        let admin = Address::random();
286        let unauthorized = Address::random();
287        StorageCtx::enter(&mut storage, || {
288            StorageCtx.set_timestamp(U256::from(1500));
289
290            let mut registry = TIP20RewardsRegistry::new();
291            registry.initialize()?;
292
293            let mut token = TIP20Setup::create("Test", "TST", admin)
294                .with_issuer(admin)
295                .with_mint(admin, U256::from(100e18 as u128))
296                .apply()?;
297            let token_addr = token.address();
298
299            // Start a reward stream that lasts 5 seconds from current time (1500)
300            let current_time = token.storage().timestamp().to::<u128>();
301            let stream_duration = 5;
302            let stream_id = token.start_reward(
303                admin,
304                ITIP20::startRewardCall {
305                    amount: U256::from(100e18 as u128),
306                    secs: stream_duration,
307                },
308            )?;
309            assert_eq!(stream_id, 1);
310
311            // Test unauthorized caller
312            let result = registry.finalize_streams(unauthorized);
313            assert!(matches!(
314                result.unwrap_err(),
315                TempoPrecompileError::TIP20RewardsRegistry(
316                    TIP20RewardsRegistryError::Unauthorized(_)
317                )
318            ));
319
320            let result = registry.finalize_streams(Address::ZERO);
321            assert!(result.is_ok());
322
323            // Verify the stream was added to registry at the correct end time
324            let end_time = current_time + stream_duration as u128;
325            let streams_before = registry.ending_streams.at(end_time).read()?;
326            assert_eq!(streams_before.len(), 1);
327            assert_eq!(streams_before[0], token_addr);
328
329            // Fast forward to the end time to simulate stream completion
330            registry.storage.set_timestamp(U256::from(end_time));
331            registry.finalize_streams(Address::ZERO)?;
332
333            let last_updated = registry.last_updated_timestamp.read()?;
334            assert_eq!(last_updated, end_time);
335
336            // Verify streams were cleared from the registry
337            let streams_after = registry.ending_streams.at(end_time).read()?;
338            assert_eq!(streams_after.len(), 0);
339
340            let result = registry.finalize_streams(Address::ZERO);
341            assert!(result.is_ok());
342
343            Ok(())
344        })
345    }
346}
tempo_precompiles/tip20_rewards_registry/mod.rs

tempo_precompiles/tip20_rewards_registry/
mod.rs
