twsnap 0.9.2

use std::fmt;
use std::marker::PhantomData;

use hashbrown::hash_map::Entry;
use hashbrown::HashMap;
use serde::de::{MapAccess, Visitor};
use serde::ser::SerializeMap;
use serde::{Deserialize, Deserializer, Serialize};

/// Key type for [`OrderedMap`].
/// Internally it is an `u32`, and the lower 16 bits are for the legacy id.
#[derive(
    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default,
)]
pub struct SnapId(pub u32);

impl SnapId {
    pub fn from_ddnet_snap_id(id: u16) -> Self {
        Self(id.into())
    }

    pub fn ddnet_snap_id(&self) -> u16 {
        self.0 as u16
    }

    pub(crate) fn snap_from_i32(inp: i32) -> Option<SnapId> {
        Some(SnapId(u32::try_from(inp).ok()?))
    }

    pub(crate) fn to_id(&self) -> i32 {
        self.0 as i32
    }

    pub(crate) fn snap_opt_to_id(snap_id: Option<SnapId>) -> i32 {
        snap_id.map(|s| s.ddnet_snap_id().into()).unwrap_or(-1)
    }
}

/// Collection type with the two primary goals:
///
/// 1. Fast lookup from ID to item
/// 2. Consistent iteration order, sorted by ID
///
/// Could maybe just be replaced with a BTreeMap, on second though
#[derive(Debug, Clone)]
pub struct OrderedMap<T> {
    map: HashMap<SnapId, T>,
    ordering: Vec<SnapId>,
}

impl<T> Default for OrderedMap<T> {
    fn default() -> Self {
        Self {
            map: HashMap::new(),
            ordering: Vec::new(),
        }
    }
}

fn id_into_ordering(ordering: &mut Vec<SnapId>, id: SnapId) {
    if ordering.last().map_or(true, |last_id| *last_id < id) {
        // Very happy path, the id is the new largest -> also last id
        ordering.push(id);
    } else if let Err(index) = ordering.binary_search(&id) {
        // Semi happy path: id position needs to be determined with binary search
        ordering.insert(index, id);
    } else {
        // Sad path: id was already present
        panic!("Snap item ID duplicate ({})", id.0);
    }
}

impl<T> OrderedMap<T> {
    fn with_capacity(capacity: usize) -> Self {
        Self {
            map: HashMap::with_capacity(capacity),
            ordering: Vec::with_capacity(capacity),
        }
    }

    pub fn clear(&mut self) {
        self.map.clear();
        self.ordering.clear();
    }

    pub fn contains_key(&self, id: SnapId) -> bool {
        self.map.contains_key(&id)
    }

    pub fn len(&self) -> usize {
        self.ordering.len()
    }

    pub fn is_empty(&self) -> bool {
        self.ordering.is_empty()
    }

    /// Inserts key-value pair into collection.
    /// Panics: If `OrderedMap` already contains `id`.
    pub fn insert(&mut self, id: SnapId, item: T) {
        self.map.insert(id, item);
        id_into_ordering(&mut self.ordering, id);
    }

    pub fn get(&self, id: SnapId) -> Option<&T> {
        self.map.get(&id)
    }

    pub fn get_mut(&mut self, id: SnapId) -> Option<&mut T> {
        self.map.get_mut(&id)
    }

    pub fn get_mut_default(&mut self, id: SnapId) -> &mut T
    where
        T: Default,
    {
        match self.map.entry(id) {
            Entry::Occupied(item) => item.into_mut(),
            Entry::Vacant(empty) => {
                id_into_ordering(&mut self.ordering, id);
                empty.insert(T::default())
            }
        }
    }

    /// Removes entry from `OrderedMap` and returns it.
    /// Panics: If the `OrderedMap` does not contain `id`.
    pub fn remove(&mut self, id: SnapId) -> T {
        if let Ok(index) = self.ordering.binary_search(&id) {
            self.ordering.remove(index);
        } else {
            panic!("Can't remove non-existant ID")
        }
        self.map.remove(&id).expect("Internal desync")
    }

    /// Iterate over all key-value pairs in the correct order.
    pub fn iter(&self) -> Iter<'_, T> {
        Iter {
            map: &self.map,
            ordering: self.ordering.iter(),
        }
    }

    /// Iterate over all inserted values in the correct order.
    pub fn values(&self) -> Values<'_, T> {
        Values(Iter {
            map: &self.map,
            ordering: self.ordering.iter(),
        })
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, T> {
        IterMut {
            map: &mut self.map,
            ordering: self.ordering.iter(),
        }
    }

    /// Iterate over all inserted values in the correct order.
    pub fn values_mut(&mut self) -> ValuesMut<'_, T> {
        ValuesMut(IterMut {
            map: &mut self.map,
            ordering: self.ordering.iter(),
        })
    }
}

