1use std::cmp::Ordering;
15use std::collections::HashMap;
16use std::collections::hash_map::Entry;
17use std::hash::Hash;
18use std::marker::PhantomData;
19use std::mem;
20use std::mem::MaybeUninit;
21use std::ops::Range;
22use std::ptr;
23use std::ptr::NonNull;
24
25use futures::Future;
26use ndslice::Slice;
27use ndslice::extent;
28use ndslice::view;
29use ndslice::view::Ranked;
30use ndslice::view::Region;
31use serde::Deserialize;
32use serde::Serialize;
33use typeuri::Named;
34
35use crate::accum::Accumulator;
36use crate::accum::CommReducer;
37use crate::accum::ReducerSpec;
38
39pub mod rle;
41mod value_overlay;
42pub use value_overlay::BuildError;
43pub use value_overlay::ValueOverlay;
44
45#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
47pub enum ValueMeshError {
48 #[error("invalid value mesh: expected {expected} ranks, but contains {actual} ranks")]
50 InvalidRankCardinality {
51 expected: usize,
53 actual: usize,
55 },
56}
57
58impl From<view::InvalidCardinality> for ValueMeshError {
59 fn from(value: view::InvalidCardinality) -> Self {
60 Self::InvalidRankCardinality {
61 expected: value.expected,
62 actual: value.actual,
63 }
64 }
65}
66
67#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct ValueMesh<T> {
76 region: Region,
81
82 rep: Rep<T>,
93}
94
95#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
106struct Run {
107 start: u64,
109 end: u64,
111 id: u32,
113}
114
115impl Run {
116 fn new(start: usize, end: usize, id: u32) -> Self {
121 Self {
122 start: start as u64,
123 end: end as u64,
124 id,
125 }
126 }
127}
128
129impl TryFrom<Run> for (Range<usize>, u32) {
130 type Error = &'static str;
131
132 #[allow(clippy::result_large_err)]
141 fn try_from(r: Run) -> Result<Self, Self::Error> {
142 let start = usize::try_from(r.start).map_err(|_| "run.start too large")?;
143 let end = usize::try_from(r.end).map_err(|_| "run.end too large")?;
144 Ok((start..end, r.id))
145 }
146}
147
148#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] enum Rep<T> {
163 Dense {
170 values: Vec<T>,
173 },
174
175 Compressed {
187 table: Vec<T>,
189
190 runs: Vec<Run>,
193 },
194}
195
196impl<T> Default for ValueMesh<T> {
201 fn default() -> Self {
202 Self::empty()
203 }
204}
205
206impl<T> ValueMesh<T> {
207 pub fn empty() -> Self {
214 let region = extent!(r = 0).into();
216 Self::new_unchecked(region, Vec::<T>::new())
217 }
218
219 #[allow(clippy::result_large_err)]
232 pub fn new(region: Region, ranks: Vec<T>) -> std::result::Result<Self, ValueMeshError> {
233 let (actual, expected) = (ranks.len(), region.num_ranks());
234 if actual != expected {
235 return Err(ValueMeshError::InvalidRankCardinality { expected, actual });
236 }
237 Ok(Self {
238 region,
239 rep: Rep::Dense { values: ranks },
240 })
241 }
242
243 #[inline]
246 pub(crate) fn new_unchecked(region: Region, ranks: Vec<T>) -> Self {
247 debug_assert_eq!(region.num_ranks(), ranks.len());
248 Self {
249 region,
250 rep: Rep::Dense { values: ranks },
251 }
252 }
253}
254
255impl<T: Clone> ValueMesh<T> {
256 pub fn from_single(region: Region, s: T) -> Self {
264 let n = region.num_ranks();
265 if n == 0 {
266 return Self {
267 region,
268 rep: Rep::Compressed {
269 table: Vec::new(),
270 runs: Vec::new(),
271 },
272 };
273 }
274
275 let table = vec![s];
276 let runs = vec![Run::new(0, n, 0)];
277 Self {
278 region,
279 rep: Rep::Compressed { table, runs },
280 }
281 }
282}
283
284impl<T: Clone + Default> ValueMesh<T> {
285 pub fn from_default(region: Region) -> Self {
291 ValueMesh::<T>::from_single(region, T::default())
292 }
293}
294
295impl<T: Eq + Hash> ValueMesh<T> {
296 #[allow(clippy::result_large_err)]
305 pub fn from_ranges_with_default(
306 region: Region,
307 default: T,
308 mut ranges: Vec<(Range<usize>, T)>,
309 ) -> std::result::Result<Self, ValueMeshError> {
310 let n = region.num_ranks();
311
312 if n == 0 {
313 return Ok(Self {
314 region,
315 rep: Rep::Compressed {
316 table: Vec::new(),
317 runs: Vec::new(),
318 },
319 });
320 }
321
322 for (r, _) in &ranges {
324 if r.is_empty() {
325 return Err(ValueMeshError::InvalidRankCardinality {
326 expected: n,
327 actual: 0,
328 }); }
330 if r.end > n {
331 return Err(ValueMeshError::InvalidRankCardinality {
332 expected: n,
333 actual: r.end,
334 });
335 }
336 }
337 ranges.sort_by_key(|(r, _)| (r.start, r.end));
338
339 for w in ranges.windows(2) {
341 let (a, _) = &w[0];
342 let (b, _) = &w[1];
343 if a.end > b.start {
344 return Err(ValueMeshError::InvalidRankCardinality {
346 expected: n,
347 actual: b.start, });
349 }
350 }
351
352 let mut index: HashMap<T, u32> = HashMap::with_capacity(1 + ranges.len());
354 let mut next_id: u32 = 0;
355
356 let id_of = |v: T, index: &mut HashMap<T, u32>, next_id: &mut u32| -> u32 {
359 match index.entry(v) {
360 Entry::Occupied(o) => *o.get(),
361 Entry::Vacant(vac) => {
362 let id = *next_id;
363 *next_id += 1;
364 vac.insert(id);
365 id
366 }
367 }
368 };
369
370 let default_id = id_of(default, &mut index, &mut next_id);
371
372 let mut runs: Vec<Run> = Vec::with_capacity(1 + 2 * ranges.len());
373 let mut cursor = 0usize;
374
375 for (r, v) in ranges.into_iter() {
376 if cursor < r.start {
378 runs.push(Run::new(cursor, r.start, default_id));
379 }
380 let id = id_of(v, &mut index, &mut next_id);
382 runs.push(Run::new(r.start, r.end, id));
383 cursor = r.end;
384 }
385
386 if cursor < n {
388 runs.push(Run::new(cursor, n, default_id));
389 }
390
391 let mut table_slots: Vec<Option<T>> = Vec::new();
394 table_slots.resize_with(next_id as usize, || None);
395
396 for (t, id) in index.into_iter() {
397 table_slots[id as usize] = Some(t);
398 }
399
400 let table: Vec<T> = table_slots
401 .into_iter()
402 .map(|o| o.expect("every id must be assigned"))
403 .collect();
404
405 Ok(Self {
406 region,
407 rep: Rep::Compressed { table, runs },
408 })
409 }
410
411 #[allow(clippy::result_large_err)]
441 pub fn from_dense(region: Region, values: Vec<T>) -> std::result::Result<Self, ValueMeshError> {
442 let mut vm = Self::new(region, values)?;
443 vm.compress_adjacent_in_place();
444 Ok(vm)
445 }
446}
447
448impl<F: Future> ValueMesh<F> {
449 pub async fn join(self) -> ValueMesh<F::Output> {
452 let ValueMesh { region, rep } = self;
453
454 match rep {
455 Rep::Dense { values } => {
456 let results = futures::future::join_all(values).await;
457 ValueMesh::new_unchecked(region, results)
458 }
459 Rep::Compressed { .. } => {
460 unreachable!("join() not implemented for compressed meshes")
461 }
462 }
463 }
464}
465
466impl<T, E> ValueMesh<Result<T, E>> {
467 pub fn transpose(self) -> Result<ValueMesh<T>, E> {
470 let ValueMesh { region, rep } = self;
471
472 match rep {
473 Rep::Dense { values } => {
474 let values = values.into_iter().collect::<Result<Vec<T>, E>>()?;
475 Ok(ValueMesh::new_unchecked(region, values))
476 }
477 Rep::Compressed { table, runs } => {
478 let table: Vec<T> = table.into_iter().collect::<Result<Vec<T>, E>>()?;
479 Ok(ValueMesh {
480 region,
481 rep: Rep::Compressed { table, runs },
482 })
483 }
484 }
485 }
486}
487
488impl<T: 'static> view::Ranked for ValueMesh<T> {
489 type Item = T;
490
491 fn region(&self) -> &Region {
494 &self.region
495 }
496
497 fn get(&self, rank: usize) -> Option<&Self::Item> {
509 if rank >= self.region.num_ranks() {
510 return None;
511 }
512
513 match &self.rep {
514 Rep::Dense { values } => values.get(rank),
515
516 Rep::Compressed { table, runs } => {
517 let rank = rank as u64;
518
519 let idx = runs
522 .binary_search_by(|run| {
523 if run.end <= rank {
524 Ordering::Less
525 } else if run.start > rank {
526 Ordering::Greater
527 } else {
528 Ordering::Equal
529 }
530 })
531 .ok()?;
532
533 let id = runs[idx].id as usize;
535 table.get(id)
536 }
537 }
538 }
539}
540
541impl<T: 'static> ValueMesh<T> {
542 pub fn get_by_base_rank(&self, base_rank: usize) -> Option<&T> {
555 if !self.region.slice().contains(base_rank) {
556 return None;
557 }
558 let ordinal = self.region.point_of_base_rank(base_rank).ok()?.rank();
559 self.get(ordinal)
560 }
561}
562
563enum SourceOrdinalMapping {
569 Affine(Slice),
570 BaseSpace { selected: Slice, source: Slice },
571}
572
573impl SourceOrdinalMapping {
574 fn new(selected: &Slice, source: &Slice) -> Self {
581 if let Some(relative) = selected.relative_ordinal_slice(source) {
586 Self::Affine(relative)
587 } else {
588 Self::BaseSpace {
589 selected: selected.clone(),
590 source: source.clone(),
591 }
592 }
593 }
594
595 fn source_ordinal_at(&self, selected_ordinal: usize) -> usize {
601 match self {
603 Self::Affine(source_ordinals) => source_ordinals
605 .get(selected_ordinal)
606 .expect("selected ordinal should belong to the relative slice"),
607 Self::BaseSpace { selected, source } => {
608 let base_rank = selected
611 .get(selected_ordinal)
612 .expect("selected ordinal should belong to the selected slice");
613
614 source
615 .index(base_rank)
616 .expect("selected base rank should belong to the source slice")
617 }
618 }
619 }
620
621 fn first_selected_ordinal_after(&self, source_ordinal: usize) -> usize {
626 match self {
629 Self::Affine(source_ordinals) => {
630 if source_ordinal < source_ordinals.offset() {
632 return 0;
633 }
634
635 let mut distance = source_ordinal - source_ordinals.offset();
636 let mut selected_stride = source_ordinals.len();
637 let mut selected_ordinal = 0;
638
639 for (&dim_size, &source_stride) in source_ordinals
640 .sizes()
641 .iter()
642 .zip(source_ordinals.strides())
643 {
644 selected_stride /= dim_size;
645 let coordinate = (distance / source_stride).min(dim_size - 1);
646 selected_ordinal += coordinate * selected_stride;
647 distance -= coordinate * source_stride;
648 }
649
650 selected_ordinal + 1
651 }
652 Self::BaseSpace { selected, .. } => {
653 let mut lower_cursor = 0;
656 let mut upper_cursor = selected.len();
657
658 while lower_cursor < upper_cursor {
659 let middle = lower_cursor + (upper_cursor - lower_cursor) / 2;
660
661 if self.source_ordinal_at(middle) <= source_ordinal {
662 lower_cursor = middle + 1;
663 } else {
664 upper_cursor = middle;
665 }
666 }
667
668 lower_cursor
669 }
670 }
671 }
672}
673
674impl<T: Clone + 'static> view::RankedSliceable for ValueMesh<T> {
675 fn sliced(&self, region: Region) -> Self {
676 debug_assert!(region.is_subset(self.region()), "sliced: not a subset");
677
678 self.sliced_unchecked(region)
679 }
680}
681
682impl<T: Clone + 'static> ValueMesh<T> {
683 pub fn sliced_unchecked(&self, region: Region) -> Self {
693 match &self.rep {
694 Rep::Dense { .. } => self.slice_dense(region),
695 Rep::Compressed { .. } if region.num_ranks() == 0 => Self {
696 region,
697 rep: Rep::Compressed {
698 table: Vec::new(),
699 runs: Vec::new(),
700 },
701 },
702 Rep::Compressed { table, runs } => {
703 let source_ordinals =
704 SourceOrdinalMapping::new(region.slice(), self.region.slice());
705
706 self.slice_compressed(region, table, runs, source_ordinals)
707 }
708 }
709 }
710}
711
712impl<T: Clone + 'static> ValueMesh<T> {
713 fn slice_dense(&self, region: Region) -> Self {
714 let ranks: Vec<T> = self
715 .region()
716 .remap(®ion)
717 .unwrap()
718 .map(|index| self.get(index).unwrap().clone())
719 .collect();
720 debug_assert_eq!(
721 region.num_ranks(),
722 ranks.len(),
723 "sliced: cardinality mismatch"
724 );
725 Self::new_unchecked(region, ranks)
726 }
727
728 fn slice_compressed(
753 &self,
754 slice_region: Region,
755 value_table: &[T],
756 source_runs: &[Run],
757 source_ordinals: SourceOrdinalMapping,
758 ) -> Self {
759 let mut slice_runs: Vec<Run> = Vec::new();
760 let mut slice_ordinal_cursor = 0usize;
761 let mut source_ordinal_cursor = source_ordinals.source_ordinal_at(0);
762 let mut source_run_index = 0usize;
763
764 while slice_ordinal_cursor < slice_region.num_ranks() {
765 while source_runs
768 .get(source_run_index)
769 .is_some_and(|run| run.end as usize <= source_ordinal_cursor)
770 {
771 source_run_index += 1;
772 }
773
774 let source_run = &source_runs[source_run_index];
775 debug_assert!(source_run.start as usize <= source_ordinal_cursor);
776
777 let next_slice_ordinal =
778 source_ordinals.first_selected_ordinal_after(source_run.end as usize - 1);
779
780 if let Some(previous_run) = slice_runs.last_mut()
781 && previous_run.id == source_run.id
782 {
783 previous_run.end = next_slice_ordinal as u64;
788 } else {
789 slice_runs.push(Run::new(
791 slice_ordinal_cursor,
792 next_slice_ordinal,
793 source_run.id,
794 ));
795 }
796
797 slice_ordinal_cursor = next_slice_ordinal;
798
799 if slice_ordinal_cursor < slice_region.num_ranks() {
800 source_ordinal_cursor = source_ordinals.source_ordinal_at(slice_ordinal_cursor);
801 }
802 }
803
804 Self {
805 region: slice_region,
806 rep: Rep::Compressed {
807 table: value_table.to_vec(),
808 runs: slice_runs,
809 },
810 }
811 }
812}
813
814impl<T> view::BuildFromRegion<T> for ValueMesh<T> {
815 type Error = ValueMeshError;
816
817 fn build_dense(region: Region, values: Vec<T>) -> Result<Self, Self::Error> {
818 Self::new(region, values)
819 }
820
821 fn build_dense_unchecked(region: Region, values: Vec<T>) -> Self {
822 Self::new_unchecked(region, values)
823 }
824}
825
826impl<T> view::BuildFromRegionIndexed<T> for ValueMesh<T> {
827 type Error = ValueMeshError;
828
829 fn build_indexed(
830 region: Region,
831 pairs: impl IntoIterator<Item = (usize, T)>,
832 ) -> Result<Self, Self::Error> {
833 let n = region.num_ranks();
834
835 let mut buf: Vec<MaybeUninit<T>> = Vec::with_capacity(n);
841 unsafe {
847 buf.set_len(n);
848 }
849
850 let words = n.div_ceil(64);
852 let mut bits = vec![0u64; words];
853 let mut filled = 0usize;
854
855 struct DropGuard<T> {
861 buf: NonNull<MaybeUninit<T>>,
862 bits: NonNull<u64>,
863 n_elems: usize,
864 n_words: usize,
865 disarm: bool,
866 }
867
868 impl<T> DropGuard<T> {
869 unsafe fn new(buf: &mut [MaybeUninit<T>], bits: &mut [u64]) -> Self {
888 let n_elems = buf.len();
889 let n_words = bits.len();
890 Self {
893 buf: NonNull::new(buf.as_mut_ptr()).unwrap_or_else(NonNull::dangling),
894 bits: NonNull::new(bits.as_mut_ptr()).unwrap_or_else(NonNull::dangling),
895 n_elems,
896 n_words,
897 disarm: false,
898 }
899 }
900
901 #[inline]
902 fn disarm(&mut self) {
903 self.disarm = true;
904 }
905 }
906
907 impl<T> Drop for DropGuard<T> {
908 fn drop(&mut self) {
909 if self.disarm {
910 return;
911 }
912
913 unsafe {
931 let buf_base = self.buf.as_ptr();
932 let bits_base = self.bits.as_ptr();
933
934 for w in 0..self.n_words {
935 let mut word = *bits_base.add(w);
937
938 if w == self.n_words.saturating_sub(1) {
941 let used_bits = self.n_elems.saturating_sub(w * 64);
942 if used_bits < 64 {
943 let mask = if used_bits == 0 {
944 0
945 } else {
946 (1u64 << used_bits) - 1
947 };
948 word &= mask;
949 }
950 }
951
952 while word != 0 {
954 let tz = word.trailing_zeros() as usize;
955 let i = w * 64 + tz;
956 debug_assert!(i < self.n_elems);
957
958 let slot = buf_base.add(i);
959 ptr::drop_in_place((*slot).as_mut_ptr());
961
962 word &= word - 1;
964 }
965 }
966 }
967 }
968 }
969
970 let mut guard = unsafe { DropGuard::new(&mut buf, &mut bits) };
979
980 for (rank, value) in pairs {
981 if rank >= guard.n_elems {
983 return Err(ValueMeshError::InvalidRankCardinality {
984 expected: guard.n_elems,
985 actual: rank + 1,
986 });
987 }
988
989 let w = rank / 64;
991 let b = rank % 64;
992 let mask = 1u64 << b;
993
994 unsafe {
1013 let bits_ptr = guard.bits.as_ptr().add(w);
1015 let buf_slot = guard.buf.as_ptr().add(rank);
1016
1017 let word = *bits_ptr;
1019
1020 if (word & mask) != 0 {
1021 core::ptr::drop_in_place((*buf_slot).as_mut_ptr());
1023 } else {
1026 *bits_ptr = word | mask;
1028 filled += 1;
1029 }
1030
1031 (*buf_slot).write(value);
1033 }
1034 }
1035
1036 if filled != n {
1037 return Err(ValueMeshError::InvalidRankCardinality {
1039 expected: n,
1040 actual: filled,
1041 });
1042 }
1043
1044 guard.disarm();
1046
1047 let ranks = unsafe {
1049 let ptr = buf.as_mut_ptr() as *mut T;
1050 let len = buf.len();
1051 let cap = buf.capacity();
1052 mem::forget(buf);
1057 Vec::from_raw_parts(ptr, len, cap)
1058 };
1059
1060 Ok(Self::new_unchecked(region, ranks))
1061 }
1062}
1063
1064impl<T: PartialEq> ValueMesh<T> {
1065 pub fn compress_adjacent_in_place(&mut self) {
1085 self.compress_adjacent_in_place_by(|a, b| a == b)
1086 }
1087}
1088
1089impl<T: Clone + PartialEq> ValueMesh<T> {
1090 fn materialized_runs(&self) -> Vec<(Range<usize>, T)> {
1106 match &self.rep {
1107 Rep::Dense { values } => {
1108 let mut out = Vec::new();
1110 if values.is_empty() {
1111 return out;
1112 }
1113 let mut start = 0usize;
1114 for i in 1..values.len() {
1115 if values[i] != values[i - 1] {
1116 out.push((start..i, values[i - 1].clone()));
1117 start = i;
1118 }
1119 }
1120 out.push((start..values.len(), values.last().unwrap().clone()));
1121 out
1122 }
1123 Rep::Compressed { table, runs } => runs
1124 .iter()
1125 .map(|r| {
1126 let id = r.id as usize;
1127 ((r.start as usize..r.end as usize), table[id].clone())
1128 })
1129 .collect(),
1130 }
1131 }
1132}
1133
1134impl<T: Clone + Eq> ValueMesh<T> {
1135 pub fn merge_from_overlay(&mut self, overlay: ValueOverlay<T>) -> Result<(), BuildError> {
1141 let n = self.region.num_ranks();
1142
1143 for (r, _) in overlay.runs() {
1146 if r.end > n {
1147 return Err(BuildError::OutOfBounds {
1148 range: r.clone(),
1149 region_len: n,
1150 });
1151 }
1152 }
1153
1154 let left = self.materialized_runs();
1156 let right: Vec<(std::ops::Range<usize>, T)> = overlay.runs().cloned().collect();
1159
1160 let merged = rle::merge_value_runs(left, right);
1163
1164 let (table, raw_runs) = rle::rle_from_value_runs(merged);
1166 let runs = raw_runs
1167 .into_iter()
1168 .map(|(r, id)| Run::new(r.start, r.end, id))
1169 .collect();
1170 self.rep = Rep::Compressed { table, runs };
1171
1172 Ok(())
1173 }
1174}
1175
1176impl<T> ValueMesh<T> {
1177 pub fn compress_adjacent_in_place_by<F>(&mut self, same: F)
1199 where
1200 F: FnMut(&T, &T) -> bool,
1201 {
1202 let values = match &mut self.rep {
1203 Rep::Dense { values } => std::mem::take(values),
1204 Rep::Compressed { .. } => return,
1205 };
1206 let (table, raw_runs) = rle::rle_from_dense(values, same);
1207 let runs = raw_runs
1208 .into_iter()
1209 .map(|(r, id)| Run::new(r.start, r.end, id))
1210 .collect();
1211 self.rep = Rep::Compressed { table, runs };
1212 }
1213}
1214
1215impl<T> Accumulator for ValueMesh<T>
1234where
1235 T: Eq + Clone + Named,
1236{
1237 type State = Self;
1238 type Update = ValueOverlay<T>;
1239
1240 fn accumulate(&self, state: &mut Self::State, update: Self::Update) -> anyhow::Result<()> {
1241 if state.region().num_ranks() == 0 {
1246 *state = self.clone();
1247 }
1248
1249 state.merge_from_overlay(update)?;
1252 Ok(())
1253 }
1254
1255 fn reducer_spec(&self) -> Option<ReducerSpec> {
1256 Some(ReducerSpec {
1257 typehash: <ValueOverlayReducer<T> as Named>::typehash(),
1258 builder_params: None,
1259 })
1260 }
1261}
1262
1263#[derive(Named)]
1270pub struct ValueOverlayReducer<T>(PhantomData<T>);
1271
1272impl<T> ValueOverlayReducer<T> {
1273 pub fn new() -> Self {
1275 Self(PhantomData)
1276 }
1277}
1278
1279impl<T> CommReducer for ValueOverlayReducer<T>
1289where
1290 T: Eq + Clone + Named,
1291{
1292 type Update = ValueOverlay<T>;
1293
1294 fn reduce(&self, left: Self::Update, right: Self::Update) -> anyhow::Result<Self::Update> {
1296 let merged = rle::merge_value_runs(
1297 left.runs().cloned().collect(),
1298 right.runs().cloned().collect(),
1299 );
1300 Ok(ValueOverlay::try_from_runs(merged)?)
1301 }
1302}
1303
1304#[cfg(test)]
1305mod tests {
1306 use std::convert::Infallible;
1307 use std::future::Future;
1308 use std::pin::Pin;
1309 use std::task::Context;
1310 use std::task::Poll;
1311 use std::task::RawWaker;
1312 use std::task::RawWakerVTable;
1313 use std::task::Waker;
1314
1315 use futures::executor::block_on;
1316 use futures::future;
1317 use ndslice::extent;
1318 use ndslice::strategy::gen_region;
1319 use ndslice::view::CollectExactMeshExt;
1320 use ndslice::view::CollectIndexedMeshExt;
1321 use ndslice::view::CollectMeshExt;
1322 use ndslice::view::MapIntoExt;
1323 use ndslice::view::Ranked;
1324 use ndslice::view::RankedSliceable;
1325 use ndslice::view::ViewExt;
1326 use proptest::prelude::*;
1327 use proptest::strategy::ValueTree;
1328 use serde::Deserialize;
1329 use serde::Serialize;
1330 use serde_json;
1331 use typeuri::Named;
1332
1333 use super::*;
1334
1335 fn rectangular_slicing_case() -> impl Strategy<Value = (Region, Region, Vec<u8>)> {
1337 prop::collection::vec(1usize..=5, 1..=4)
1340 .prop_filter("rank count must stay small", |sizes| {
1342 sizes.iter().product::<usize>() <= 256
1343 })
1344 .prop_flat_map(|sizes| {
1347 let num_dims = sizes.len();
1348 let num_ranks: usize = sizes.iter().product();
1349 (
1350 Just(sizes),
1351 0usize..=16,
1353 1usize..=4,
1355 prop::collection::vec((0usize..=32, 0usize..=32, 0usize..=32), num_dims),
1357 prop::collection::vec(0u8..=3, num_ranks),
1359 )
1360 })
1361 .prop_map(|(sizes, offset, scale, ranges, values)| {
1364 let labels: Vec<_> = (0..sizes.len()).map(|dim| format!("d{dim}")).collect();
1365 let mut strides = vec![0; sizes.len()];
1366 let mut stride = scale;
1367 for dim in (0..sizes.len()).rev() {
1368 strides[dim] = stride;
1369 stride *= sizes[dim];
1370 }
1371
1372 let region = Region::new(
1373 labels.clone(),
1374 Slice::new(offset, sizes.clone(), strides)
1375 .expect("generated parent slice should be valid"),
1376 );
1377 let selected = ranges.into_iter().enumerate().fold(
1378 region.clone(),
1379 |selected, (dim, (begin_seed, end_seed, step_seed))| {
1380 let begin = begin_seed % sizes[dim];
1381 let end = begin + 1 + end_seed % (sizes[dim] - begin);
1382 let valid_steps: Vec<_> = (1..=sizes[dim])
1383 .filter(|step| sizes[dim].is_multiple_of(*step))
1384 .collect();
1385 let step = valid_steps[step_seed % valid_steps.len()];
1386 selected
1387 .range(&labels[dim], ndslice::Range(begin, Some(end), step))
1388 .expect("generated range should be valid")
1389 },
1390 );
1391
1392 (region, selected, values)
1393 })
1394 }
1395
1396 proptest! {
1397 #![proptest_config(ProptestConfig::with_cases(512))]
1398
1399 #[test]
1400 fn compressed_slicing_matches_dense_slicing(
1401 (region, selected, values) in rectangular_slicing_case()
1402 ) {
1403 let dense = ValueMesh::new_unchecked(region, values);
1404 let mut compressed = dense.clone();
1405 compressed.compress_adjacent_in_place();
1406
1407 let expected = dense.sliced(selected.clone());
1408 let actual = compressed.sliced(selected);
1409
1410 prop_assert_eq!(actual.region(), expected.region());
1411 prop_assert_eq!(
1412 actual.values().collect::<Vec<_>>(),
1413 expected.values().collect::<Vec<_>>(),
1414 );
1415 prop_assert!(
1416 matches!(actual.rep, Rep::Compressed { .. }),
1417 "the supported rectangular selection should stay compressed",
1418 );
1419 }
1420 }
1421
1422 #[test]
1423 fn compressed_slicing_supports_reshape() {
1424 let matrix: Region = extent!(host = 2, gpu = 2).into();
1425 let mesh = ValueMesh::from_dense(matrix.clone(), vec![0, 0, 1, 1])
1426 .expect("values should match the source region");
1427 let flat = Region::new(vec!["rank".to_string()], Slice::new_row_major(vec![4]));
1428
1429 let flattened = mesh.sliced(flat);
1430
1431 assert_eq!(flattened.values().collect::<Vec<_>>(), vec![0, 0, 1, 1]);
1432 assert!(matches!(flattened.rep, Rep::Compressed { .. }));
1433
1434 let reshaped = flattened.sliced(matrix);
1435
1436 assert_eq!(reshaped.values().collect::<Vec<_>>(), vec![0, 0, 1, 1]);
1437 assert!(matches!(reshaped.rep, Rep::Compressed { .. }));
1438 }
1439
1440 #[test]
1441 fn compressed_slicing_matches_dense_for_dimension_changing_subset() {
1442 let parent = Region::new(vec!["rank".to_string()], Slice::new_row_major(vec![8]));
1446
1447 let dense = ValueMesh::new(parent, vec![0, 0, 1, 1, 2, 2, 3, 3])
1448 .expect("values should match the parent region");
1449
1450 let mut compressed = dense.clone();
1451
1452 compressed.compress_adjacent_in_place();
1453
1454 let selected = Region::new(
1455 vec!["row".to_string(), "column".to_string()],
1456 Slice::new(2, vec![2, 2], vec![2, 1]).expect("selection should be valid"),
1457 );
1458
1459 assert!(selected.is_subset(dense.region()));
1460
1461 let expected = dense.sliced(selected.clone());
1462 let actual = compressed.sliced(selected);
1463
1464 assert_eq!(expected.values().collect::<Vec<_>>(), vec![1, 1, 2, 2]);
1465
1466 assert_eq!(actual.region(), expected.region());
1467
1468 assert_eq!(
1469 actual.values().collect::<Vec<_>>(),
1470 expected.values().collect::<Vec<_>>()
1471 );
1472
1473 assert!(matches!(actual.rep, Rep::Compressed { .. }));
1474 }
1475
1476 #[test]
1477 fn compressed_slicing_matches_dense_for_reshaped_proper_subset() {
1478 let parent = Region::new(
1483 vec!["row".to_string(), "column".to_string()],
1484 Slice::new(0, vec![2, 6], vec![8, 1]).expect("parent should be valid"),
1485 );
1486 let dense = ValueMesh::new(parent, vec![0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5])
1487 .expect("values should match the parent region");
1488 let mut compressed = dense.clone();
1489 compressed.compress_adjacent_in_place();
1490 let selected = Region::new(
1491 vec!["outer".to_string(), "row".to_string(), "column".to_string()],
1492 Slice::new(1, vec![2, 2, 2], vec![8, 2, 1]).expect("selection should be valid"),
1493 );
1494
1495 let expected = dense.sliced(selected.clone());
1496 let actual = compressed.sliced(selected);
1497
1498 assert_eq!(
1499 expected.values().collect::<Vec<_>>(),
1500 vec![0, 1, 1, 2, 3, 4, 4, 5]
1501 );
1502 assert_eq!(actual.region(), expected.region());
1503 assert_eq!(
1504 actual.values().collect::<Vec<_>>(),
1505 expected.values().collect::<Vec<_>>()
1506 );
1507 assert!(matches!(actual.rep, Rep::Compressed { .. }));
1508 }
1509
1510 #[test]
1511 fn value_mesh_new_ok() {
1512 let region: Region = extent!(replica = 2, gpu = 3).into();
1513 let mesh = ValueMesh::new(region.clone(), (0..6).collect()).expect("new should succeed");
1514 assert_eq!(mesh.region().num_ranks(), 6);
1515 assert_eq!(mesh.values().count(), 6);
1516 assert_eq!(mesh.values().collect::<Vec<_>>(), vec![0, 1, 2, 3, 4, 5]);
1517 }
1518
1519 #[test]
1520 fn get_by_base_rank_resolves_through_region_geometry() {
1521 let region: Region = extent!(row = 4, col = 2).into();
1524 let mesh: ValueMesh<i32> =
1525 ValueMesh::new(region.clone(), (0..8).map(|i| i * 10).collect()).unwrap();
1526 assert_eq!(mesh.get_by_base_rank(0), Some(&0));
1527 assert_eq!(mesh.get_by_base_rank(5), Some(&50));
1528 assert_eq!(mesh.get_by_base_rank(8), None);
1529
1530 let sub = mesh.sliced(region.range("row", ndslice::Range(1, Some(3), 1)).unwrap());
1533
1534 assert_eq!(sub.get_by_base_rank(2), Some(&20)); assert_eq!(sub.get_by_base_rank(5), Some(&50)); assert_eq!(sub.get_by_base_rank(0), None); assert_eq!(sub.get_by_base_rank(6), None); assert_eq!(sub.get(0), Some(&20)); }
1543
1544 #[test]
1545 fn value_mesh_new_len_mismatch_is_error() {
1546 let region: Region = extent!(replica = 2, gpu = 3).into();
1547 let err = ValueMesh::new(region, vec![0_i32; 5]).unwrap_err();
1548 let ValueMeshError::InvalidRankCardinality { expected, actual } = err;
1549 assert_eq!(expected, 6);
1550 assert_eq!(actual, 5);
1551 }
1552
1553 #[test]
1554 fn value_mesh_transpose_ok_and_err() {
1555 let region: Region = extent!(x = 2).into();
1556
1557 let ok_mesh = ValueMesh::new(region.clone(), vec![Ok::<_, Infallible>(1), Ok(2)]).unwrap();
1559 let ok = ok_mesh.transpose().unwrap();
1560 assert_eq!(ok.values().collect::<Vec<_>>(), vec![1, 2]);
1561
1562 #[derive(Debug, PartialEq)]
1564 enum E {
1565 Boom,
1566 }
1567 let err_mesh = ValueMesh::new(region, vec![Ok(1), Err(E::Boom)]).unwrap();
1568 let err = err_mesh.transpose().unwrap_err();
1569 assert_eq!(err, E::Boom);
1570 }
1571
1572 #[test]
1573 fn value_mesh_join_preserves_region_and_values() {
1574 let region: Region = extent!(x = 2, y = 2).into();
1575 let futs = vec![
1576 future::ready(10),
1577 future::ready(11),
1578 future::ready(12),
1579 future::ready(13),
1580 ];
1581 let mesh = ValueMesh::new(region.clone(), futs).unwrap();
1582
1583 let joined = block_on(mesh.join());
1584 assert_eq!(joined.region().num_ranks(), 4);
1585 assert_eq!(joined.values().collect::<Vec<_>>(), vec![10, 11, 12, 13]);
1586 }
1587
1588 #[test]
1589 fn collect_mesh_ok() {
1590 let region: Region = extent!(x = 2, y = 3).into();
1591 let mesh = (0..6)
1592 .collect_mesh::<ValueMesh<_>>(region.clone())
1593 .expect("collect_mesh should succeed");
1594
1595 assert_eq!(mesh.region().num_ranks(), 6);
1596 assert_eq!(mesh.values().collect::<Vec<_>>(), vec![0, 1, 2, 3, 4, 5]);
1597 }
1598
1599 #[test]
1600 fn collect_mesh_len_too_short_is_error() {
1601 let region: Region = extent!(x = 2, y = 3).into();
1602 let err = (0..5).collect_mesh::<ValueMesh<_>>(region).unwrap_err();
1603
1604 let ValueMeshError::InvalidRankCardinality { expected, actual } = err;
1605 assert_eq!(expected, 6);
1606 assert_eq!(actual, 5);
1607 }
1608
1609 #[test]
1610 fn collect_mesh_len_too_long_is_error() {
1611 let region: Region = extent!(x = 2, y = 3).into();
1612 let err = (0..7).collect_mesh::<ValueMesh<_>>(region).unwrap_err();
1613 let ValueMeshError::InvalidRankCardinality { expected, actual } = err;
1614 assert_eq!(expected, 6);
1615 assert_eq!(actual, 7);
1616 }
1617
1618 #[test]
1619 fn collect_mesh_from_map_pipeline() {
1620 let region: Region = extent!(x = 2, y = 2).into();
1621 let mesh = (0..4)
1622 .map(|i| i * 10)
1623 .collect_mesh::<ValueMesh<_>>(region.clone())
1624 .unwrap();
1625
1626 assert_eq!(mesh.region().num_ranks(), 4);
1627 assert_eq!(mesh.values().collect::<Vec<_>>(), vec![0, 10, 20, 30]);
1628 }
1629
1630 #[test]
1631 fn collect_exact_mesh_ok() {
1632 let region: Region = extent!(x = 2, y = 3).into();
1633 let mesh = (0..6)
1634 .collect_exact_mesh::<ValueMesh<_>>(region.clone())
1635 .expect("collect_exact_mesh should succeed");
1636
1637 assert_eq!(mesh.region().num_ranks(), 6);
1638 assert_eq!(mesh.values().collect::<Vec<_>>(), vec![0, 1, 2, 3, 4, 5]);
1639 }
1640
1641 #[test]
1642 fn collect_exact_mesh_len_too_short_is_error() {
1643 let region: Region = extent!(x = 2, y = 3).into();
1644 let err = (0..5)
1645 .collect_exact_mesh::<ValueMesh<_>>(region)
1646 .unwrap_err();
1647
1648 let ValueMeshError::InvalidRankCardinality { expected, actual } = err;
1649 assert_eq!(expected, 6);
1650 assert_eq!(actual, 5);
1651 }
1652
1653 #[test]
1654 fn collect_exact_mesh_len_too_long_is_error() {
1655 let region: Region = extent!(x = 2, y = 3).into();
1656 let err = (0..7)
1657 .collect_exact_mesh::<ValueMesh<_>>(region)
1658 .unwrap_err();
1659
1660 let ValueMeshError::InvalidRankCardinality { expected, actual } = err;
1661 assert_eq!(expected, 6);
1662 assert_eq!(actual, 7);
1663 }
1664
1665 #[test]
1666 fn collect_exact_mesh_from_map_pipeline() {
1667 let region: Region = extent!(x = 2, y = 2).into();
1668 let mesh = (0..4)
1669 .map(|i| i * 10)
1670 .collect_exact_mesh::<ValueMesh<_>>(region.clone())
1671 .unwrap();
1672
1673 assert_eq!(mesh.region().num_ranks(), 4);
1674 assert_eq!(mesh.values().collect::<Vec<_>>(), vec![0, 10, 20, 30]);
1675 }
1676
1677 #[test]
1678 fn collect_indexed_ok_shuffled() {
1679 let region: Region = extent!(x = 2, y = 3).into();
1680 let pairs = vec![(3, 30), (0, 0), (5, 50), (2, 20), (1, 10), (4, 40)];
1682 let mesh = pairs
1683 .into_iter()
1684 .collect_indexed::<ValueMesh<_>>(region.clone())
1685 .unwrap();
1686
1687 assert_eq!(mesh.region().num_ranks(), 6);
1688 assert_eq!(
1689 mesh.values().collect::<Vec<_>>(),
1690 vec![0, 10, 20, 30, 40, 50]
1691 );
1692 }
1693
1694 #[test]
1695 fn collect_indexed_missing_rank_is_error() {
1696 let region: Region = extent!(x = 2, y = 2).into(); let pairs = vec![(0, 100), (1, 101), (2, 102)];
1699 let err = pairs
1700 .into_iter()
1701 .collect_indexed::<ValueMesh<_>>(region)
1702 .unwrap_err();
1703
1704 let ValueMeshError::InvalidRankCardinality { expected, actual } = err;
1705 assert_eq!(expected, 4);
1706 assert_eq!(actual, 3); }
1708
1709 #[test]
1710 fn collect_indexed_out_of_bounds_is_error() {
1711 let region: Region = extent!(x = 2, y = 2).into(); let pairs = vec![(0, 1), (4, 9)]; let err = pairs
1714 .into_iter()
1715 .collect_indexed::<ValueMesh<_>>(region)
1716 .unwrap_err();
1717
1718 let ValueMeshError::InvalidRankCardinality { expected, actual } = err;
1719 assert_eq!(expected, 4);
1720 assert_eq!(actual, 5); }
1722
1723 #[test]
1724 fn collect_indexed_duplicate_last_write_wins() {
1725 let region: Region = extent!(x = 1, y = 3).into(); let pairs = vec![(0, 7), (1, 8), (1, 88), (2, 9)];
1728 let mesh = pairs
1729 .into_iter()
1730 .collect_indexed::<ValueMesh<_>>(region.clone())
1731 .unwrap();
1732
1733 assert_eq!(mesh.values().collect::<Vec<_>>(), vec![7, 88, 9]);
1734 }
1735
1736 #[allow(clippy::result_large_err)]
1738 fn build_value_mesh_indexed<T>(
1739 region: Region,
1740 pairs: impl IntoIterator<Item = (usize, T)>,
1741 ) -> std::result::Result<ValueMesh<T>, ValueMeshError> {
1742 let n = region.num_ranks();
1743
1744 let mut buf: Vec<Option<T>> = std::iter::repeat_with(|| None).take(n).collect();
1746 let mut filled = 0usize;
1747
1748 for (rank, value) in pairs {
1749 if rank >= n {
1750 return Err(ValueMeshError::InvalidRankCardinality {
1754 expected: n,
1755 actual: rank + 1,
1756 });
1757 }
1758 if buf[rank].is_none() {
1759 filled += 1;
1760 }
1761 buf[rank] = Some(value); }
1763
1764 if filled != n {
1765 return Err(ValueMeshError::InvalidRankCardinality {
1767 expected: n,
1768 actual: filled,
1769 });
1770 }
1771
1772 let ranks: Vec<T> = buf.into_iter().map(Option::unwrap).collect();
1774 Ok(ValueMesh::new_unchecked(region, ranks))
1775 }
1776
1777 fn hash_key(x: usize) -> u64 {
1785 let mut z = x as u64 ^ 0x9E3779B97F4A7C15;
1786 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
1787 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
1788 z ^ (z >> 31)
1789 }
1790
1791 fn pseudo_shuffle<'a, T: 'a>(v: &'a mut [T], key: impl Fn(usize) -> u64 + Copy) {
1808 let mut with_keys: Vec<(u64, usize)> = (0..v.len()).map(|i| (key(i), i)).collect();
1810 with_keys.sort_by_key(|&(k, _)| k);
1811 let perm: Vec<usize> = with_keys.into_iter().map(|(_, i)| i).collect();
1812
1813 let mut seen = vec![false; v.len()];
1816 for i in 0..v.len() {
1817 if seen[i] {
1818 continue;
1819 }
1820 let mut a = i;
1821 while !seen[a] {
1822 seen[a] = true;
1823 let b = perm[a];
1824 if b == i {
1826 break;
1827 }
1828 v.swap(a, b);
1829 a = b;
1830 }
1831 }
1832 }
1833
1834 proptest! {
1850 #[test]
1851 fn try_collect_opt_equivalence(region in gen_region(1..=4, 6), extra_len in 0usize..=12) {
1852 let n = region.num_ranks();
1853
1854 let mut pairs: Vec<(usize, i64)> = (0..n).map(|r| (r, r as i64)).collect();
1856
1857 let extras = proptest::collection::vec(0..n, extra_len)
1860 .new_tree(&mut proptest::test_runner::TestRunner::default())
1861 .unwrap()
1862 .current();
1863 for (k, r) in extras.into_iter().enumerate() {
1864 pairs.push((r, (n as i64) + (k as i64)));
1865 }
1866
1867 pseudo_shuffle(&mut pairs, hash_key);
1870
1871 let mesh_ref = build_value_mesh_indexed(region.clone(), pairs.clone()).unwrap();
1873 let mesh_opt = pairs.into_iter().collect_indexed::<ValueMesh<_>>(region.clone()).unwrap();
1874
1875 prop_assert_eq!(mesh_ref.region(), mesh_opt.region());
1876 prop_assert_eq!(mesh_ref.values().collect::<Vec<_>>(), mesh_opt.values().collect::<Vec<_>>());
1877 }
1878 }
1879
1880 proptest! {
1891 #[test]
1892 fn try_collect_opt_missing_rank_errors_match(region in gen_region(1..=4, 6)) {
1893 let n = region.num_ranks();
1894 let mut pairs: Vec<(usize, i64)> = (0..n).map(|r| (r, r as i64)).collect();
1896 if n > 0 {
1898 let drop_idx = 0usize; pairs.remove(drop_idx);
1900 }
1901 pseudo_shuffle(&mut pairs, hash_key);
1903
1904 let ref_err = build_value_mesh_indexed(region.clone(), pairs.clone()).unwrap_err();
1905 let opt_err = pairs.into_iter().collect_indexed::<ValueMesh<_>>(region).unwrap_err();
1906 assert_eq!(format!("{ref_err:?}"), format!("{opt_err:?}"));
1907 }
1908 }
1909
1910 proptest! {
1922 #[test]
1923 fn try_collect_opt_out_of_bound_errors_match(region in gen_region(1..=4, 6)) {
1924 let n = region.num_ranks();
1925 let mut pairs = vec![(0usize, 0i64), (n, 123i64)];
1927 pseudo_shuffle(&mut pairs, hash_key);
1928
1929 let ref_err = build_value_mesh_indexed(region.clone(), pairs.clone()).unwrap_err();
1930 let opt_err = pairs.into_iter().collect_indexed::<ValueMesh<_>>(region).unwrap_err();
1931 assert_eq!(format!("{ref_err:?}"), format!("{opt_err:?}"));
1932 }
1933 }
1934
1935 #[test]
1936 fn map_into_preserves_region_and_order() {
1937 let region: Region = extent!(rows = 2, cols = 3).into();
1938 let vm = ValueMesh::new_unchecked(region.clone(), vec![0, 1, 2, 3, 4, 5]);
1939
1940 let doubled: ValueMesh<_> = vm.map_into(|x| x * 2);
1941 assert_eq!(doubled.region, region);
1942 assert_eq!(
1943 doubled.values().collect::<Vec<_>>(),
1944 vec![0, 2, 4, 6, 8, 10]
1945 );
1946 }
1947
1948 #[test]
1949 fn map_into_ref_borrows_and_preserves() {
1950 let region: Region = extent!(n = 4).into();
1951 let vm = ValueMesh::new_unchecked(
1952 region.clone(),
1953 vec!["a".to_string(), "b".into(), "c".into(), "d".into()],
1954 );
1955
1956 let lens: ValueMesh<_> = vm.map_into(|s| s.len());
1957 assert_eq!(lens.region, region);
1958 assert_eq!(lens.values().collect::<Vec<_>>(), vec![1, 1, 1, 1]);
1959 }
1960
1961 #[test]
1962 fn try_map_into_short_circuits_on_error() {
1963 let region = extent!(n = 4).into();
1964 let vm = ValueMesh::new_unchecked(region, vec![1, 2, 3, 4]);
1965
1966 let res: Result<ValueMesh<i32>, &'static str> =
1967 vm.try_map_into(|x| if x == &3 { Err("boom") } else { Ok(x + 10) });
1968
1969 assert!(res.is_err());
1970 assert_eq!(res.unwrap_err(), "boom");
1971 }
1972
1973 #[test]
1974 fn try_map_into_ref_short_circuits_on_error() {
1975 let region = extent!(n = 4).into();
1976 let vm = ValueMesh::new_unchecked(region, vec![1, 2, 3, 4]);
1977
1978 let res: Result<ValueMesh<i32>, &'static str> =
1979 vm.try_map_into(|x| if x == &3 { Err("boom") } else { Ok(x + 10) });
1980
1981 assert!(res.is_err());
1982 assert_eq!(res.unwrap_err(), "boom");
1983 }
1984
1985 fn noop_waker() -> Waker {
1987 fn clone(_: *const ()) -> RawWaker {
1988 RawWaker::new(std::ptr::null(), &VTABLE)
1989 }
1990 fn wake(_: *const ()) {}
1991 fn wake_by_ref(_: *const ()) {}
1992 fn drop(_: *const ()) {}
1993 static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);
1994 unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) }
1999 }
2000
2001 fn poll_now<F: Future>(mut fut: F) -> F::Output {
2002 let waker = noop_waker();
2003 let mut cx = Context::from_waker(&waker);
2004 let mut fut = unsafe { Pin::new_unchecked(&mut fut) };
2009 match fut.as_mut().poll(&mut cx) {
2010 Poll::Ready(v) => v,
2011 Poll::Pending => unreachable!("Ready futures must complete immediately"),
2012 }
2013 }
2014 #[test]
2017 fn map_into_ready_futures() {
2018 let region: Region = extent!(r = 2, c = 2).into();
2019 let vm = ValueMesh::new_unchecked(region.clone(), vec![10, 20, 30, 40]);
2020
2021 let pending: ValueMesh<core::future::Ready<_>> =
2023 vm.map_into(|x| core::future::ready(x + 1));
2024 assert_eq!(pending.region, region);
2025
2026 let results: Vec<_> = pending.values().map(|f| poll_now(f.clone())).collect();
2028 assert_eq!(results, vec![11, 21, 31, 41]);
2029 }
2030
2031 #[test]
2032 fn map_into_single_element_mesh() {
2033 let region: Region = extent!(n = 1).into();
2034 let vm = ValueMesh::new_unchecked(region.clone(), vec![7]);
2035
2036 let out: ValueMesh<_> = vm.map_into(|x| x * x);
2037 assert_eq!(out.region, region);
2038 assert_eq!(out.values().collect::<Vec<_>>(), vec![49]);
2039 }
2040
2041 #[test]
2042 fn map_into_ref_with_non_clone_field() {
2043 #[derive(Debug, PartialEq, Eq)]
2045 struct NotClone(i32);
2046
2047 let region: Region = extent!(x = 3).into();
2048 let values = vec![(10, NotClone(1)), (20, NotClone(2)), (30, NotClone(3))];
2049 let mesh: ValueMesh<(i32, NotClone)> =
2050 values.into_iter().collect_mesh(region.clone()).unwrap();
2051
2052 let projected: ValueMesh<i32> = mesh.map_into(|t| t.0);
2053 assert_eq!(projected.values().collect::<Vec<_>>(), vec![10, 20, 30]);
2054 assert_eq!(projected.region(), ®ion);
2055 }
2056
2057 #[test]
2058 fn rle_roundtrip_all_equal() {
2059 let region: Region = extent!(n = 6).into();
2060 let mut vm = ValueMesh::new_unchecked(region.clone(), vec![42; 6]);
2061
2062 vm.compress_adjacent_in_place();
2064 let collected: Vec<_> = vm.values().collect();
2065 assert_eq!(collected, vec![42, 42, 42, 42, 42, 42]);
2066
2067 for i in 0..region.num_ranks() {
2069 assert_eq!(vm.get(i), Some(&42));
2070 }
2071 assert_eq!(vm.get(region.num_ranks()), None); }
2073
2074 #[test]
2075 fn rle_roundtrip_alternating() {
2076 let region: Region = extent!(n = 6).into();
2077 let original = vec![1, 2, 1, 2, 1, 2];
2078 let mut vm = ValueMesh::new_unchecked(region.clone(), original.clone());
2079
2080 vm.compress_adjacent_in_place();
2081 let collected: Vec<_> = vm.values().collect();
2082 assert_eq!(collected, original);
2083
2084 assert_eq!(vm.get(0), Some(&1));
2086 assert_eq!(vm.get(1), Some(&2));
2087 assert_eq!(vm.get(3), Some(&2));
2088 assert_eq!(vm.get(5), Some(&2));
2089 }
2090
2091 #[test]
2092 fn rle_roundtrip_blocky_and_slice() {
2093 let region: Region = extent!(n = 10).into();
2095 let original = vec![0, 0, 0, 1, 1, 2, 2, 2, 2, 3];
2096 let mut vm = ValueMesh::new_unchecked(region.clone(), original.clone());
2097
2098 vm.compress_adjacent_in_place();
2099 let collected: Vec<_> = vm.values().collect();
2100 assert_eq!(collected, original);
2101
2102 let sub_region = region.range("n", 3..8).unwrap();
2104 let sliced = vm.sliced(sub_region);
2105 let sliced_vec: Vec<_> = sliced.values().collect();
2106 assert_eq!(sliced_vec, vec![1, 1, 2, 2, 2]);
2107 }
2108
2109 #[test]
2110 fn rle_idempotent_noop_on_second_call() {
2111 let region: Region = extent!(n = 7).into();
2112 let original = vec![9, 9, 9, 8, 8, 9, 9];
2113 let mut vm = ValueMesh::new_unchecked(region.clone(), original.clone());
2114
2115 vm.compress_adjacent_in_place();
2116 let once: Vec<_> = vm.values().collect();
2117 assert_eq!(once, original);
2118
2119 vm.compress_adjacent_in_place();
2122 let twice: Vec<_> = vm.values().collect();
2123 assert_eq!(twice, original);
2124 }
2125
2126 #[test]
2127 fn rle_works_after_build_indexed() {
2128 let region: Region = extent!(x = 2, y = 3).into(); let pairs = vec![(3, 30), (0, 0), (5, 50), (2, 20), (1, 10), (4, 40)];
2132 let mut vm = pairs
2133 .into_iter()
2134 .collect_indexed::<ValueMesh<_>>(region.clone())
2135 .unwrap();
2136
2137 vm.compress_adjacent_in_place();
2140 let collected: Vec<_> = vm.values().collect();
2141 assert_eq!(collected, vec![0, 10, 20, 30, 40, 50]);
2142 assert_eq!(vm.get(4), Some(&40));
2144 }
2145
2146 #[test]
2147 fn rle_handles_singleton_mesh() {
2148 let region: Region = extent!(n = 1).into();
2149 let mut vm = ValueMesh::new_unchecked(region.clone(), vec![123]);
2150
2151 vm.compress_adjacent_in_place();
2152 let collected: Vec<_> = vm.values().collect();
2153 assert_eq!(collected, vec![123]);
2154 assert_eq!(vm.get(0), Some(&123));
2155 assert_eq!(vm.get(1), None);
2156 }
2157
2158 #[test]
2159 fn test_dense_round_trip_json() {
2160 let region: Region = extent!(x = 5).into();
2162 let dense = ValueMesh::new(region.clone(), vec![1, 2, 3, 4, 5]).unwrap();
2163
2164 let json = serde_json::to_string_pretty(&dense).unwrap();
2165 let restored: ValueMesh<i32> = serde_json::from_str(&json).unwrap();
2166
2167 assert_eq!(dense, restored);
2168
2169 let v: serde_json::Value = serde_json::from_str(&json).unwrap();
2172 let tag = v.get("rep").and_then(|o| o.get("Dense"));
2174 assert!(tag.is_some(), "json is {}", json);
2175 }
2176
2177 #[test]
2178 fn test_dense_round_trip_bincode() {
2179 let region: Region = extent!(x = 5).into();
2181 let dense = ValueMesh::new(region.clone(), vec![1, 2, 3, 4, 5]).unwrap();
2182
2183 let encoded = bincode::serde::encode_to_vec(&dense, bincode::config::legacy()).unwrap();
2184 let restored: ValueMesh<i32> =
2185 bincode::serde::decode_from_slice(&encoded, bincode::config::legacy())
2186 .map(|(v, _)| v)
2187 .unwrap();
2188
2189 assert_eq!(dense, restored);
2190 assert!(matches!(restored.rep, Rep::Dense { .. }));
2192 }
2193
2194 #[test]
2195 fn test_compressed_round_trip_json() {
2196 let region: Region = extent!(x = 10).into();
2199 let mut mesh = ValueMesh::new(region.clone(), vec![1, 1, 1, 2, 2, 3, 3, 3, 3, 3]).unwrap();
2200 mesh.compress_adjacent_in_place();
2201
2202 let json = serde_json::to_string_pretty(&mesh).unwrap();
2203 let restored: ValueMesh<i32> = serde_json::from_str(&json).unwrap();
2204
2205 assert_eq!(mesh, restored);
2207
2208 let v: serde_json::Value = serde_json::from_str(&json).unwrap();
2210 let tag = v.get("rep").and_then(|o| o.get("Compressed"));
2212 assert!(tag.is_some(), "json is {}", json);
2213 }
2214
2215 #[test]
2216 fn test_compressed_round_trip_bincode() {
2217 let region: Region = extent!(x = 10).into();
2220 let mut mesh = ValueMesh::new(region.clone(), vec![1, 1, 1, 2, 2, 3, 3, 3, 3, 3]).unwrap();
2221 mesh.compress_adjacent_in_place();
2222
2223 let encoded = bincode::serde::encode_to_vec(&mesh, bincode::config::legacy()).unwrap();
2224 let restored: ValueMesh<i32> =
2225 bincode::serde::decode_from_slice(&encoded, bincode::config::legacy())
2226 .map(|(v, _)| v)
2227 .unwrap();
2228
2229 assert_eq!(mesh, restored);
2231 assert!(matches!(restored.rep, Rep::Compressed { .. }));
2233 }
2234
2235 #[test]
2236 fn test_stable_run_encoding() {
2237 let run = Run::new(0, 10, 42);
2238 let json = serde_json::to_string(&run).unwrap();
2239 let decoded: Run = serde_json::from_str(&json).unwrap();
2240
2241 assert_eq!(run, decoded);
2242 assert_eq!(run.start, 0);
2243 assert_eq!(run.end, 10);
2244 assert_eq!(run.id, 42);
2245
2246 let (range, id): (Range<usize>, u32) = run.try_into().unwrap();
2248 assert_eq!(range, 0..10);
2249 assert_eq!(id, 42);
2250 }
2251
2252 #[test]
2253 fn from_single_builds_single_run() {
2254 let region: Region = extent!(n = 6).into();
2255 let vm = ValueMesh::from_single(region.clone(), 7);
2256
2257 assert_eq!(vm.region(), ®ion);
2258 assert_eq!(vm.values().collect::<Vec<_>>(), vec![7, 7, 7, 7, 7, 7]);
2259 assert_eq!(vm.get(0), Some(&7));
2260 assert_eq!(vm.get(5), Some(&7));
2261 assert_eq!(vm.get(6), None);
2262 }
2263
2264 #[test]
2265 fn from_default_builds_with_default_value() {
2266 let region: Region = extent!(n = 6).into();
2267 let vm = ValueMesh::<i32>::from_default(region.clone());
2268
2269 assert_eq!(vm.region(), ®ion);
2270 assert_eq!(vm.values().collect::<Vec<_>>(), vec![0, 0, 0, 0, 0, 0]);
2272 assert_eq!(vm.get(0), Some(&0));
2273 assert_eq!(vm.get(5), Some(&0));
2274 }
2275
2276 #[test]
2277 fn test_default_vs_single_equivalence() {
2278 let region: Region = extent!(x = 4).into();
2279 let d1 = ValueMesh::<i32>::from_default(region.clone());
2280 let d2 = ValueMesh::from_single(region.clone(), 0);
2281 assert_eq!(d1, d2);
2282 }
2283
2284 #[test]
2285 fn build_from_ranges_with_default_basic() {
2286 let region: Region = extent!(n = 10).into();
2287 let vm = ValueMesh::from_ranges_with_default(
2288 region.clone(),
2289 0, vec![(2..4, 1), (6..9, 2)],
2291 )
2292 .unwrap();
2293
2294 assert_eq!(vm.region(), ®ion);
2295 assert_eq!(
2296 vm.values().collect::<Vec<_>>(),
2297 vec![0, 0, 1, 1, 0, 0, 2, 2, 2, 0]
2298 );
2299
2300 if let Rep::Compressed { table, runs } = &vm.rep {
2303 assert!(table.len() <= 3);
2305 assert_eq!(runs.len(), 5);
2306 } else {
2307 panic!("expected compressed");
2308 }
2309 }
2310
2311 #[test]
2312 fn build_from_ranges_with_default_edge_cases() {
2313 let region: Region = extent!(n = 5).into();
2314
2315 let vm = ValueMesh::from_ranges_with_default(region.clone(), 9, vec![(0..5, 3)]).unwrap();
2317 assert_eq!(vm.values().collect::<Vec<_>>(), vec![3, 3, 3, 3, 3]);
2318
2319 let vm = ValueMesh::from_ranges_with_default(region.clone(), 0, vec![(1..2, 7), (2..4, 7)])
2321 .unwrap();
2322 assert_eq!(vm.values().collect::<Vec<_>>(), vec![0, 7, 7, 7, 0]);
2323
2324 let empty_region: Region = extent!(n = 0).into();
2326 let vm = ValueMesh::from_ranges_with_default(empty_region.clone(), 42, vec![]).unwrap();
2327 assert_eq!(vm.values().collect::<Vec<_>>(), Vec::<i32>::new());
2328 }
2329
2330 #[test]
2331 fn from_dense_builds_and_compresses() {
2332 let region: Region = extent!(n = 6).into();
2333 let mesh = ValueMesh::from_dense(region.clone(), vec![1, 1, 2, 2, 3, 3]).unwrap();
2334
2335 assert_eq!(mesh.region(), ®ion);
2336 assert!(matches!(mesh.rep, Rep::Compressed { .. }));
2337 assert_eq!(mesh.values().collect::<Vec<_>>(), vec![1, 1, 2, 2, 3, 3]);
2338
2339 assert_eq!(mesh.get(0), Some(&1));
2341 assert_eq!(mesh.get(3), Some(&2));
2342 assert_eq!(mesh.get(5), Some(&3));
2343 }
2344
2345 #[test]
2346 fn merge_from_overlay_basic() {
2347 let region: Region = extent!(n = 8).into();
2349 let mut mesh = ValueMesh::from_dense(region.clone(), vec![1, 1, 1, 2, 2, 2, 3, 3]).unwrap();
2350
2351 let overlay = ValueOverlay::try_from_runs(vec![(2..6, 9)]).unwrap();
2353
2354 mesh.merge_from_overlay(overlay).unwrap();
2355
2356 let out = mesh.materialized_runs();
2358
2359 assert_eq!(out, vec![(0..2, 1), (2..6, 9), (6..8, 3)]);
2361 }
2362
2363 #[test]
2364 fn merge_from_overlay_multiple_spans() {
2365 let region: Region = extent!(m = 12).into();
2367 let mut mesh =
2368 ValueMesh::from_dense(region.clone(), vec![1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4])
2369 .unwrap();
2370
2371 let overlay = ValueOverlay::try_from_runs(vec![(2..6, 9), (9..11, 8)]).unwrap();
2374
2375 mesh.merge_from_overlay(overlay).unwrap();
2376 let out = mesh.materialized_runs();
2377
2378 assert_eq!(
2385 out,
2386 vec![(0..2, 1), (2..6, 9), (6..9, 3), (9..11, 8), (11..12, 4)]
2387 );
2388 }
2389
2390 #[test]
2391 fn merge_from_overlay_crosses_row_boundary() {
2392 let region: Region = extent!(rows = 2, cols = 5).into();
2394
2395 let mut mesh =
2399 ValueMesh::from_dense(region.clone(), vec![1, 1, 1, 2, 2, 3, 3, 4, 4, 4]).unwrap();
2400
2401 let overlay = ValueOverlay::try_from_runs(vec![(3..7, 9)]).unwrap();
2406
2407 mesh.merge_from_overlay(overlay).unwrap();
2408
2409 let flat: Vec<_> = mesh.values().collect();
2412 assert_eq!(flat, vec![1, 1, 1, 9, 9, 9, 9, 4, 4, 4]);
2413
2414 let runs = mesh.materialized_runs();
2417 assert_eq!(runs, vec![(0..3, 1), (3..7, 9), (7..10, 4)]);
2418 }
2419
2420 #[test]
2421 fn accumulator_initializes_and_merges_overlays() {
2422 #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Named)]
2425 enum Val {
2426 Init,
2427 Running,
2428 Stopped,
2429 }
2430
2431 let region: Region = extent!(m = 8).into();
2432 let template = ValueMesh::from_single(region.clone(), Val::Init);
2433
2434 let mut state = ValueMesh::<Val>::default();
2436 assert_eq!(state.region().num_ranks(), 0);
2437
2438 let overlay1 = ValueOverlay::try_from_runs(vec![(2..5, Val::Running)]).unwrap();
2440 template.accumulate(&mut state, overlay1).unwrap();
2441
2442 assert_eq!(state.region().num_ranks(), 8);
2444 let vals: Vec<_> = state.values().collect();
2445 assert_eq!(
2446 vals,
2447 vec![
2448 Val::Init,
2449 Val::Init,
2450 Val::Running,
2451 Val::Running,
2452 Val::Running,
2453 Val::Init,
2454 Val::Init,
2455 Val::Init,
2456 ]
2457 );
2458
2459 let overlay2 = ValueOverlay::try_from_runs(vec![(5..7, Val::Stopped)]).unwrap();
2461 template.accumulate(&mut state, overlay2).unwrap();
2462
2463 let vals: Vec<_> = state.values().collect();
2464 assert_eq!(
2465 vals,
2466 vec![
2467 Val::Init,
2468 Val::Init,
2469 Val::Running,
2470 Val::Running,
2471 Val::Running,
2472 Val::Stopped,
2473 Val::Stopped,
2474 Val::Init,
2475 ]
2476 );
2477
2478 let spec = template
2480 .reducer_spec()
2481 .expect("reducer_spec should be Some");
2482 assert_eq!(
2483 spec.typehash,
2484 <ValueOverlayReducer<Val> as Named>::typehash()
2485 );
2486 }
2487
2488 #[test]
2489 fn value_overlay_reducer_merges_with_right_wins() {
2490 #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Named)]
2493 enum Val {
2494 Init,
2495 Running,
2496 Stopped,
2497 }
2498
2499 let reducer = ValueOverlayReducer::<Val>(PhantomData);
2500
2501 let left =
2503 ValueOverlay::try_from_runs(vec![(0..4, Val::Running), (6..8, Val::Stopped)]).unwrap();
2504
2505 let right = ValueOverlay::try_from_runs(vec![(2..6, Val::Init)]).unwrap();
2508
2509 let merged = reducer.reduce(left, right).unwrap();
2510 let runs: Vec<_> = merged.runs().cloned().collect();
2511
2512 assert_eq!(
2517 runs,
2518 vec![
2519 (0..2, Val::Running),
2520 (2..6, Val::Init),
2521 (6..8, Val::Stopped),
2522 ]
2523 );
2524 }
2525}