yoink snowflake out of clearcut

5 files changed, 227 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..4fffb2f
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+/target
+/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..6421915
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,17 @@
+[package]
+name = "snowflake"
+authors = ["Galen Guyer <galen@galenguyer.com>"]
+license = "MIT"
+version = "1.0.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dev-dependencies]
+criterion = "0.4.0"
+
+[[bench]]
+name = "snowflake"
+harness = false
+
+[dependencies]
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..60bc4c4
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,19 @@
+Copyright (c) 2023 Galen Guyer <galen@galenguyer.com>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
\ No newline at end of file
diff --git a/benches/snowflake.rs b/benches/snowflake.rs
new file mode 100644
index 0000000..be6fc9f
--- /dev/null
+++ b/benches/snowflake.rs
@@ -0,0 +1,18 @@
+use criterion::{criterion_group, criterion_main, Criterion, Throughput};
+
+fn bench_generate(c: &mut Criterion) {
+    let mut group = c.benchmark_group("snowflake");
+    group.throughput(Throughput::Elements(1)).sample_size(1000);
+    group.bench_function("generate", |b| {
+        let mut generator = snowflake::SnowflakeGenerator::new(0, 0);
+        b.iter(|| generator.generate())
+    });
+    group.bench_function("generate_fuzzy", |b| {
+        let mut generator = snowflake::SnowflakeGenerator::new(0, 1);
+        b.iter(|| generator.generate_fuzzy())
+    });
+    group.finish();
+}
+
+criterion_group!(benches, bench_generate);
+criterion_main!(benches);
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..6a234af
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,171 @@
+use std::time::{SystemTime, UNIX_EPOCH};
+
+pub struct SnowflakeGenerator {
+    epoch: SystemTime,
+    last_time: u64,
+    machine_id: u8,
+    thread_id: u8,
+    counter: u16,
+}
+
+#[derive(Debug)]
+pub struct Snowflake {
+    /// The time in milliseconds since the epoch.
+    /// This field does not automatically compensate if an epoc other than UNIX_EPOCH is used.
+    pub time: u64,
+    /// The machine ID the snowflake was generated on.
+    pub machine_id: u8,
+    /// The thread ID the snowflake was generated on.
+    pub thread_id: u8,
+    /// The counter for the snowflake. This is incremented every time a snowflake is generated
+    /// and reset if the time has changed
+    pub counter: u16,
+}
+
+impl SnowflakeGenerator {
+    /// Creates a new SnowflakeGenerator with the given machine ID and thread ID.
+    /// The machine ID must be less than 32 and the thread ID must be less than 32.
+    ///
+    /// # Examples
+    /// ```
+    /// # use snowflake::SnowflakeGenerator;
+    /// let mut generator = SnowflakeGenerator::new(0, 0);
+    /// ```
+    /// # Panics
+    /// This function will panic if the machine ID or thread ID is greater than 31.
+    ///
+    /// ```should_panic
+    /// # use snowflake::SnowflakeGenerator;
+    /// let mut generator = SnowflakeGenerator::new(32, 32);
+    /// ```
+    pub fn new(machine_id: u8, thread_id: u8) -> Self {
+        SnowflakeGenerator::with_epoch(UNIX_EPOCH, machine_id, thread_id)
+    }
+
+    /// Creates a new SnowflakeGenerator with the given epoch, machine ID, and thread ID.
+    /// The machine ID must be less than 32 and the thread ID must be less than 32.
+    /// The epoch is the time that the SnowflakeGenerator will use as the start of time.
+    /// This is useful if you want to use a different epoch than the Unix epoch.
+    ///
+    /// # Examples
+    /// ```
+    /// # use snowflake::SnowflakeGenerator;
+    /// # use std::time::UNIX_EPOCH;
+    /// let mut generator = SnowflakeGenerator::with_epoch(UNIX_EPOCH, 0, 0);
+    /// ```
+    ///
+    /// # Panics
+    /// This function will panic if the machine ID or thread ID is greater than 31.
+    ///
+    /// ```should_panic
+    /// # use snowflake::SnowflakeGenerator;
+    /// # use std::time::UNIX_EPOCH;
+    /// let mut generator = SnowflakeGenerator::with_epoch(UNIX_EPOCH, 32, 32);
+    /// ```
+    pub fn with_epoch(epoch: SystemTime, machine_id: u8, thread_id: u8) -> Self {
+        assert!(machine_id < 32, "machine_id must be less than 32");
+        assert!(thread_id < 32, "thread_id must be less than 32");
+        SnowflakeGenerator {
+            epoch,
+            last_time: get_time_millis(epoch),
+            machine_id,
+            thread_id,
+            counter: 0,
+        }
+    }
+
+    /// Generates a new Snowflake ID.
+    /// This function will block until it can generate a new ID.
+    ///
+    /// # Examples
+    /// ```
+    /// # use snowflake::SnowflakeGenerator;
+    /// let mut generator = SnowflakeGenerator::new(0, 0);
+    /// let id = generator.generate();
+    /// ```
+    pub fn generate(&mut self) -> u64 {
+        let mut now = get_time_millis(self.epoch);
+
+        // If the time is the same as the last time we generated an ID, we need to increment our counter
+        if now == self.last_time {
+            self.counter = (self.counter + 1) % 4096;
+            if self.counter == 0 {
+                // If we've reached the maximum number of IDs we can generate in a single millisecond,
+                // we need to wait until the next millisecond
+                while now <= self.last_time {
+                    now = get_time_millis(self.epoch);
+                }
+            }
+        } else {
+            // This is a new millisecond so we reset our counter
+            self.counter = 0;
+        }
+
+        self.last_time = now;
+
+        self.last_time << 22
+            | ((self.machine_id as u64) << 17)
+            | ((self.thread_id as u64) << 12)
+            | (self.counter as u64)
+    }
+
+    /// Generates a new Snowflake ID.
+    /// This function will not block and will increment the timestamp if the counter is full.
+    ///
+    /// # Examples
+    /// ```
+    /// # use snowflake::SnowflakeGenerator;
+    /// let mut generator = SnowflakeGenerator::new(0, 0);
+    /// let id = generator.generate_fuzzy();
+    /// ```
+    pub fn generate_fuzzy(&mut self) -> u64 {
+        let mut now = get_time_millis(self.epoch);
+
+        // If the actual time is less than or the same as the last time we generated an ID,
+        // we need to increment our counter
+        if now <= self.last_time {
+            self.counter = (self.counter + 1) % 4096;
+            if self.counter == 0 {
+                // If we've reached the maximum number of IDs we can generate in a single millisecond,
+                // we need to increment the current millisecond
+                now += 1;
+            }
+        } else {
+            // This is a new millisecond so we reset our counter
+            self.counter = 0;
+        }
+
+        self.last_time = now;
+
+        self.last_time << 22
+            | ((self.machine_id as u64) << 17)
+            | ((self.thread_id as u64) << 12)
+            | (self.counter as u64)
+    }
+}
+
+impl From<u64> for Snowflake {
+    fn from(value: u64) -> Self {
+        Snowflake {
+            time: value >> 22,
+            machine_id: ((value & 0x3E0000) >> 17) as u8,
+            thread_id: ((value & 0x1F000) >> 12) as u8,
+            counter: (value & 0xFFF) as u16,
+        }
+    }
+}
+impl From<Snowflake> for u64 {
+    fn from(value: Snowflake) -> Self {
+        value.time << 22
+            | ((value.machine_id as u64) << 17)
+            | ((value.thread_id as u64) << 12)
+            | (value.counter as u64)
+    }
+}
+
+fn get_time_millis(epoch: SystemTime) -> u64 {
+    SystemTime::now()
+        .duration_since(epoch)
+        .expect("time is before epoch")
+        .as_millis() as u64
+}