use core::fmt;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::collections::HashMap;
use std::marker::PhantomData;
use std::ops::{Index, IndexMut};
pub trait Id: Copy + Clone + Eq + PartialEq + fmt::Debug + Into<u64> + From<u64> {
const NONE: Self;
fn none() -> Self {
Self::NONE
}
fn is_none(&self) -> bool {
self.eq(&Self::NONE)
}
fn as_u64(&self) -> u64 {
(*self).into()
}
fn from_u64(val: u64) -> Self {
Self::from(val)
}
}
#[macro_export]
macro_rules! new_id_type {
() => {};
(
$(#[$meta:meta])*
$vis:vis struct $name:ident;
$($rest:tt)*
) => {
$(#[$meta])*
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
$vis struct $name(pub u64);
impl From<u64> for $name {
#[inline]
fn from(val: u64) -> Self {
Self(val)
}
}
impl From<$name> for u64 {
#[inline]
fn from(id: $name) -> Self {
id.0
}
}
impl $crate::Id for $name {
const NONE: Self = Self(0);
}
impl serde::Serialize for $name {
#[inline]
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
self.0.serialize(serializer)
}
}
impl<'de> serde::Deserialize<'de> for $name {
#[inline]
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let val = u64::deserialize(deserializer)?;
Ok(Self(val))
}
}
$crate::new_id_type!($($rest)*);
};
}
new_id_type!{
pub struct DefaultId;
}
#[derive(Debug, Clone)]
pub struct IdMap<K: Id, V> {
pub(crate) inner: HashMap<u64, V>,
max_id: u64,
_marker: PhantomData<K>,
}
impl<V> IdMap<DefaultId, V> {
pub fn new() -> Self { Self::with_id_capacity(0) }
pub fn with_capacity(capacity: usize) -> Self { Self::with_id_capacity(capacity) }
}
impl<K: Id ,V> Default for IdMap<K, V> {
fn default() -> Self {
Self::with_id()
}
}
impl<K: Id, V> IdMap<K, V> {
pub fn with_id() -> Self {
Self {
inner: HashMap::new(),
max_id: 0,
_marker: PhantomData,
}
}
pub fn with_id_capacity(capacity: usize) -> Self {
Self {
inner: HashMap::with_capacity(capacity),
max_id: 0,
_marker: PhantomData,
}
}
pub fn insert(&mut self, value: V) -> K {
self.max_id += 1; let id_u64 = self.max_id;
self.inner.insert(id_u64, value); K::from_u64(id_u64) }
pub fn insert_with_id(&mut self, id: K, value: V) -> Option<V> {
let id_u64 = id.as_u64();
if id_u64 > self.max_id {
self.max_id = id_u64;
}
self.inner.insert(id_u64, value)
}
pub fn from_vec(values: Vec<V>) -> (Self, Vec<K>) {
let mut map = Self {
inner: HashMap::with_capacity(values.len()),
max_id: 0,
_marker: PhantomData,
};
let ids = values
.into_iter()
.map(|val| {
map.max_id += 1;
let id_u64 = map.max_id;
map.inner.insert(id_u64, val);
K::from_u64(id_u64)
})
.collect();
(map, ids)
}
pub fn insert_cyclic<F>(&mut self, f: F) -> K
where
F: FnOnce(K) -> V,
{
self.max_id += 1;
let new_id = K::from_u64(self.max_id);
let value = f(new_id);
self.inner.insert(self.max_id, value);
new_id
}
pub fn get(&self, id: K) -> Option<&V> {
self.inner.get(&id.as_u64())
}
pub fn get_mut(&mut self, id: K) -> Option<&mut V> {
self.inner.get_mut(&id.as_u64())
}
pub fn remove(&mut self, id: K) -> Option<V> {
self.inner.remove(&id.as_u64())
}
pub fn contains_id(&self, id: K) -> bool {
self.inner.contains_key(&id.as_u64())
}
pub fn max_id(&self) -> K {
K::from_u64(self.max_id)
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
pub fn clear(&mut self) {
self.inner.clear();
}
}
impl<K: Id, V: Serialize> Serialize for IdMap<K, V> {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
self.inner.serialize(serializer)
}
}
impl<'de, K: Id, V: Deserialize<'de>> Deserialize<'de> for IdMap<K, V> {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let inner: HashMap<u64, V> = HashMap::deserialize(deserializer)?;
let max_id = inner.keys().copied().max().unwrap_or(0);
Ok(IdMap {
inner,
max_id,
_marker: PhantomData,
})
}
}
use std::collections::hash_map;
use std::iter::{Map};
impl<K: Id, V> IntoIterator for IdMap<K, V> {
type Item = (K, V);
type IntoIter = Map<hash_map::IntoIter<u64, V>, fn((u64, V)) -> (K, V)>;
fn into_iter(self) -> Self::IntoIter {
self.inner.into_iter().map(|(id_u64, v)| (K::from_u64(id_u64), v))
}
}
impl<'a, K: Id, V> IntoIterator for &'a IdMap<K, V> {
type Item = (K, &'a V);
type IntoIter = Map<hash_map::Iter<'a, u64, V>, fn((&'a u64, &'a V)) -> (K, &'a V)>;
fn into_iter(self) -> Self::IntoIter {
self.inner.iter().map(|(id_u64, v)| (K::from_u64(*id_u64), v))
}
}
impl<'a, K: Id, V> IntoIterator for &'a mut IdMap<K, V> {
type Item = (K, &'a mut V);
type IntoIter = Map<hash_map::IterMut<'a, u64, V>, fn((&'a u64, &'a mut V)) -> (K, &'a mut V)>;
fn into_iter(self) -> Self::IntoIter {
self.inner.iter_mut().map(|(id_u64, v)| (K::from_u64(*id_u64), v))
}
}
impl<K: Id, V> IdMap<K, V> {
pub fn iter<'a>(&'a self) -> Map<hash_map::Iter<'a, u64, V>, fn((&'a u64, &'a V)) -> (K, &'a V)> {
self.inner.iter().map(|(id_u64, v)| (K::from_u64(*id_u64), v))
}
pub fn iter_mut<'a>(&'a mut self) -> Map<hash_map::IterMut<'a, u64, V>, fn((&'a u64, &'a mut V)) -> (K, &'a mut V)> {
self.inner.iter_mut().map(|(id_u64, v)| (K::from_u64(*id_u64), v))
}
pub fn keys(&self) -> Map<hash_map::Keys<'_, u64, V>, fn(&u64) -> K> {
self.inner.keys().map(|id_u64| K::from_u64(*id_u64))
}
pub fn values(&self) -> hash_map::Values<'_, u64, V> {
self.inner.values()
}
pub fn values_mut(&mut self) -> hash_map::ValuesMut<'_, u64, V> {
self.inner.values_mut()
}
}
impl<K: Id, V> Index<K> for IdMap<K, V> {
type Output = V;
fn index(&self, id: K) -> &Self::Output {
self.get(id).expect("invalid IdMap id")
}
}
impl<K: Id, V> IndexMut<K> for IdMap<K, V> {
fn index_mut(&mut self, id: K) -> &mut Self::Output {
self.get_mut(id).expect("invalid IdMap id")
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json;
#[test]
fn test_default_id_auto_generate() {
let mut map = IdMap::new();
let id1 = map.insert("hello");
let id2 = map.insert("world");
let id3 = map.insert("rust");
assert_eq!(id1, DefaultId(1));
assert_eq!(id2, DefaultId(2));
assert_eq!(id3, DefaultId(3));
assert_eq!(map.get(id1), Some(&"hello"));
assert_eq!(map[id2], "world");
assert_eq!(map.max_id(), DefaultId(3));
map.remove(id2);
assert_eq!(map.max_id(), DefaultId(3));
let id4 = map.insert("new value");
assert_eq!(id4, DefaultId(4));
assert_eq!(map.len(), 3);
map.clear();
assert!(map.is_empty());
}
new_id_type! {
struct MyId;
}
#[test]
fn test_custom_id() {
let mut map = IdMap::<MyId, u32>::with_id();
let id1 = map.insert(42);
let id2 = map.insert(100);
assert_eq!(id1, MyId(1));
assert_eq!(id2, MyId(2));
assert_eq!(map.get(id1), Some(&42));
map.remove(id1);
assert!(!map.contains_id(id1));
}
#[test]
fn test_id_serde() {
let id = DefaultId(123456789);
let json = serde_json::to_string(&id).unwrap();
assert_eq!(json, "123456789"); let id2: DefaultId = serde_json::from_str(&json).unwrap();
assert_eq!(id2, id);
let my_id = MyId(987654321);
let json = serde_json::to_string(&my_id).unwrap();
let my_id2: MyId = serde_json::from_str(&json).unwrap();
assert_eq!(my_id2, my_id);
}
}