kabu_defi_uniswap_v3_math/
tick_bitmap.rs

1use crate::tick_provider::TickProvider;
2use crate::U256_1;
3use crate::{bit_math, error::UniswapV3MathError};
4use alloy::primitives::{Address, BlockNumber, U256};
5use alloy::providers::Provider;
6use alloy::sol;
7use std::{collections::HashMap, sync::Arc};
8
9sol! {
10    #[sol(rpc)]
11    interface IUniswapV3Pool {
12        function tick_bitmap(int16) external returns (uint256);
13    }
14}
15
16//Flips the initialized state for a given tick from false to true, or vice versa
17pub fn flip_tick(tick_bitmap: &mut HashMap<i16, U256>, tick: i32, tick_spacing: i32) -> Result<(), UniswapV3MathError> {
18    if (tick % tick_spacing) != 0 {
19        return Err(UniswapV3MathError::TickSpacingError);
20    }
21
22    let (word_pos, bit_pos) = position(tick / tick_spacing);
23    let mask = U256_1 << bit_pos;
24    let word = *tick_bitmap.get(&word_pos).unwrap_or(&U256::ZERO);
25    tick_bitmap.insert(word_pos, word ^ mask);
26    Ok(())
27}
28
29//Returns the next initialized tick contained in the same word (or adjacent word) as the tick that is either
30//to the left (less than or equal to) or right (greater than) of the given tick
31pub fn next_initialized_tick_within_one_word<T: TickProvider>(
32    tick_provider: &T,
33    tick: i32,
34    tick_spacing: i32,
35    lte: bool,
36) -> Result<(i32, bool), UniswapV3MathError> {
37    let compressed = if tick < 0 && tick % tick_spacing != 0 { (tick / tick_spacing) - 1 } else { tick / tick_spacing };
38
39    if lte {
40        let (word_pos, bit_pos) = position(compressed);
41
42        let mask = (U256_1 << bit_pos) - U256_1 + (U256_1 << bit_pos);
43
44        let masked = tick_provider.get_tick(word_pos).unwrap_or(U256::ZERO) & mask;
45
46        let initialized = !masked.is_zero();
47
48        let next = if initialized {
49            (compressed - (bit_pos.overflowing_sub(bit_math::most_significant_bit(masked)?).0) as i32) * tick_spacing
50        } else {
51            (compressed - bit_pos as i32) * tick_spacing
52        };
53
54        Ok((next, initialized))
55    } else {
56        let (word_pos, bit_pos) = position(compressed + 1);
57
58        let mask = !((U256_1 << bit_pos) - U256_1);
59
60        let masked = tick_provider.get_tick(word_pos).unwrap_or(U256::ZERO) & mask;
61
62        let initialized = !masked.is_zero();
63
64        let next = if initialized {
65            (compressed + 1 + (bit_math::least_significant_bit(masked)?.overflowing_sub(bit_pos).0) as i32) * tick_spacing
66        } else {
67            (compressed + 1 + ((0xFF - bit_pos) as i32)) * tick_spacing
68        };
69
70        Ok((next, initialized))
71    }
72}
73
74//Returns next and initialized. This function calls the node to get the word at the word_pos.
75//current_word is the current word in the TickBitmap of the pool based on `tick`. TickBitmap[word_pos] = current_word
76//Where word_pos is the 256 bit offset of the ticks word_pos.. word_pos := tick >> 8
77pub async fn next_initialized_tick_within_one_word_from_provider<P: Provider>(
78    tick: i32,
79    tick_spacing: i32,
80    lte: bool,
81    pool_address: Address,
82    block_number: Option<BlockNumber>,
83    provider: Arc<P>,
84) -> Result<(i32, bool), UniswapV3MathError> {
85    let compressed = if tick < 0 && tick % tick_spacing != 0 { (tick / tick_spacing) - 1 } else { tick / tick_spacing };
86
87    if lte {
88        let (word_pos, bit_pos) = position(compressed);
89        let mask = (U256_1 << bit_pos) - U256_1 + (U256_1 << bit_pos);
90
91        let word = if let Some(block_number) = block_number {
92            match IUniswapV3Pool::new(pool_address, provider).tick_bitmap(word_pos).block(block_number.into()).call().await {
93                Ok(word) => word,
94                Err(err) => return Err(UniswapV3MathError::MiddlewareError(err.to_string())),
95            }
96        } else {
97            match IUniswapV3Pool::new(pool_address, provider).tick_bitmap(word_pos).call().await {
98                Ok(word) => word,
99                Err(err) => return Err(UniswapV3MathError::MiddlewareError(err.to_string())),
100            }
101        };
102
103        let masked = word & mask;
104
105        let initialized = !masked.is_zero();
106
107        let next = if initialized {
108            (compressed - (bit_pos.overflowing_sub(bit_math::most_significant_bit(masked)?).0) as i32) * tick_spacing
109        } else {
110            (compressed - bit_pos as i32) * tick_spacing
111        };
112
113        Ok((next, initialized))
114    } else {
115        let (word_pos, bit_pos) = position(compressed + 1);
116        let mask = !((U256_1 << bit_pos) - U256_1);
117
118        let word = if let Some(block_number) = block_number {
119            match IUniswapV3Pool::new(pool_address, provider).tick_bitmap(word_pos).block(block_number.into()).call().await {
120                Ok(word) => word,
121                Err(err) => return Err(UniswapV3MathError::MiddlewareError(err.to_string())),
122            }
123        } else {
124            match IUniswapV3Pool::new(pool_address, provider).tick_bitmap(word_pos).call().await {
125                Ok(word) => word,
126                Err(err) => return Err(UniswapV3MathError::MiddlewareError(err.to_string())),
127            }
128        };
129
130        let masked = word & mask;
131        let initialized = !masked.is_zero();
132
133        let next = if initialized {
134            (compressed + 1 + (bit_math::least_significant_bit(masked)?.overflowing_sub(bit_pos).0) as i32) * tick_spacing
135        } else {
136            (compressed + 1 + ((0xFF - bit_pos) as i32)) * tick_spacing
137        };
138
139        Ok((next, initialized))
140    }
141}
142
143//Computes the position in the mapping where the initialized bit for a tick lives
144pub fn position(tick: i32) -> (i16, u8) {
145    ((tick >> 8) as i16, (tick % 256) as u8)
146}
147
148#[cfg(test)]
149mod test {
150    use std::{collections::HashMap, vec};
151
152    use super::{flip_tick, next_initialized_tick_within_one_word};
153    use crate::tick_provider::TickProvider;
154    use alloy::primitives::U256;
155
156    pub struct TickProviderHashMap {
157        tick_bitmap: HashMap<i16, U256>,
158    }
159    impl TickProviderHashMap {
160        pub fn new(tick_bitmap: HashMap<i16, U256>) -> Self {
161            Self { tick_bitmap }
162        }
163        pub fn mut_ref(&mut self) -> &mut HashMap<i16, U256> {
164            &mut self.tick_bitmap
165        }
166        pub fn clear(&mut self) {
167            self.tick_bitmap.clear();
168        }
169    }
170    impl TickProvider for TickProviderHashMap {
171        fn get_tick(&self, word_pos: i16) -> eyre::Result<U256> {
172            let val = self.tick_bitmap.get(&word_pos).cloned();
173            Ok(val.unwrap_or(U256::ZERO))
174        }
175    }
176
177    pub fn init_test_ticks() -> eyre::Result<TickProviderHashMap> {
178        let test_ticks = vec![-200, -55, -4, 70, 78, 84, 139, 240, 535];
179        let mut tick_bitmap: HashMap<i16, U256> = HashMap::new();
180        for tick in test_ticks {
181            flip_tick(&mut tick_bitmap, tick, 1)?;
182        }
183        Ok(TickProviderHashMap::new(tick_bitmap))
184    }
185
186    pub fn initialized<T: TickProvider>(tick: i32, tick_provider: &T) -> eyre::Result<bool> {
187        let (next, initialized) = next_initialized_tick_within_one_word(tick_provider, tick, 1, true)?;
188        if next == tick {
189            Ok(initialized)
190        } else {
191            Ok(false)
192        }
193    }
194
195    #[test]
196    pub fn test_next_initialized_tick_within_one_word_lte_false() -> eyre::Result<()> {
197        let mut tick_bitmap = init_test_ticks()?;
198        //returns tick to right if at initialized tick
199        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 78, 1, false)?;
200        assert_eq!(next, 84);
201        assert!(initialized);
202        tick_bitmap = init_test_ticks()?;
203        // //returns the tick directly to the right
204        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 77, 1, false)?;
205
206        assert_eq!(next, 78);
207        assert!(initialized);
208        tick_bitmap = init_test_ticks()?;
209        // //returns the tick directly to the right
210        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, -56, 1, false)?;
211
212        assert_eq!(next, -55);
213        assert!(initialized);
214        tick_bitmap = init_test_ticks()?;
215        //returns the next words initialized tick if on the right boundary
216        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 255, 1, false)?;
217
218        assert_eq!(next, 511);
219        assert!(!initialized);
220        tick_bitmap = init_test_ticks()?;
221        //returns the next words initialized tick if on the right boundary
222        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, -257, 1, false)?;
223
224        assert_eq!(next, -200);
225        assert!(initialized);
226        tick_bitmap = init_test_ticks()?;
227        //returns the next initialized tick from the next word
228        flip_tick(tick_bitmap.mut_ref(), 340, 1)?;
229        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 328, 1, false)?;
230
231        assert_eq!(next, 340);
232        assert!(initialized);
233        tick_bitmap = init_test_ticks()?;
234        //does not exceed boundary
235        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 508, 1, false)?;
236
237        assert_eq!(next, 511);
238        assert!(!initialized);
239        tick_bitmap = init_test_ticks()?;
240        //skips entire word
241        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 255, 1, false)?;
242
243        assert_eq!(next, 511);
244        assert!(!initialized);
245        tick_bitmap = init_test_ticks()?;
246        //skips half word
247        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 383, 1, false)?;
248
249        assert_eq!(next, 511);
250        assert!(!initialized);
251        Ok(())
252    }
253
254    #[test]
255    pub fn test_next_initialized_tick_within_one_word_lte_true() -> eyre::Result<()> {
256        let mut tick_bitmap = init_test_ticks()?;
257        //returns same tick if initialized
258        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 78, 1, true)?;
259        assert_eq!(next, 78);
260        assert!(initialized);
261        tick_bitmap = init_test_ticks()?;
262        //returns tick directly to the left of input tick if not initialized
263        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 79, 1, true)?;
264
265        assert_eq!(next, 78);
266        assert!(initialized);
267        tick_bitmap = init_test_ticks()?;
268        //will not exceed the word boundary
269        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 258, 1, true)?;
270
271        assert_eq!(next, 256);
272        assert!(!initialized);
273        tick_bitmap = init_test_ticks()?;
274        //at the word boundary
275        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 256, 1, true)?;
276
277        assert_eq!(next, 256);
278        assert!(!initialized);
279        tick_bitmap = init_test_ticks()?;
280        //word boundary less 1 (next initialized tick in next word)',
281        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 72, 1, true)?;
282
283        assert_eq!(next, 70);
284        assert!(initialized);
285        tick_bitmap = init_test_ticks()?;
286        //word boundary
287        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, -257, 1, true)?;
288
289        assert_eq!(next, -512);
290        assert!(!initialized);
291        tick_bitmap = init_test_ticks()?;
292        //entire empty word
293        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 1023, 1, true)?;
294
295        assert_eq!(next, 768);
296        assert!(!initialized);
297        tick_bitmap = init_test_ticks()?;
298        //halfway through empty word
299        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 900, 1, true)?;
300
301        assert_eq!(next, 768);
302        assert!(!initialized);
303        tick_bitmap = init_test_ticks()?;
304        //boundary is initialized
305        flip_tick(tick_bitmap.mut_ref(), 329, 1)?;
306        let (next, initialized) = next_initialized_tick_within_one_word(&tick_bitmap, 456, 1, true)?;
307
308        assert_eq!(next, 329);
309        assert!(initialized);
310        Ok(())
311    }
312
313    #[test]
314    pub fn test_initialized() -> eyre::Result<()> {
315        //is false at first
316        let mut tick_bitmap = TickProviderHashMap::new(HashMap::new());
317        let is_initialized = initialized(1, &tick_bitmap)?;
318
319        assert!(!is_initialized);
320        //is flipped by #flipTick
321        flip_tick(tick_bitmap.mut_ref(), 1, 1)?;
322        let is_initialized: bool = initialized(1, &tick_bitmap)?;
323        assert!(is_initialized);
324
325        //is flipped back by #flipTick
326        tick_bitmap.clear();
327        flip_tick(tick_bitmap.mut_ref(), 1, 1)?;
328        flip_tick(tick_bitmap.mut_ref(), 1, 1)?;
329        let is_initialized = initialized(1, &tick_bitmap)?;
330        assert!(!is_initialized);
331
332        //is not changed by another flip to a different tick
333        tick_bitmap.clear();
334        flip_tick(tick_bitmap.mut_ref(), 2, 1)?;
335        let is_initialized = initialized(1, &tick_bitmap)?;
336        assert!(!is_initialized);
337
338        //is not changed by another flip to a different tick on another word
339        tick_bitmap.clear();
340        flip_tick(tick_bitmap.mut_ref(), 1 + 256, 1)?;
341        let is_initialized = initialized(257, &tick_bitmap)?;
342        assert!(is_initialized);
343        let is_initialized = initialized(1, &tick_bitmap)?;
344        assert!(!is_initialized);
345        Ok(())
346    }
347
348    #[test]
349    pub fn test_flip_tick() -> eyre::Result<()> {
350        //flips only the specified tick
351        let mut tick_bitmap = TickProviderHashMap::new(HashMap::new());
352        flip_tick(tick_bitmap.mut_ref(), -230, 1)?;
353        let is_initialized = initialized(-230, &tick_bitmap)?;
354        assert!(is_initialized);
355        let is_initialized = initialized(-231, &tick_bitmap)?;
356        assert!(!is_initialized);
357        let is_initialized = initialized(-229, &tick_bitmap)?;
358        assert!(!is_initialized);
359        let is_initialized = initialized(-230 + 256, &tick_bitmap)?;
360        assert!(!is_initialized);
361        let is_initialized = initialized(-230 - 256, &tick_bitmap)?;
362        assert!(!is_initialized);
363        flip_tick(tick_bitmap.mut_ref(), -230, 1)?;
364        let is_initialized = initialized(-230, &tick_bitmap)?;
365        assert!(!is_initialized);
366        let is_initialized = initialized(-231, &tick_bitmap)?;
367        assert!(!is_initialized);
368        let is_initialized = initialized(-229, &tick_bitmap)?;
369        assert!(!is_initialized);
370        let is_initialized = initialized(-230 + 256, &tick_bitmap)?;
371        assert!(!is_initialized);
372        let is_initialized = initialized(-230 - 256, &tick_bitmap)?;
373        assert!(!is_initialized);
374        //reverts only itself
375        tick_bitmap.clear();
376        flip_tick(tick_bitmap.mut_ref(), -230, 1)?;
377        flip_tick(tick_bitmap.mut_ref(), -259, 1)?;
378        flip_tick(tick_bitmap.mut_ref(), -229, 1)?;
379        flip_tick(tick_bitmap.mut_ref(), 500, 1)?;
380        flip_tick(tick_bitmap.mut_ref(), -259, 1)?;
381        flip_tick(tick_bitmap.mut_ref(), -229, 1)?;
382        flip_tick(tick_bitmap.mut_ref(), -259, 1)?;
383        let is_initialized = initialized(-259, &tick_bitmap)?;
384        assert!(is_initialized);
385        let is_initialized = initialized(-229, &tick_bitmap)?;
386        assert!(!is_initialized);
387
388        Ok(())
389    }
390}