kabu_types_entities/
pool_id.rs

1use alloy_primitives::{Address, B256};
2use serde::{Deserialize, Serialize};
3
4#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize, strum_macros::Display)]
5#[serde(untagged)]
6pub enum PoolId {
7    #[strum(to_string = "{0}")]
8    Address(Address),
9    #[strum(to_string = "{0}")]
10    B256(B256),
11}
12
13impl From<Address> for PoolId {
14    fn from(address: Address) -> Self {
15        PoolId::Address(address)
16    }
17}
18
19impl From<B256> for PoolId {
20    fn from(b256: B256) -> Self {
21        PoolId::B256(b256)
22    }
23}
24
25impl PoolId {
26    pub fn address_or_zero(&self) -> Address {
27        match self {
28            PoolId::Address(addr) => *addr,
29            PoolId::B256(_) => Address::ZERO,
30        }
31    }
32}
33
34#[cfg(test)]
35mod tests {
36    use super::*;
37    use alloy_primitives::Address;
38
39    #[test]
40    fn test_deserialize_address() {
41        let deserialized: PoolId = serde_json::from_str("\"0x0101010101010101010101010101010101010101\"").unwrap();
42        assert_eq!(PoolId::Address(Address::repeat_byte(0x01)), deserialized);
43    }
44
45    #[test]
46    fn test_serialize_address() {
47        let addr = Address::repeat_byte(0x01);
48        let pool_id = PoolId::Address(addr);
49        let serialized = serde_json::to_string(&pool_id).unwrap();
50        assert_eq!(serialized, "\"0x0101010101010101010101010101010101010101\"");
51    }
52
53    #[test]
54    fn test_deserialize_b256() {
55        let deserialized: PoolId = serde_json::from_str("\"0x0101010101010101010101010101010101010101010101010101010101010101\"").unwrap();
56        assert_eq!(PoolId::B256(B256::repeat_byte(0x01)), deserialized);
57    }
58
59    #[test]
60    fn test_serialize_b256() {
61        let b256 = B256::repeat_byte(0x01);
62        let pool_id = PoolId::B256(b256);
63        let serialized = serde_json::to_string(&pool_id).unwrap();
64        assert_eq!(serialized, "\"0x0101010101010101010101010101010101010101010101010101010101010101\"");
65    }
66
67    #[test]
68    fn test_display() {
69        let addr = Address::repeat_byte(0x01);
70        let pool_id = PoolId::Address(addr);
71        assert_eq!(format!("{pool_id}"), "0x0101010101010101010101010101010101010101");
72
73        let b256 = B256::repeat_byte(0x01);
74        let pool_id = PoolId::B256(b256);
75        assert_eq!(format!("{pool_id}"), "0x0101010101010101010101010101010101010101010101010101010101010101");
76    }
77
78    #[test]
79    fn test_ord() {
80        let addr_small = Address::repeat_byte(0x01);
81        let addr_large = Address::repeat_byte(0x02);
82        let b256_small = B256::repeat_byte(0x01);
83        let b256_large = B256::repeat_byte(0x02);
84
85        let pool_id_addr_small = PoolId::Address(addr_small);
86        let pool_id_addr_large = PoolId::Address(addr_large);
87        let pool_id_byte32_small = PoolId::B256(b256_small);
88        let pool_id_byte32_large = PoolId::B256(b256_large);
89
90        // Address small < large
91        assert!(pool_id_addr_small < pool_id_addr_large);
92        assert!(pool_id_addr_large > pool_id_addr_small);
93
94        // B256 small < large
95        assert!(pool_id_byte32_small < pool_id_byte32_large);
96        assert!(pool_id_byte32_large > pool_id_byte32_small);
97
98        // Address < B256
99        assert!(pool_id_addr_small < pool_id_byte32_large);
100        assert!(pool_id_byte32_large > pool_id_addr_small);
101    }
102}
kabu_types_entities/pool_id.rs

kabu_types_entities/
pool_id.rs
