polyphony_types/utils/
snowflake.rs

1use std::fmt::Display;
2
3use atomic::Atomic;
4use bigdecimal::{Num, ToPrimitive, Zero};
5use num_bigint::{BigInt, ToBigInt};
6use serde::{Deserialize, Serialize};
7#[cfg(feature = "sqlx")]
8use sqlx::Type;
9
10const EPOCH: i64 = 1420070400000;
11static WORKER_ID: u128 = 0;
12static PROCESS_ID: u128 = 1;
13lazy_static::lazy_static! {
14    static ref INCREMENT: Atomic<u128> = Atomic::default();
15}
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
17#[cfg_attr(feature = "sqlx", derive(Type))]
18#[cfg_attr(feature = "sqlx", sqlx(transparent))]
19pub struct Snowflake(String);
20
21impl Default for Snowflake {
22    fn default() -> Self {
23        Self::generate()
24    }
25}
26
27impl Display for Snowflake {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        write!(f, "{}", self.0)
30    }
31}
32
33impl Snowflake {
34    pub fn to_binary(&self) -> String {
35        let self_len = self.0.len();
36        let high = self.0[..self_len - 10].parse::<u64>().unwrap_or(0);
37        let low = self.0[self_len - 10..].parse::<u64>().unwrap();
38        let mut low = low;
39        let mut high = high;
40        let mut bin = Vec::with_capacity(64);
41
42        while low > 0 || high > 0 {
43            bin.push((low & 1) as u8);
44            low >>= 1;
45
46            if high > 0 {
47                low += 5_000_000_000 * (high % 2);
48                high >>= 1;
49            }
50        }
51
52        bin.iter()
53            .rev()
54            .map(|b| char::from_digit(*b as u32, 10).unwrap())
55            .collect()
56    }
57
58    pub fn from_binary(num: &str) -> String {
59        let mut num = BigInt::from_str_radix(num, 2).unwrap();
60        let mut dec = Vec::with_capacity(18);
61
62        let ten = 10.to_bigint().unwrap();
63        let _two = 2.to_bigint().unwrap();
64        let _thirty_two = 32.to_bigint().unwrap();
65
66        while num.bits() > 50 {
67            let high: BigInt = &num >> 32;
68            let low: BigInt = (high.clone() % &ten) << 32 | &num & BigInt::from((1u64 << 32) - 1);
69
70            let next: BigInt = low.clone() % &ten;
71            dec.push(next.to_u8().unwrap());
72            num = (high / &ten) << 32 | (low / &ten);
73        }
74
75        while !num.is_zero() {
76            dec.push((num.clone() % &ten).to_u8().unwrap());
77            num /= &ten;
78        }
79
80        dec.iter()
81            .rev()
82            .map(|d| char::from_digit(*d as u32, 10).unwrap())
83            .collect()
84    }
85
86    pub fn generate_worker_process() -> u128 {
87        let time = (chrono::Utc::now().naive_utc().timestamp_millis() - EPOCH) << 22;
88        let worker = WORKER_ID << 17;
89        let process = PROCESS_ID << 12;
90        let increment = INCREMENT.load(atomic::Ordering::Relaxed);
91
92        INCREMENT.store(increment + 1, atomic::Ordering::Relaxed);
93
94        time as u128 | worker | process | increment
95    }
96
97    pub fn generate() -> Self {
98        Self(Self::generate_worker_process().to_string())
99    }
100
101    pub fn deconstruct(&self) -> DeconstructedSnowflake {
102        let binary = format!("{:0>64}", self.to_binary());
103
104        let ts = i64::from_str_radix(&binary[0..42], 2).unwrap() + EPOCH;
105        let wid = u64::from_str_radix(&binary[42..47], 2).unwrap();
106        let pid = u64::from_str_radix(&binary[47..52], 2).unwrap();
107        let increment = BigInt::from_str_radix(&binary[52..64], 2).unwrap();
108
109        DeconstructedSnowflake {
110            timestamp: ts,
111            worker_id: wid,
112            process_id: pid,
113            increment,
114            binary,
115        }
116    }
117}
118
119#[derive(Debug, Clone, PartialEq, Eq)]
120pub struct DeconstructedSnowflake {
121    pub timestamp: i64,
122    pub worker_id: u64,
123    pub process_id: u64,
124    pub increment: BigInt,
125    pub binary: String,
126}
127
128#[cfg(test)]
129mod test {
130    use super::Snowflake;
131
132    #[test]
133    fn test_new_snowflake() {
134        let snow = Snowflake::generate();
135        println!("{snow}");
136    }
137
138    #[test]
139    fn snowflake_to_binary() {
140        let snowflake = super::Snowflake("1104339392517902336".to_string());
141
142        let bin = snowflake.to_binary();
143        println!("{bin}");
144    }
145
146    #[test]
147    fn binary_to_snowflake() {
148        let snowflake = super::Snowflake::from_binary(
149            "111101010011011001101101001110010010100000000001000000000000",
150        );
151        println!("{snowflake}");
152    }
153
154    #[test]
155    fn test_deconstruct() {
156        let new = super::Snowflake::generate();
157
158        println!("{:?}", new.deconstruct());
159    }
160}