use std::{
collections::{HashMap, HashSet},
hash::Hash,
vec::IntoIter,
};
pub trait UnionZip<K, M> {
fn union_zip<'a>(&'a self, other: &'a M) -> UnionZipped<'a, K, M>;
}
pub trait IntoUnionZip<'a, K, M> {
fn into_union_zip(self, other: &'a M) -> UnionZipped<'a, K, M>;
}
pub trait AsUnionZipped<'a, K, M> {
fn as_union_zipped(&'a self) -> UnionZipped<'a, K, M>;
}
pub struct UnionZipped<'a, K, M> {
keys: Vec<K>,
maps: Vec<&'a M>,
}
pub struct UnionZippedIter<'a, K, M, I: Iterator<Item = K>> {
keys: I,
maps: Vec<&'a M>,
}
impl<K, V> UnionZip<K, HashMap<K, V>> for HashMap<K, V>
where
K: Hash + Eq + Clone + Ord,
{
fn union_zip<'a>(
&'a self,
other: &'a Self,
) -> UnionZipped<'a, K, HashMap<K, V>> {
let mut keys: Vec<K> = self
.keys()
.collect::<HashSet<&K>>()
.union(&other.keys().collect::<HashSet<&K>>())
.map(|&k| k.clone())
.collect();
keys.sort();
UnionZipped {
keys,
maps: vec![&self, other],
}
}
}
impl<'a, K, V> IntoUnionZip<'a, K, HashMap<K, V>>
for UnionZipped<'a, K, HashMap<K, V>>
where
K: Hash + Eq + Clone + Ord,
{
fn into_union_zip(
self,
other: &'a HashMap<K, V>,
) -> UnionZipped<'a, K, HashMap<K, V>> {
let mut keys: Vec<K> = self
.keys
.iter()
.collect::<HashSet<&K>>()
.union(&other.keys().collect::<HashSet<&K>>())
.map(|&k| k.clone())
.collect();
keys.sort();
let mut maps = self.maps;
maps.push(other);
UnionZipped {
keys,
maps,
}
}
}
impl<'a, K, V> AsUnionZipped<'a, K, HashMap<K, V>> for HashMap<K, V>
where
K: Clone + Ord,
{
fn as_union_zipped(&'a self) -> UnionZipped<'a, K, HashMap<K, V>> {
let mut keys: Vec<K> = self.keys().cloned().collect();
keys.sort();
UnionZipped {
keys,
maps: vec![&self],
}
}
}
impl<'a, K, V> IntoIterator for UnionZipped<'a, K, HashMap<K, V>>
where
K: Hash + Eq,
{
type IntoIter = UnionZippedIter<'a, K, HashMap<K, V>, IntoIter<K>>;
type Item =
<UnionZippedIter<'a, K, HashMap<K, V>, IntoIter<K>> as Iterator>::Item;
fn into_iter(self) -> Self::IntoIter {
UnionZippedIter {
keys: self.keys.into_iter(),
maps: self.maps,
}
}
}
impl<'a, K, V> Iterator for UnionZippedIter<'a, K, HashMap<K, V>, IntoIter<K>>
where
K: Hash + Eq,
{
type Item = (K, Vec<Option<&'a V>>);
fn next(&mut self) -> Option<Self::Item> {
match self.keys.next() {
None => None,
Some(k) => {
let mapped: Vec<Option<&'a V>> = self
.maps
.iter()
.map(|m| {
if m.contains_key(&k) {
Some(&m[&k])
} else {
None
}
})
.collect();
Some((k, mapped))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_union_zip_iter_hashmap() {
let m1: HashMap<i32, i32> = vec![
(1, 10),
(3, 23),
(4, 20),
(12, 6),
]
.into_iter()
.collect();
let m2: HashMap<i32, i32> = vec![
(0, 4),
(1, 20),
(4, 20),
(9, 29),
]
.into_iter()
.collect();
let mut iter = m1.union_zip(&m2).into_iter();
assert_eq!(Some((0, vec![None, Some(&4)])), iter.next());
assert_eq!(Some((1, vec![Some(&10), Some(&20)])), iter.next());
assert_eq!(Some((3, vec![Some(&23), None])), iter.next());
assert_eq!(Some((4, vec![Some(&20), Some(&20)])), iter.next());
assert_eq!(Some((9, vec![None, Some(&29)])), iter.next());
assert_eq!(Some((12, vec![Some(&6), None])), iter.next());
assert_eq!(None, iter.next());
let m3: HashMap<i32, i32> = vec![
(0, 9),
(3, 43),
(4, 8),
(14, 68),
]
.into_iter()
.collect();
let mut iter2 = m1.union_zip(&m2).into_union_zip(&m3).into_iter();
assert_eq!(Some((0, vec![None, Some(&4), Some(&9)])), iter2.next());
assert_eq!(Some((1, vec![Some(&10), Some(&20), None])), iter2.next());
assert_eq!(Some((3, vec![Some(&23), None, Some(&43)])), iter2.next());
assert_eq!(
Some((4, vec![Some(&20), Some(&20), Some(&8)])),
iter2.next()
);
assert_eq!(Some((9, vec![None, Some(&29), None])), iter2.next());
assert_eq!(Some((12, vec![Some(&6), None, None])), iter2.next());
assert_eq!(Some((14, vec![None, None, Some(&68)])), iter2.next());
assert_eq!(None, iter2.next());
let m4: HashMap<i32, i32> = vec![
(4, 73),
(14, 64),
]
.into_iter()
.collect();
let mut iter3 = m1
.union_zip(&m2)
.into_union_zip(&m3)
.into_union_zip(&m4)
.into_iter();
assert_eq!(
Some((0, vec![None, Some(&4), Some(&9), None])),
iter3.next()
);
assert_eq!(
Some((1, vec![Some(&10), Some(&20), None, None])),
iter3.next()
);
assert_eq!(
Some((3, vec![Some(&23), None, Some(&43), None])),
iter3.next()
);
assert_eq!(
Some((4, vec![Some(&20), Some(&20), Some(&8), Some(&73)])),
iter3.next()
);
assert_eq!(Some((9, vec![None, Some(&29), None, None])), iter3.next());
assert_eq!(Some((12, vec![Some(&6), None, None, None])), iter3.next());
assert_eq!(
Some((14, vec![None, None, Some(&68), Some(&64)])),
iter3.next()
);
assert_eq!(None, iter3.next());
}
#[test]
fn test_as_union_zipped() {
let m: HashMap<i32, i32> =
vec![(2, 10), (3, 0), (4, 1)].into_iter().collect();
let union_zipped = m.as_union_zipped();
assert_eq!(union_zipped.maps[0], &m);
}
}