use once_cell::sync::OnceCell;
use crate::{
arq::*,
error::{GossipError, GossipResult},
op::OpRegion,
spacetime::*,
};
use derivative::Derivative;
use super::{Region, RegionCoords, RegionData, RegionDataConstraints};
#[derive(Debug, PartialEq, Eq, derive_more::Constructor, serde::Serialize, serde::Deserialize)]
#[cfg_attr(feature = "test_utils", derive(Clone))]
pub struct RegionCoordSetLtcs {
pub(super) times: TelescopingTimes,
pub(super) arq_set: ArqSet,
}
impl RegionCoordSetLtcs {
#[cfg_attr(not(feature = "test_utils"), deprecated = "use into_region_set")]
pub(crate) fn region_coords_flat(
&self,
) -> impl Iterator<Item = ((usize, usize, usize), RegionCoords)> + '_ {
self.region_coords_nested().flatten().flatten()
}
#[cfg_attr(not(feature = "test_utils"), deprecated = "use into_region_set")]
pub(crate) fn region_coords_nested(
&self,
) -> impl Iterator<
Item = impl Iterator<Item = impl Iterator<Item = ((usize, usize, usize), RegionCoords)>> + '_,
> + '_ {
let arqs = self.arq_set.arqs();
arqs.iter().enumerate().map(move |(ia, arq)| {
arq.segments().enumerate().map(move |(ix, x)| {
self.times
.segments()
.into_iter()
.enumerate()
.map(move |(it, t)| ((ia, ix, it), RegionCoords::new(x, t)))
})
})
}
pub fn into_region_set<D, F, E>(self, mut f: F) -> Result<RegionSetLtcs<D>, E>
where
D: RegionDataConstraints,
F: FnMut(((usize, usize, usize), RegionCoords)) -> Result<D, E>,
{
let data = self
.region_coords_nested()
.map(|arqdata| {
arqdata
.map(|column| column.map(&mut f).collect::<Result<Vec<D>, E>>())
.collect::<Result<Vec<Vec<D>>, E>>()
})
.collect::<Result<Vec<Vec<Vec<D>>>, E>>()?;
Ok(RegionSetLtcs::from_data(self, data))
}
pub fn into_region_set_infallible<D, F>(self, mut f: F) -> RegionSetLtcs<D>
where
D: RegionDataConstraints,
F: FnMut(((usize, usize, usize), RegionCoords)) -> D,
{
self.into_region_set(|c| Result::<D, std::convert::Infallible>::Ok(f(c)))
.unwrap()
}
pub fn empty() -> Self {
Self {
times: TelescopingTimes::empty(),
arq_set: ArqSet::empty(),
}
}
pub fn num_space_chunks(&self) -> usize {
self.arq_set.arqs().len()
}
pub fn count(&self) -> usize {
let nt = self.times.segments().len();
self.arq_set.arqs().iter().map(|a| a.count()).sum::<u32>() as usize * nt
}
}
#[derive(serde::Serialize, serde::Deserialize, Derivative)]
#[derivative(PartialEq, Eq)]
#[cfg_attr(feature = "test_utils", derive(Clone))]
pub struct RegionSetLtcs<D: RegionDataConstraints = RegionData> {
pub coords: RegionCoordSetLtcs,
#[derivative(PartialEq = "ignore")]
#[serde(skip)]
pub(crate) _region_coords: OnceCell<Vec<RegionCoords>>,
#[serde(bound(deserialize = "D: serde::de::DeserializeOwned"))]
data: Vec<Vec<Vec<D>>>,
}
impl<D: RegionDataConstraints> std::fmt::Debug for RegionSetLtcs<D> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RegionSetLtcs")
.field(
"nonzero_regions",
&self.nonzero_regions().collect::<Vec<_>>(),
)
.finish()
}
}
impl<D: RegionDataConstraints> RegionSetLtcs<D> {
pub fn empty() -> Self {
Self {
coords: RegionCoordSetLtcs::empty(),
data: vec![],
_region_coords: OnceCell::new(),
}
}
pub fn from_data(coords: RegionCoordSetLtcs, data: Vec<Vec<Vec<D>>>) -> Self {
Self {
coords,
data,
_region_coords: OnceCell::new(),
}
}
pub fn count(&self) -> usize {
self.data
.iter()
.map(|d| {
if d.is_empty() {
0
} else {
debug_assert!(d.iter().all(|i| i.len() == d[0].len()));
d.len() * d[0].len()
}
})
.sum()
}
pub fn regions(&self) -> impl Iterator<Item = Region<D>> + '_ {
self.coords
.region_coords_flat()
.map(|((ia, ix, it), coords)| Region::new(coords, self.data[ia][ix][it].clone()))
}
pub fn rectify(&mut self, other: &mut Self) -> GossipResult<()> {
if self.coords.arq_set != other.coords.arq_set {
return Err(GossipError::ArqSetMismatchForDiff);
}
if self.coords.times > other.coords.times {
std::mem::swap(self, other);
}
let mut len = 0;
for (da, db) in self.data.iter_mut().zip(other.data.iter_mut()) {
for (dda, ddb) in da.iter_mut().zip(db.iter_mut()) {
TelescopingTimes::rectify((&self.coords.times, dda), (&other.coords.times, ddb));
len = dda.len();
}
}
let times = other.coords.times.limit(len as u32);
self.coords.times = times;
other.coords.times = times;
Ok(())
}
pub fn diff(mut self, mut other: Self) -> GossipResult<Vec<Region<D>>> {
self.rectify(&mut other)?;
let regions = self
.regions()
.zip(other.regions())
.filter_map(|(a, b)| (a.data != b.data).then_some(a))
.collect();
Ok(regions)
}
pub fn nonzero_regions(
&self,
) -> impl '_ + Iterator<Item = ((usize, usize, usize), RegionCoords, D)> {
self.coords
.region_coords_flat()
.filter_map(|((a, x, y), c)| {
let d = &self
.data
.get(a)
.and_then(|d| d.get(x))
.and_then(|d| d.get(y));
d.filter(|d| d.count() > 0)
.map(|d| ((a, x, y), c, d.clone()))
})
}
pub fn data(&self) -> &[Vec<Vec<D>>] {
self.data.as_ref()
}
}
#[cfg(feature = "test_utils")]
impl<D: RegionDataConstraints> RegionSetLtcs<D> {
pub fn from_store<O: OpRegion<D>, S: crate::persistence::AccessOpStore<O, D>>(
store: &S,
coords: RegionCoordSetLtcs,
) -> Self {
coords.into_region_set_infallible(|(_, coords)| store.query_region_data(&coords))
}
}