pub struct Iter<'a, T> {
    map: &'a HashMap<SnapId, T>,
    ordering: std::slice::Iter<'a, SnapId>,
}

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = (SnapId, &'a T);

    fn next(&mut self) -> Option<Self::Item> {
        self.ordering.next().map(|id| (*id, &self.map[id]))
    }
}

pub struct IterMut<'a, T> {
    map: &'a mut HashMap<SnapId, T>,
    ordering: std::slice::Iter<'a, SnapId>,
}

impl<'a, T> Iterator for IterMut<'a, T> {
    type Item = (SnapId, &'a mut T);

    fn next(&mut self) -> Option<Self::Item> {
        let id = self.ordering.next()?;
        let entry: *mut T = self.map.get_mut(id).unwrap() as *mut _;
        // Safety: We need to guarantee that we don't return two mut references to the same entry.
        // For that, `ordering` must not contain the same index twice.
        // We do that in `insert` where we panic if the same ID is inserted twice.
        Some((*id, unsafe { &mut *entry }))
    }
}

pub struct Values<'a, T>(Iter<'a, T>);

impl<'a, T> Iterator for Values<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(_id, item)| item)
    }
}

pub struct ValuesMut<'a, T>(IterMut<'a, T>);

impl<'a, T> Iterator for ValuesMut<'a, T> {
    type Item = &'a mut T;

    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(_id, item)| item)
    }
}

impl<T: Serialize> Serialize for OrderedMap<T> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let mut map = serializer.serialize_map(Some(self.len()))?;
        for (k, v) in self.iter() {
            map.serialize_entry(&k.0, v)?;
        }
        map.end()
    }
}

struct OrderedMapVisitor<T> {
    marker: PhantomData<fn() -> OrderedMap<T>>,
}

impl<T> OrderedMapVisitor<T> {
    fn new() -> Self {
        OrderedMapVisitor {
            marker: PhantomData,
        }
    }
}

impl<'de, T> Visitor<'de> for OrderedMapVisitor<T>
where
    T: Deserialize<'de>,
{
    type Value = OrderedMap<T>;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("Values sorted by their respective unique")
    }

    fn visit_map<M>(self, mut access: M) -> Result<Self::Value, M::Error>
    where
        M: MapAccess<'de>,
    {
        let mut map = OrderedMap::with_capacity(access.size_hint().unwrap_or(0));
        while let Some((key, value)) = access.next_entry()? {
            map.insert(SnapId(key), value);
        }
        Ok(map)
    }
}

impl<'de, T> Deserialize<'de> for OrderedMap<T>
where
    T: Deserialize<'de>,
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_map(OrderedMapVisitor::new())
    }
}

#[test]
fn test() {
    let mut map = OrderedMap::default();
    for n in [8, 5, 0, 3, 6, 4, 9, 7, 2, 1] {
        map.insert(SnapId(n), n);
    }
    for (i, n) in map.values().enumerate() {
        assert_eq!(i, *n as usize);
    }
    map.remove(SnapId(9));
    for (i, n) in map.values().enumerate() {
        assert_eq!(i, *n as usize);
    }
    for (k, v) in map.iter() {
        assert_eq!(map.get(k).unwrap(), v);
    }
    let mut map2 = OrderedMap::default();
    for n in [8, 5, 6, 9, 7] {
        map2.insert(SnapId(n), n);
    }
    let mut common = vec![8, 5, 6, 7];
    for (k, v) in map.iter() {
        if map2.get(k).is_some() {
            let position = common.iter().position(|e| e == v).unwrap();
            common.remove(position);
        }
    }
    assert!(common.is_empty());
}

#[test]
fn serialize_empty() {
    let map = OrderedMap::<String>::default();
    let serialized = serde_json::to_string(&map).unwrap();
    let deserialized: OrderedMap<u32> = serde_json::from_str(&serialized).unwrap();
    assert_eq!(deserialized.len(), 0);
}

#[test]
fn serialize() {
    const N: u32 = 1000;
    let mut map = OrderedMap::default();
    for n in 0..N {
        map.insert(SnapId(n), n * n);
    }
    let serialized = serde_json::to_string(&map).unwrap();
    let deserialized: OrderedMap<u32> = serde_json::from_str(&serialized).unwrap();
    assert_eq!(deserialized.len(), N as usize);
    for (n, (k, v)) in (0..N).zip(deserialized.iter()) {
        assert_eq!(k.0, n);
        assert_eq!(*v, n * n);
    }
}