HollowPS_TV_experiments/qwer/src/collection.rs
2025-06-11 05:50:41 +03:00

782 lines
23 KiB
Rust

use std::{
collections::{HashMap, HashSet},
io::Result,
};
use crate::OctData;
pub type DoubleKeyHashMap<K1, K2, V> = HashMap<K1, HashMap<K2, V>>;
#[macro_export]
macro_rules! phashmap {
($(($key:expr, $value:expr)),*) => {
PropertyHashMap::Base(
{
let mut map = std::collections::HashMap::new();
$(
map.insert($key, $value);
)*
map
}
)
};
}
#[macro_export]
macro_rules! phashset {
($($value:expr),*) => {
PropertyHashSet::Base(
{
let mut set = std::collections::HashSet::new();
$(
set.insert($value);
)*
set
}
)
};
}
#[macro_export]
macro_rules! pdkhashmap {
($(($key1:expr, $key2:expr, $value:expr)),*) => {
PropertyDoubleKeyHashMap::Base(
{
let mut map = qwer::DoubleKeyHashMap::new();
$(
map.entry($key1).or_insert(HashMap::new()).insert($key2, $value);
)*
map
}
)
};
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PropertyHashMap<K, V>
where
K: OctData + Eq + std::hash::Hash,
V: OctData,
{
Base(HashMap<K, V>),
Modify {
to_add: Vec<(K, V)>,
to_remove: Vec<K>,
},
}
impl<K, V> PropertyHashMap<K, V>
where
K: OctData + Eq + std::hash::Hash,
V: OctData,
{
pub fn insert(&mut self, key: K, value: V) {
match self {
Self::Base(base) => {
base.insert(key, value);
}
Self::Modify { to_add, .. } => {
to_add.push((key, value));
}
}
}
pub fn remove(&mut self, key: K) -> Option<V> {
match self {
Self::Base(base) => base.remove(&key),
Self::Modify { to_remove, .. } => {
to_remove.push(key);
None
}
}
}
pub fn get(&self, key: &K) -> Option<&V> {
match self {
Self::Base(base) => base.get(key),
Self::Modify { .. } => None,
}
}
pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
match self {
Self::Base(base) => base.get_mut(key),
Self::Modify { .. } => None,
}
}
#[must_use]
pub fn len(&self) -> usize {
match self {
Self::Base(base) => base.len(),
Self::Modify { to_add, to_remove } => to_add.len() + to_remove.len(),
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
match self {
Self::Base(base) => base.is_empty(),
Self::Modify { to_add, to_remove } => to_add.is_empty() && to_remove.is_empty(),
}
}
#[must_use]
pub fn iter(&self) -> std::collections::hash_map::Iter<K, V> {
match self {
Self::Base(base) => base.iter(),
Self::Modify { .. } => unreachable!("PropertyHashMap::Modify::iter()"),
}
}
pub fn iter_mut(&mut self) -> std::collections::hash_map::IterMut<K, V> {
match self {
Self::Base(base) => base.iter_mut(),
Self::Modify { .. } => unreachable!("PropertyHashMap::Modify::iter_mut()"),
}
}
}
impl<K, V> IntoIterator for PropertyHashMap<K, V>
where
K: OctData + Eq + std::hash::Hash,
V: OctData,
{
type IntoIter = std::collections::hash_map::IntoIter<K, V>;
type Item = (K, V);
fn into_iter(self) -> Self::IntoIter {
match self {
Self::Base(base) => base.into_iter(),
Self::Modify { .. } => unreachable!("PropertyHashMap::Modify::into_iter()"),
}
}
}
impl<'a, K, V> IntoIterator for &'a PropertyHashMap<K, V>
where
K: OctData + Eq + std::hash::Hash,
V: OctData,
{
type IntoIter = std::collections::hash_map::Iter<'a, K, V>;
type Item = (&'a K, &'a V);
fn into_iter(self) -> Self::IntoIter {
match self {
PropertyHashMap::Base(base) => base.iter(),
PropertyHashMap::Modify { .. } => unreachable!("PropertyHashMap::Modify::into_iter()"),
}
}
}
impl<'a, K, V> IntoIterator for &'a mut PropertyHashMap<K, V>
where
K: OctData + Eq + std::hash::Hash,
V: OctData,
{
type IntoIter = std::collections::hash_map::IterMut<'a, K, V>;
type Item = (&'a K, &'a mut V);
fn into_iter(self) -> Self::IntoIter {
match self {
PropertyHashMap::Base(base) => base.iter_mut(),
PropertyHashMap::Modify { .. } => unreachable!("PropertyHashMap::Modify::into_iter()"),
}
}
}
#[test]
fn test_hashmap_iter() {
let mut map = phashmap![(1, 2), (3, 4)];
let mut expecting = vec![(1, 2), (3, 4)];
let iter = map.iter_mut();
for (key, value) in iter {
assert!(expecting.contains(&(*key, *value)));
expecting.retain(|x| x != &(*key, *value));
*value += 1;
}
}
impl<T> PropertyHashSet<T>
where
T: OctData + Eq + std::hash::Hash,
{
pub fn insert(&mut self, value: T) {
match self {
Self::Base(base) => {
base.insert(value);
}
Self::Modify { to_add, .. } => {
to_add.push(value);
}
}
}
pub fn remove(&mut self, value: T) {
match self {
Self::Base(base) => {
base.remove(&value);
}
Self::Modify { to_remove, .. } => {
to_remove.push(value);
}
}
}
#[must_use]
pub fn len(&self) -> usize {
match self {
Self::Base(base) => base.len(),
Self::Modify { to_add, to_remove } => to_add.len() + to_remove.len(),
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
match self {
Self::Base(base) => base.is_empty(),
Self::Modify { to_add, to_remove } => to_add.is_empty() && to_remove.is_empty(),
}
}
#[must_use]
pub fn iter(&self) -> std::collections::hash_set::Iter<T> {
match self {
Self::Base(base) => base.iter(),
Self::Modify { .. } => unreachable!("PropertyHashSet::Modify::iter()"),
}
}
#[must_use]
pub fn iter_mut(&mut self) -> std::collections::hash_set::Iter<T> {
self.into_iter()
}
}
impl<T> IntoIterator for PropertyHashSet<T>
where
T: OctData + Eq + std::hash::Hash,
{
type IntoIter = std::collections::hash_set::IntoIter<T>;
type Item = T;
fn into_iter(self) -> Self::IntoIter {
match self {
Self::Base(base) => base.into_iter(),
Self::Modify { .. } => unreachable!("PropertyHashSet::Modify::into_iter()"),
}
}
}
impl<'a, T> IntoIterator for &'a PropertyHashSet<T>
where
T: OctData + Eq + std::hash::Hash,
{
type IntoIter = std::collections::hash_set::Iter<'a, T>;
type Item = &'a T;
fn into_iter(self) -> Self::IntoIter {
match self {
PropertyHashSet::Base(base) => base.iter(),
PropertyHashSet::Modify { .. } => unreachable!("PropertyHashSet::Modify::into_iter()"),
}
}
}
impl<'a, T> IntoIterator for &'a mut PropertyHashSet<T>
where
T: OctData + Eq + std::hash::Hash,
{
type IntoIter = std::collections::hash_set::Iter<'a, T>;
type Item = &'a T;
fn into_iter(self) -> Self::IntoIter {
match self {
PropertyHashSet::Base(base) => base.iter(),
PropertyHashSet::Modify { .. } => unreachable!("PropertyHashSet::Modify::into_iter()"),
}
}
}
#[test]
fn test_phashmap_macro() {
let map = phashmap![(1, 2), (3, 4)];
assert_eq!(
map,
PropertyHashMap::Base([(1, 2), (3, 4)].into_iter().collect())
);
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PropertyHashSet<T>
where
T: OctData + Eq + std::hash::Hash,
{
Base(HashSet<T>),
Modify { to_add: Vec<T>, to_remove: Vec<T> },
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PropertyDoubleKeyHashMap<K1, K2, V>
where
K1: OctData + Eq + std::hash::Hash,
K2: OctData + Eq + std::hash::Hash,
V: OctData,
{
Base(DoubleKeyHashMap<K1, K2, V>),
Modify {
to_add: Vec<(K1, K2, V)>,
to_remove: Vec<(K1, K2)>,
},
}
pub struct PropertyDoubleKeyHashMapIterMut<'a, K1, K2, V> {
outer_iter: std::collections::hash_map::IterMut<'a, K1, HashMap<K2, V>>,
inner_iter: Option<std::collections::hash_map::IterMut<'a, K2, V>>,
current_outer_key: Option<&'a K1>,
}
pub struct PropertyDoubleKeyHashMapIter<'a, K1, K2, V> {
outer_iter: std::collections::hash_map::Iter<'a, K1, HashMap<K2, V>>,
inner_iter: Option<std::collections::hash_map::Iter<'a, K2, V>>,
current_outer_key: Option<&'a K1>,
}
pub struct PropertyDoubleKeyHashMapIntoIter<K1, K2, V> {
outer_iter: std::collections::hash_map::IntoIter<K1, HashMap<K2, V>>,
inner_iter: Option<std::collections::hash_map::IntoIter<K2, V>>,
current_outer_key: Option<K1>,
}
impl<'a, K1, K2, V> Iterator for PropertyDoubleKeyHashMapIterMut<'a, K1, K2, V>
where
K1: OctData + Eq + std::hash::Hash,
K2: OctData + Eq + std::hash::Hash,
V: OctData,
{
type Item = (&'a K1, &'a K2, &'a mut V);
fn next(&mut self) -> Option<Self::Item> {
loop {
if self.inner_iter.is_none() {
if let Some((key1, sub_map)) = self.outer_iter.next() {
self.current_outer_key = Some(key1);
self.inner_iter = Some(sub_map.iter_mut());
} else {
return None;
}
}
if let Some(inner_iter) = &mut self.inner_iter {
if let Some((key2, value)) = inner_iter.next() {
return Some((self.current_outer_key.unwrap(), key2, value));
}
self.inner_iter = None;
}
}
}
}
impl<'a, K1, K2, V> Iterator for PropertyDoubleKeyHashMapIter<'a, K1, K2, V>
where
K1: OctData + Eq + std::hash::Hash,
K2: OctData + Eq + std::hash::Hash,
V: OctData,
{
type Item = (&'a K1, &'a K2, &'a V);
fn next(&mut self) -> Option<Self::Item> {
loop {
if self.inner_iter.is_none() {
if let Some((key1, sub_map)) = self.outer_iter.next() {
self.current_outer_key = Some(key1);
self.inner_iter = Some(sub_map.iter());
} else {
return None;
}
}
if let Some(inner_iter) = &mut self.inner_iter {
if let Some((key2, value)) = inner_iter.next() {
return Some((self.current_outer_key.unwrap(), key2, value));
}
self.inner_iter = None;
}
}
}
}
impl<K1, K2, V> Iterator for PropertyDoubleKeyHashMapIntoIter<K1, K2, V>
where
K1: OctData + Eq + std::hash::Hash + Copy,
K2: OctData + Eq + std::hash::Hash,
V: OctData,
{
type Item = (K1, K2, V);
fn next(&mut self) -> Option<Self::Item> {
loop {
if self.inner_iter.is_none() {
if let Some((key1, sub_map)) = self.outer_iter.next() {
self.current_outer_key = Some(key1);
self.inner_iter = Some(sub_map.into_iter());
} else {
return None;
}
}
if let Some(inner_iter) = &mut self.inner_iter {
if let Some((key2, value)) = inner_iter.next() {
return Some((self.current_outer_key.unwrap(), key2, value));
}
self.inner_iter = None;
}
}
}
}
impl<K1, K2, V> IntoIterator for PropertyDoubleKeyHashMap<K1, K2, V>
where
K1: OctData + Eq + std::hash::Hash + Copy,
K2: OctData + Eq + std::hash::Hash,
V: OctData,
{
type IntoIter = PropertyDoubleKeyHashMapIntoIter<K1, K2, V>;
type Item = (K1, K2, V);
fn into_iter(self) -> Self::IntoIter {
match self {
Self::Base(base) => PropertyDoubleKeyHashMapIntoIter {
outer_iter: base.into_iter(),
inner_iter: None,
current_outer_key: None,
},
Self::Modify { .. } => {
unreachable!("PropertyDoubleKeyHashMap::Modify::into_iter() is not implemented")
}
}
}
}
impl<'a, K1, K2, V> IntoIterator for &'a PropertyDoubleKeyHashMap<K1, K2, V>
where
K1: OctData + Eq + std::hash::Hash,
K2: OctData + Eq + std::hash::Hash,
V: OctData,
{
type IntoIter = PropertyDoubleKeyHashMapIter<'a, K1, K2, V>;
type Item = (&'a K1, &'a K2, &'a V);
fn into_iter(self) -> Self::IntoIter {
match self {
PropertyDoubleKeyHashMap::Base(base) => PropertyDoubleKeyHashMapIter {
outer_iter: base.iter(),
inner_iter: None,
current_outer_key: None,
},
PropertyDoubleKeyHashMap::Modify { .. } => {
unreachable!("PropertyDoubleKeyHashMap::Modify::into_iter() is not implemented")
}
}
}
}
impl<'a, K1, K2, V> IntoIterator for &'a mut PropertyDoubleKeyHashMap<K1, K2, V>
where
K1: OctData + Eq + std::hash::Hash,
K2: OctData + Eq + std::hash::Hash,
V: OctData,
{
type IntoIter = PropertyDoubleKeyHashMapIterMut<'a, K1, K2, V>;
type Item = (&'a K1, &'a K2, &'a mut V);
fn into_iter(self) -> Self::IntoIter {
match self {
PropertyDoubleKeyHashMap::Base(base) => PropertyDoubleKeyHashMapIterMut {
outer_iter: base.iter_mut(),
inner_iter: None,
current_outer_key: None,
},
PropertyDoubleKeyHashMap::Modify { .. } => {
unreachable!("PropertyDoubleKeyHashMap::Modify::into_iter() is not implemented")
}
}
}
}
impl<K1, K2, V> PropertyDoubleKeyHashMap<K1, K2, V>
where
K1: OctData + Eq + std::hash::Hash,
K2: OctData + Eq + std::hash::Hash,
V: OctData,
{
pub fn insert(&mut self, key: K1, sub_key: K2, value: V) {
match self {
Self::Base(base) => {
base.entry(key)
.or_insert_with(HashMap::new)
.insert(sub_key, value);
}
Self::Modify { to_add, .. } => {
to_add.push((key, sub_key, value));
}
}
}
pub fn get(&self, key: &K1, sub_key: &K2) -> Option<&V> {
match self {
Self::Base(base) => base.get(key).and_then(|sub_map| sub_map.get(sub_key)),
Self::Modify { .. } => {
unreachable!("PropertyDoubleKeyHashMap::Modify::get() is not implemented")
}
}
}
pub fn get_mut(&mut self, key: &K1, sub_key: &K2) -> Option<&mut V> {
match self {
Self::Base(base) => base
.get_mut(key)
.and_then(|sub_map| sub_map.get_mut(sub_key)),
Self::Modify { .. } => {
unreachable!("PropertyDoubleKeyHashMap::Modify::get_mut() is not implemented")
}
}
}
pub fn remove(&mut self, key: K1, sub_key: K2) -> Option<V> {
match self {
Self::Base(base) => {
let mut removed = None;
base.entry(key)
.and_modify(|sub_map| removed = sub_map.remove(&sub_key));
removed
}
Self::Modify { to_remove, .. } => {
to_remove.push((key, sub_key));
None
}
}
}
#[must_use]
pub fn len(&self) -> usize {
match self {
Self::Base(base) => base.len(),
Self::Modify { to_add, to_remove } => to_add.len() + to_remove.len(),
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
match self {
Self::Base(base) => base.is_empty(),
Self::Modify { to_add, to_remove } => to_add.is_empty() && to_remove.is_empty(),
}
}
#[must_use]
pub fn iter(&self) -> PropertyDoubleKeyHashMapIter<K1, K2, V> {
match self {
Self::Base(base) => PropertyDoubleKeyHashMapIter {
outer_iter: base.iter(),
inner_iter: None,
current_outer_key: None,
},
Self::Modify { .. } => {
unreachable!("PropertyDoubleKeyHashMap::Modify::iter() is not implemented")
}
}
}
pub fn iter_mut(&mut self) -> PropertyDoubleKeyHashMapIterMut<K1, K2, V> {
match self {
Self::Base(base) => PropertyDoubleKeyHashMapIterMut {
outer_iter: base.iter_mut(),
inner_iter: None,
current_outer_key: None,
},
Self::Modify { .. } => {
unreachable!("PropertyDoubleKeyHashMap::Modify::iter_mut() is not implemented")
}
}
}
}
#[test]
fn test_dkhashmap_iter() {
let mut map = PropertyDoubleKeyHashMap::Base(
std::iter::once((1, [(2, 3), (4, 5)].into_iter().collect())).collect(),
);
let mut expecting = vec![(1, 2, 3), (1, 4, 5)];
let iter = map.iter_mut();
for (key1, key2, value) in iter {
assert!(expecting.contains(&(*key1, *key2, *value)));
expecting.retain(|x| x != &(*key1, *key2, *value));
*value += 1;
}
}
impl<K, V> OctData for PropertyHashMap<K, V>
where
K: OctData + Eq + Ord + std::hash::Hash,
V: OctData,
{
fn marshal_to<W: std::io::Write>(&self, w: &mut W, bt_property_tag: u16) -> Result<()> {
match self {
Self::Base(map) => {
map.marshal_to(w, bt_property_tag)?;
}
Self::Modify { to_add, to_remove } => {
let len = -(to_add.len() as i32 + to_remove.len() as i32);
len.marshal_to(w, bt_property_tag)?;
for (key, value) in to_add {
key.marshal_to(w, bt_property_tag)?;
false.marshal_to(w, bt_property_tag)?;
value.marshal_to(w, bt_property_tag)?;
}
for key in to_remove {
key.marshal_to(w, bt_property_tag)?;
true.marshal_to(w, bt_property_tag)?;
}
}
}
Ok(())
}
fn unmarshal_from<R: std::io::Read>(r: &mut R, bt_property_tag: u16) -> Result<Self> {
let len = i32::unmarshal_from(r, bt_property_tag)?;
if len >= 0 {
let mut map = HashMap::with_capacity(len as usize);
for _ in 0..len {
map.insert(
K::unmarshal_from(r, bt_property_tag)?,
V::unmarshal_from(r, bt_property_tag)?,
);
}
Ok(Self::Base(map))
} else {
let mut to_add = Vec::new();
let mut to_remove = Vec::new();
for _ in 0..-len {
let key = K::unmarshal_from(r, bt_property_tag)?;
if !bool::unmarshal_from(r, bt_property_tag)? {
to_add.push((key, V::unmarshal_from(r, bt_property_tag)?));
} else {
to_remove.push(key);
}
}
Ok(Self::Modify { to_add, to_remove })
}
}
}
impl<K> OctData for PropertyHashSet<K>
where
K: OctData + Eq + std::hash::Hash,
{
fn marshal_to<W: std::io::Write>(&self, w: &mut W, bt_property_tag: u16) -> Result<()> {
match self {
Self::Base(set) => {
set.marshal_to(w, bt_property_tag)?;
}
Self::Modify { to_add, to_remove } => {
let len = -(to_add.len() as i32 + to_remove.len() as i32);
len.marshal_to(w, bt_property_tag)?;
for value in to_add {
value.marshal_to(w, bt_property_tag)?;
false.marshal_to(w, bt_property_tag)?;
}
for value in to_remove {
value.marshal_to(w, bt_property_tag)?;
true.marshal_to(w, bt_property_tag)?;
}
}
}
Ok(())
}
fn unmarshal_from<R: std::io::Read>(r: &mut R, bt_property_tag: u16) -> Result<Self> {
let len = i32::unmarshal_from(r, bt_property_tag)?;
if len >= 0 {
let mut set = HashSet::with_capacity(len as usize);
for _ in 0..len {
set.insert(K::unmarshal_from(r, bt_property_tag)?);
}
Ok(Self::Base(set))
} else {
let mut to_add = Vec::new();
let mut to_remove = Vec::new();
for _ in 0..-len {
let value = K::unmarshal_from(r, bt_property_tag)?;
if !bool::unmarshal_from(r, bt_property_tag)? {
to_add.push(value);
} else {
to_remove.push(value);
}
}
Ok(Self::Modify { to_add, to_remove })
}
}
}
impl<K1, K2, V> OctData for PropertyDoubleKeyHashMap<K1, K2, V>
where
K1: OctData + Eq + Ord + std::hash::Hash,
K2: OctData + Eq + Ord + std::hash::Hash,
V: OctData,
{
fn marshal_to<W: std::io::Write>(&self, w: &mut W, bt_property_tag: u16) -> Result<()> {
match self {
Self::Base(map) => {
map.marshal_to(w, bt_property_tag)?;
}
Self::Modify { to_add, to_remove } => {
let len = -(to_add.len() as i32 + to_remove.len() as i32);
len.marshal_to(w, bt_property_tag)?;
for (key1, key2, value) in to_add {
key1.marshal_to(w, bt_property_tag)?;
key2.marshal_to(w, bt_property_tag)?;
false.marshal_to(w, bt_property_tag)?;
value.marshal_to(w, bt_property_tag)?;
}
for (key1, key2) in to_remove {
key1.marshal_to(w, bt_property_tag)?;
key2.marshal_to(w, bt_property_tag)?;
true.marshal_to(w, bt_property_tag)?;
}
}
}
Ok(())
}
fn unmarshal_from<R: std::io::Read>(r: &mut R, bt_property_tag: u16) -> Result<Self> {
let len = i32::unmarshal_from(r, bt_property_tag)?;
if len >= 0 {
let mut map = HashMap::new();
for _ in 0..len {
let key1 = K1::unmarshal_from(r, bt_property_tag)?;
let key2 = K2::unmarshal_from(r, bt_property_tag)?;
let value = V::unmarshal_from(r, bt_property_tag)?;
map.entry(key1)
.or_insert_with(HashMap::new)
.insert(key2, value);
}
Ok(Self::Base(map))
} else {
let mut to_add = Vec::new();
let mut to_remove = Vec::new();
for _ in 0..-len {
let key1 = K1::unmarshal_from(r, bt_property_tag)?;
let key2 = K2::unmarshal_from(r, bt_property_tag)?;
if !bool::unmarshal_from(r, bt_property_tag)? {
to_add.push((key1, key2, V::unmarshal_from(r, bt_property_tag)?));
} else {
to_remove.push((key1, key2));
}
}
Ok(Self::Modify { to_add, to_remove })
}
}
}