1use std::str::FromStr;
33use std::sync::Arc;
34
35use rankspace::Dim;
36use rankspace::Rank;
37use rankspace::RankRect;
38use rankspace::RankSpace;
39use serde::Deserialize;
40use serde::Serialize;
41use thiserror::Error;
42
43use crate::Range;
44use crate::Shape; use crate::Slice;
46use crate::SliceError;
47use crate::SliceIterator;
48use crate::parse::Parser;
49use crate::parse::ParserError;
50
51#[derive(Debug, thiserror::Error)]
53pub enum ExtentError {
54 #[error("label/sizes dimension mismatch: {num_labels} != {num_sizes}")]
60 DimMismatch {
61 num_labels: usize,
63 num_sizes: usize,
65 },
66
67 #[error("overlapping label found: {label}")]
72 OverlappingLabel {
73 label: String,
75 },
76}
77
78#[derive(Clone, Deserialize, Serialize, PartialEq, Eq, Hash, Debug)]
86pub struct Extent {
87 inner: Arc<ExtentData>,
88}
89
90fn _assert_extent_traits()
91where
92 Extent: Send + Sync + 'static,
93{
94}
95
96#[derive(Clone, Deserialize, Serialize, PartialEq, Eq, Hash, Debug)]
101struct ExtentData {
102 labels: Vec<String>,
103 sizes: Vec<usize>,
104}
105
106impl From<&Shape> for Extent {
107 fn from(s: &Shape) -> Self {
108 Extent::new(s.labels().to_vec(), s.slice().sizes().to_vec()).unwrap()
110 }
111}
112
113impl From<Shape> for Extent {
114 fn from(s: Shape) -> Self {
115 Extent::from(&s)
116 }
117}
118
119impl Extent {
120 pub fn new(labels: Vec<String>, sizes: Vec<usize>) -> Result<Self, ExtentError> {
122 if labels.len() != sizes.len() {
123 return Err(ExtentError::DimMismatch {
124 num_labels: labels.len(),
125 num_sizes: sizes.len(),
126 });
127 }
128
129 Ok(Self {
130 inner: Arc::new(ExtentData { labels, sizes }),
131 })
132 }
133
134 pub fn unity() -> Extent {
135 Extent::new(vec![], vec![]).unwrap()
136 }
137
138 pub fn labels(&self) -> &[String] {
140 &self.inner.labels
141 }
142
143 pub fn sizes(&self) -> &[usize] {
145 &self.inner.sizes
146 }
147
148 pub fn size(&self, label: &str) -> Option<usize> {
151 self.position(label).map(|pos| self.sizes()[pos])
152 }
153
154 pub fn position(&self, label: &str) -> Option<usize> {
157 self.labels().iter().position(|l| l == label)
158 }
159
160 pub fn rank_of_coords(&self, coords: &[usize]) -> Result<usize, PointError> {
170 let sizes = self.sizes();
171 if coords.len() != sizes.len() {
172 return Err(PointError::DimMismatch {
173 expected: sizes.len(),
174 actual: coords.len(),
175 });
176 }
177 let mut stride = 1;
178 let mut result = 0;
179 for (&c, &size) in coords.iter().rev().zip(sizes.iter().rev()) {
180 if c >= size {
181 return Err(PointError::OutOfRangeIndex { size, index: c });
182 }
183 result += c * stride;
184 stride *= size;
185 }
186 Ok(result)
187 }
188
189 pub fn point(&self, coords: Vec<usize>) -> Result<Point, PointError> {
235 Ok(Point {
236 rank: self.rank_of_coords(&coords)?,
237 extent: self.clone(),
238 })
239 }
240
241 pub fn point_of_rank(&self, rank: usize) -> Result<Point, PointError> {
265 let total = self.num_ranks();
266 if rank >= total {
267 return Err(PointError::OutOfRangeRank { total, rank });
268 }
269 Ok(Point {
270 rank,
271 extent: self.clone(),
272 })
273 }
274
275 pub fn len(&self) -> usize {
280 self.sizes().len()
281 }
282
283 pub fn is_empty(&self) -> bool {
288 self.sizes().is_empty()
289 }
290
291 pub fn num_ranks(&self) -> usize {
296 self.sizes().iter().product()
297 }
298
299 pub fn into_inner(self) -> (Vec<String>, Vec<usize>) {
301 match Arc::try_unwrap(self.inner) {
302 Ok(data) => (data.labels, data.sizes),
303 Err(shared) => (shared.labels.clone(), shared.sizes.clone()),
304 }
305 }
306
307 pub fn to_slice(&self) -> Slice {
309 Slice::new_row_major(self.sizes())
310 }
311
312 pub fn iter(&self) -> impl Iterator<Item = (String, usize)> + use<'_> {
314 self.labels()
315 .iter()
316 .zip(self.sizes().iter())
317 .map(|(l, s)| (l.clone(), *s))
318 }
319
320 pub fn points(&self) -> ExtentPointsIterator<'_> {
322 ExtentPointsIterator::new(self)
323 }
324
325 pub fn concat(&self, other: &Extent) -> Result<Self, ExtentError> {
333 use std::collections::HashSet;
334 let lhs: HashSet<&str> = self.labels().iter().map(|s| s.as_str()).collect();
336 if let Some(dup) = other.labels().iter().find(|l| lhs.contains(l.as_str())) {
337 return Err(ExtentError::OverlappingLabel { label: dup.clone() });
338 }
339 let mut labels = self.labels().to_vec();
341 let mut sizes = self.sizes().to_vec();
342 labels.reserve(other.labels().len());
343 sizes.reserve(other.sizes().len());
344 labels.extend(other.labels().iter().cloned());
345 sizes.extend(other.sizes().iter().copied());
346 Extent::new(labels, sizes)
347 }
348}
349
350mod labels {
362 pub(super) fn is_safe_ident(s: &str) -> bool {
366 s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
367 }
368
369 pub(super) fn fmt_label(s: &str) -> String {
374 if is_safe_ident(s) {
375 s.to_string()
376 } else {
377 format!("{:?}", s)
378 }
379 }
380}
381
382impl std::fmt::Display for Extent {
418 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
419 write!(f, "{{")?;
420 for i in 0..self.sizes().len() {
421 write!(
422 f,
423 "{}: {}",
424 labels::fmt_label(&self.labels()[i]),
425 self.sizes()[i]
426 )?;
427 if i + 1 != self.sizes().len() {
428 write!(f, ", ")?;
429 }
430 }
431 write!(f, "}}")
432 }
433}
434
435pub struct ExtentPointsIterator<'a> {
437 extent: &'a Extent,
438 next_rank: usize,
439}
440
441impl<'a> ExtentPointsIterator<'a> {
442 pub fn new(extent: &'a Extent) -> Self {
443 Self {
444 extent,
445 next_rank: 0,
446 }
447 }
448}
449
450impl Iterator for ExtentPointsIterator<'_> {
451 type Item = Point;
452
453 fn next(&mut self) -> Option<Self::Item> {
456 if self.next_rank == self.extent.num_ranks() {
457 return None;
458 }
459
460 let p = Point {
461 rank: self.next_rank,
462 extent: self.extent.clone(),
463 };
464 self.next_rank += 1;
465 Some(p)
466 }
467}
468
469#[derive(Debug, Error)]
471pub enum PointError {
472 #[error("dimension mismatch: expected {expected}, got {actual}")]
479 DimMismatch {
480 expected: usize,
482 actual: usize,
484 },
485
486 #[error("out of range: total ranks {total}; does not contain rank {rank}")]
493 OutOfRangeRank {
494 total: usize,
496 rank: usize,
498 },
499
500 #[error("out of range: dim size {size}; does not contain index {index}")]
508 OutOfRangeIndex {
509 size: usize,
511 index: usize,
513 },
514
515 #[error("failed to parse point: {reason}")]
517 ParseError { reason: String },
518}
519
520#[derive(Clone, Deserialize, Serialize, PartialEq, Eq, Hash, Debug)]
553pub struct Point {
554 rank: usize,
555 extent: Extent,
556}
557
558pub struct CoordIter<'a> {
578 sizes: &'a [usize],
579 rank: usize,
580 stride: usize,
581 axis: usize,
582}
583
584impl Iterator for CoordIter<'_> {
585 type Item = usize;
586
587 fn next(&mut self) -> Option<Self::Item> {
592 if self.axis >= self.sizes.len() {
593 return None;
594 }
595 self.stride /= self.sizes[self.axis];
596 let q = self.rank / self.stride;
597 self.rank %= self.stride;
598 self.axis += 1;
599 Some(q)
600 }
601
602 fn size_hint(&self) -> (usize, Option<usize>) {
608 let rem = self.sizes.len().saturating_sub(self.axis);
609 (rem, Some(rem))
610 }
611}
612
613impl ExactSizeIterator for CoordIter<'_> {}
614
615impl<'a> IntoIterator for &'a Point {
616 type Item = usize;
617 type IntoIter = CoordIter<'a>;
618
619 fn into_iter(self) -> Self::IntoIter {
638 self.coords_iter()
639 }
640}
641
642fn _assert_point_traits()
643where
644 Point: Send + Sync + 'static,
645{
646}
647
648pub trait InExtent {
664 fn in_(self, extent: &Extent) -> Result<Point, PointError>;
665}
666
667impl InExtent for Vec<usize> {
668 fn in_(self, extent: &Extent) -> Result<Point, PointError> {
673 extent.point(self)
674 }
675}
676
677impl Point {
678 pub fn coords_iter(&self) -> CoordIter<'_> {
679 CoordIter {
680 sizes: self.extent.sizes(),
681 rank: self.rank,
682 stride: self.extent.sizes().iter().product(),
683 axis: 0,
684 }
685 }
686
687 pub fn coord(&self, i: usize) -> usize {
705 self.coords_iter()
706 .nth(i)
707 .expect("coord(i): axis out of bounds")
708 }
709
710 pub fn coords(&self) -> Vec<usize> {
725 self.coords_iter().collect()
726 }
727
728 pub fn rank(&self) -> usize {
731 self.rank
732 }
733
734 pub fn extent(&self) -> &Extent {
737 &self.extent
738 }
739
740 pub fn len(&self) -> usize {
755 self.extent.len()
756 }
757
758 pub fn is_empty(&self) -> bool {
774 self.extent.len() == 0
775 }
776
777 pub fn format_as_dict(&self) -> String {
789 format!(
790 "{{{}}}",
791 self.extent()
792 .labels()
793 .iter()
794 .zip(self.coords_iter())
795 .zip(self.extent().sizes())
796 .map(|((label, coord), size)| format!("'{}': {}/{}", label, coord, size))
797 .collect::<Vec<_>>()
798 .join(", ")
799 )
800 }
801}
802
803impl std::fmt::Display for Point {
839 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
840 let labels = self.extent.labels();
841 let sizes = self.extent.sizes();
842 let coords = self.coords();
843
844 for i in 0..labels.len() {
845 write!(
846 f,
847 "{}={}/{}",
848 labels::fmt_label(&labels[i]),
849 coords[i],
850 sizes[i]
851 )?;
852 if i + 1 != labels.len() {
853 write!(f, ",")?;
854 }
855 }
856 Ok(())
857 }
858}
859
860impl FromStr for Point {
861 type Err = PointError;
862
863 fn from_str(s: &str) -> Result<Self, Self::Err> {
864 let s = s.trim();
865
866 if s.is_empty() {
867 let empty_extent = Extent::unity();
868 return empty_extent.point(vec![]);
869 }
870
871 let mut labels = Vec::new();
872 let mut coords = Vec::new();
873 let mut sizes = Vec::new();
874
875 let mut chars = s.chars().peekable();
876
877 while chars.peek().is_some() {
878 while chars.peek() == Some(&' ') {
879 chars.next();
880 }
881
882 if chars.peek().is_none() {
883 break;
884 }
885
886 let label = if chars.peek() == Some(&'"') {
887 chars.next(); let mut label = String::new();
889 let mut escaped = false;
890
891 for ch in chars.by_ref() {
893 if escaped {
894 match ch {
895 '"' => label.push('"'),
896 '\\' => label.push('\\'),
897 _ => {
898 label.push('\\');
899 label.push(ch);
900 }
901 }
902 escaped = false;
903 } else if ch == '\\' {
904 escaped = true;
905 } else if ch == '"' {
906 break;
907 } else {
908 label.push(ch);
909 }
910 }
911
912 if label.is_empty() {
913 return Err(PointError::ParseError {
914 reason: "empty quoted label".to_string(),
915 });
916 }
917
918 label
919 } else {
920 let mut label = String::new();
921 while let Some(&ch) = chars.peek() {
922 if ch == '=' || ch == ' ' {
923 break;
924 }
925 label.push(chars.next().unwrap());
926 }
927
928 if label.is_empty() {
929 return Err(PointError::ParseError {
930 reason: "missing label".to_string(),
931 });
932 }
933
934 label
935 };
936
937 while chars.peek() == Some(&' ') {
938 chars.next();
939 }
940
941 if chars.next() != Some('=') {
942 return Err(PointError::ParseError {
943 reason: format!("expected '=' after label '{}'", label),
944 });
945 }
946
947 while chars.peek() == Some(&' ') {
948 chars.next();
949 }
950
951 let mut coord = String::new();
952 while let Some(&ch) = chars.peek() {
953 if ch == '/' || ch == ' ' {
954 break;
955 }
956 coord.push(chars.next().unwrap());
957 }
958
959 if coord.is_empty() {
960 return Err(PointError::ParseError {
961 reason: format!("missing coordinate for dimension '{}'", label),
962 });
963 }
964
965 while chars.peek() == Some(&' ') {
966 chars.next();
967 }
968
969 if chars.next() != Some('/') {
970 return Err(PointError::ParseError {
971 reason: format!("expected '/' after coordinate for dimension '{}'", label),
972 });
973 }
974
975 while chars.peek() == Some(&' ') {
976 chars.next();
977 }
978
979 let mut size = String::new();
980 while let Some(&ch) = chars.peek() {
981 if ch == ',' || ch == ' ' {
982 break;
983 }
984 size.push(chars.next().unwrap());
985 }
986
987 if size.is_empty() {
988 return Err(PointError::ParseError {
989 reason: format!("missing size for dimension '{}'", label),
990 });
991 }
992
993 let coord = coord.parse::<usize>().map_err(|e| PointError::ParseError {
994 reason: format!(
995 "invalid coordinate '{}' for dimension '{}': {}",
996 coord, label, e
997 ),
998 })?;
999
1000 let size = size.parse::<usize>().map_err(|e| PointError::ParseError {
1001 reason: format!("invalid size '{}' for dimension '{}': {}", size, label, e),
1002 })?;
1003
1004 labels.push(label);
1005 coords.push(coord);
1006 sizes.push(size);
1007
1008 while chars.peek() == Some(&' ') {
1009 chars.next();
1010 }
1011
1012 if chars.peek() == Some(&',') {
1013 chars.next(); while chars.peek() == Some(&' ') {
1015 chars.next();
1016 }
1017 if chars.peek().is_none() {
1019 return Err(PointError::ParseError {
1020 reason: "trailing comma".to_string(),
1021 });
1022 }
1023 }
1024 }
1025
1026 let extent = Extent::new(labels, sizes).map_err(|e| PointError::ParseError {
1027 reason: format!("failed to create extent: {}", e),
1028 })?;
1029
1030 extent.point(coords)
1031 }
1032}
1033
1034impl typeuri::Named for Point {
1035 fn typename() -> &'static str {
1036 "ndslice::Point"
1037 }
1038}
1039
1040hyperactor_config::impl_attrvalue!(Point);
1041
1042#[derive(Debug, Error)]
1044pub enum ViewError {
1045 #[error("no such dimension: {0}")]
1047 InvalidDim(String),
1048
1049 #[error("empty range: {range} for dimension {dim} of size {size}")]
1051 EmptyRange {
1052 range: Range,
1053 dim: String,
1054 size: usize,
1055 },
1056
1057 #[error(transparent)]
1058 ExtentError(#[from] ExtentError),
1059
1060 #[error("invalid range: selected ranks {selected} not a subset of base {base} ")]
1061 InvalidRange {
1062 base: Box<Region>,
1063 selected: Box<Region>,
1064 },
1065}
1066
1067#[derive(Debug, Error)]
1069pub enum RegionError {
1070 #[error("invalid point: this point does not belong to this region: {0}")]
1071 InvalidPoint(String),
1072
1073 #[error("out of range base rank: this base rank {0} does not belong to this region: {0}")]
1074 OutOfRangeBaseRank(usize, String),
1075}
1076
1077#[derive(Debug, Clone, Serialize, Deserialize, Hash, PartialEq, Eq)]
1084pub struct Region {
1085 labels: Vec<String>,
1086 slice: Slice,
1087}
1088
1089impl From<Region> for RankSpace {
1090 fn from(region: Region) -> Self {
1091 let Region { labels, slice } = region;
1092 let dims = labels
1093 .into_iter()
1094 .zip(slice.sizes())
1095 .map(|(label, &size)| Dim::new(label, size))
1096 .collect();
1097 let extent =
1098 rankspace::Extent::new(dims).expect("region labels and sizes must form a valid extent");
1099 let rect = RankRect::affine(extent, Rank(slice.offset()), slice.strides().to_vec())
1100 .expect("region slice must form a valid rank rect");
1101 RankSpace::dense(rect)
1102 }
1103}
1104
1105impl Region {
1106 #[allow(dead_code)]
1107 fn empty() -> Region {
1108 Region {
1109 labels: Vec::new(),
1110 slice: Slice::new(0, Vec::new(), Vec::new()).unwrap(),
1111 }
1112 }
1113
1114 #[allow(dead_code)]
1117 pub fn new(labels: Vec<String>, slice: Slice) -> Self {
1118 Self { labels, slice }
1119 }
1120
1121 pub fn labels(&self) -> &[String] {
1123 &self.labels
1124 }
1125
1126 pub fn slice(&self) -> &Slice {
1129 &self.slice
1130 }
1131
1132 pub fn into_inner(self) -> (Vec<String>, Slice) {
1134 (self.labels, self.slice)
1135 }
1136
1137 pub fn extent(&self) -> Extent {
1139 Extent::new(self.labels.clone(), self.slice.sizes().to_vec()).unwrap()
1140 }
1141
1142 pub fn is_subset(&self, other: &Region) -> bool {
1145 let mut left = self.slice.iter().peekable();
1146 let mut right = other.slice.iter().peekable();
1147
1148 loop {
1149 match (left.peek(), right.peek()) {
1150 (Some(l), Some(r)) => {
1151 if l < r {
1152 return false;
1153 } else if l == r {
1154 left.next();
1155 right.next();
1156 } else {
1157 right.next();
1159 }
1160 }
1161 (Some(_), None) => return false,
1162 (None, _) => return true,
1163 }
1164 }
1165 }
1166
1167 pub fn remap(&self, target: &Region) -> Option<impl Iterator<Item = usize> + '_> {
1187 if !target.is_subset(self) {
1188 return None;
1189 }
1190
1191 let mut ours = self.slice.iter().enumerate();
1192 let mut theirs = target.slice.iter();
1193
1194 Some(std::iter::from_fn(move || {
1195 let needle = theirs.next()?;
1196 loop {
1197 let (index, value) = ours.next().unwrap();
1198 if value == needle {
1199 break Some(index);
1200 }
1201 }
1202 }))
1203 }
1204
1205 pub fn num_ranks(&self) -> usize {
1207 self.slice.len()
1208 }
1209
1210 pub fn base_rank_of_point(&self, p: Point) -> Result<usize, RegionError> {
1213 if p.extent() != &self.extent() {
1214 return Err(RegionError::InvalidPoint(
1215 "mismatched extent: p must be a point in this region’s extent".to_string(),
1216 ));
1217 }
1218
1219 Ok(self
1220 .slice()
1221 .location(&p.coords())
1222 .expect("should have valid location since extent is checked"))
1223 }
1224
1225 pub fn point_of_base_rank(&self, rank: usize) -> Result<Point, RegionError> {
1228 let coords = self
1229 .slice()
1230 .coordinates(rank)
1231 .map_err(|e| RegionError::OutOfRangeBaseRank(rank, e.to_string()))?;
1232 Ok(self
1233 .extent()
1234 .point(coords)
1235 .expect("should have valid point since coords is from this region"))
1236 }
1237}
1238
1239impl From<Extent> for Region {
1242 fn from(extent: Extent) -> Self {
1243 Region {
1244 labels: extent.labels().to_vec(),
1245 slice: extent.to_slice(),
1246 }
1247 }
1248}
1249
1250impl From<&Shape> for Region {
1251 fn from(s: &Shape) -> Self {
1252 Region {
1253 labels: s.labels().to_vec(),
1254 slice: s.slice().clone(),
1255 }
1256 }
1257}
1258
1259impl From<Shape> for Region {
1260 fn from(s: Shape) -> Self {
1261 Region::from(&s)
1262 }
1263}
1264
1265impl std::fmt::Display for Region {
1305 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1306 if self.slice.offset() != 0 {
1307 write!(f, "{}+", self.slice.offset())?;
1308 }
1309 for i in 0..self.labels.len() {
1310 write!(
1311 f,
1312 "{}={}/{}",
1313 labels::fmt_label(&self.labels[i]),
1314 self.slice.sizes()[i],
1315 self.slice.strides()[i]
1316 )?;
1317 if i + 1 != self.labels.len() {
1318 write!(f, ",")?;
1319 }
1320 }
1321 Ok(())
1322 }
1323}
1324
1325#[derive(Debug, thiserror::Error)]
1326pub enum RegionParseError {
1327 #[error(transparent)]
1328 ParserError(#[from] ParserError),
1329
1330 #[error(transparent)]
1331 SliceError(#[from] SliceError),
1332}
1333
1334impl std::str::FromStr for Region {
1355 type Err = RegionParseError;
1356
1357 fn from_str(s: &str) -> Result<Self, Self::Err> {
1358 let mut parser = Parser::new(s, &["+", "=", ",", "/"]);
1359
1360 let offset: usize = if let Ok(offset) = parser.try_parse() {
1361 parser.expect("+")?;
1362 offset
1363 } else {
1364 0
1365 };
1366
1367 let mut labels = Vec::new();
1368 let mut sizes = Vec::new();
1369 let mut strides = Vec::new();
1370
1371 while !parser.is_empty() {
1372 if !labels.is_empty() {
1373 parser.expect(",")?;
1374 }
1375
1376 let label = if parser.peek_char() == Some('"') {
1378 parser.parse_string_literal()?
1379 } else {
1380 parser.next_or_err("label")?.to_string()
1381 };
1382 labels.push(label);
1383
1384 parser.expect("=")?;
1385 sizes.push(parser.try_parse()?);
1386 parser.expect("/")?;
1387 strides.push(parser.try_parse()?);
1388 }
1389
1390 Ok(Region {
1391 labels,
1392 slice: Slice::new(offset, sizes, strides)?,
1393 })
1394 }
1395}
1396
1397pub trait BuildFromRegion<T>: Sized {
1410 type Error;
1411
1412 fn build_dense(region: Region, values: Vec<T>) -> Result<Self, Self::Error>;
1414
1415 fn build_dense_unchecked(region: Region, values: Vec<T>) -> Self;
1417}
1418
1419pub trait BuildFromRegionIndexed<T>: Sized {
1433 type Error;
1434
1435 fn build_indexed(
1438 region: Region,
1439 pairs: impl IntoIterator<Item = (usize, T)>,
1440 ) -> Result<Self, Self::Error>;
1441}
1442
1443pub trait CollectMeshExt<T>: Iterator<Item = T> + Sized {
1457 fn collect_mesh<M>(self, region: Region) -> Result<M, M::Error>
1458 where
1459 M: BuildFromRegion<T>;
1460}
1461
1462impl<I, T> CollectMeshExt<T> for I
1465where
1466 I: Iterator<Item = T> + Sized,
1467{
1468 fn collect_mesh<M>(self, region: Region) -> Result<M, M::Error>
1469 where
1470 M: BuildFromRegion<T>,
1471 {
1472 M::build_dense(region, self.collect())
1473 }
1474}
1475
1476#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1478pub struct InvalidCardinality {
1479 pub expected: usize,
1480 pub actual: usize,
1481}
1482
1483pub trait CollectExactMeshExt<T>: ExactSizeIterator<Item = T> + Sized {
1490 fn collect_exact_mesh<M>(self, region: Region) -> Result<M, M::Error>
1491 where
1492 M: BuildFromRegion<T>,
1493 M::Error: From<InvalidCardinality>;
1494}
1495
1496impl<I, T> CollectExactMeshExt<T> for I
1499where
1500 I: ExactSizeIterator<Item = T> + Sized,
1501{
1502 fn collect_exact_mesh<M>(self, region: Region) -> Result<M, M::Error>
1503 where
1504 M: BuildFromRegion<T>,
1505 M::Error: From<InvalidCardinality>,
1506 {
1507 let expected = region.num_ranks();
1508 let actual = self.len();
1509 if actual != expected {
1510 return Err(M::Error::from(InvalidCardinality { expected, actual }));
1511 }
1512 Ok(M::build_dense_unchecked(region, self.collect()))
1513 }
1514}
1515
1516pub trait CollectIndexedMeshExt<T>: Iterator<Item = (usize, T)> + Sized {
1538 fn collect_indexed<M>(self, region: Region) -> Result<M, M::Error>
1539 where
1540 M: BuildFromRegionIndexed<T>;
1541}
1542
1543impl<I, T> CollectIndexedMeshExt<T> for I
1546where
1547 I: Iterator<Item = (usize, T)> + Sized,
1548{
1549 #[inline]
1550 fn collect_indexed<M>(self, region: Region) -> Result<M, M::Error>
1551 where
1552 M: BuildFromRegionIndexed<T>,
1553 {
1554 M::build_indexed(region, self)
1555 }
1556}
1557
1558pub trait MapIntoExt: Ranked {
1560 fn map_into<M, U>(&self, f: impl Fn(&Self::Item) -> U) -> M
1561 where
1562 Self: Sized,
1563 M: BuildFromRegion<U>,
1564 {
1565 let region = self.region().clone();
1566 let n = region.num_ranks();
1567 let values: Vec<U> = (0..n).map(|i| f(self.get(i).unwrap())).collect();
1568 M::build_dense_unchecked(region, values)
1569 }
1570
1571 fn try_map_into<M, U, E>(self, f: impl Fn(&Self::Item) -> Result<U, E>) -> Result<M, E>
1572 where
1573 Self: Sized,
1574 M: BuildFromRegion<U>,
1575 {
1576 let region = self.region().clone();
1577 let n = region.num_ranks();
1578 let mut out = Vec::with_capacity(n);
1579 for i in 0..n {
1580 out.push(f(self.get(i).unwrap())?);
1581 }
1582 Ok(M::build_dense_unchecked(region, out))
1583 }
1584}
1585
1586impl<T: Ranked> MapIntoExt for T {}
1589
1590pub trait View: Sized {
1592 type Item;
1594
1595 type View: View;
1597
1598 fn region(&self) -> Region;
1600
1601 fn get(&self, rank: usize) -> Option<Self::Item>;
1605
1606 #[allow(clippy::result_large_err)] fn subset(&self, region: Region) -> Result<Self::View, ViewError>;
1611}
1612
1613impl View for Region {
1615 type Item = usize;
1617
1618 type View = Region;
1620
1621 fn region(&self) -> Region {
1622 self.clone()
1623 }
1624
1625 fn subset(&self, region: Region) -> Result<Region, ViewError> {
1626 if region.is_subset(self) {
1627 Ok(region)
1628 } else {
1629 Err(ViewError::InvalidRange {
1630 base: Box::new(self.clone()),
1631 selected: Box::new(region),
1632 })
1633 }
1634 }
1635
1636 fn get(&self, rank: usize) -> Option<Self::Item> {
1637 self.slice.get(rank).ok()
1638 }
1639}
1640
1641impl View for Extent {
1643 type Item = usize;
1645
1646 type View = Region;
1649
1650 fn region(&self) -> Region {
1651 Region {
1652 labels: self.labels().to_vec(),
1653 slice: self.to_slice(),
1654 }
1655 }
1656
1657 fn subset(&self, region: Region) -> Result<Region, ViewError> {
1658 self.region().subset(region)
1659 }
1660
1661 fn get(&self, rank: usize) -> Option<Self::Item> {
1662 if rank < self.num_ranks() {
1663 Some(rank)
1664 } else {
1665 None
1666 }
1667 }
1668}
1669
1670pub trait Ranked: Sized {
1673 type Item: 'static;
1675
1676 fn region(&self) -> &Region;
1678
1679 fn get(&self, rank: usize) -> Option<&Self::Item>;
1681}
1682
1683pub trait RankedSliceable: Ranked {
1691 fn sliced(&self, region: Region) -> Self;
1695}
1696
1697impl<T: RankedSliceable> View for T
1698where
1699 T::Item: Clone + 'static,
1700{
1701 type Item = T::Item;
1702 type View = Self;
1703
1704 fn region(&self) -> Region {
1705 <Self as Ranked>::region(self).clone()
1706 }
1707
1708 fn get(&self, rank: usize) -> Option<Self::Item> {
1709 <Self as Ranked>::get(self, rank).cloned()
1710 }
1711
1712 fn subset(&self, region: Region) -> Result<Self, ViewError> {
1713 if !region.is_subset(self.region()) {
1714 return Err(ViewError::InvalidRange {
1715 base: Box::new(self.region().clone()),
1716 selected: Box::new(region.clone()),
1717 });
1718 }
1719
1720 Ok(self.sliced(region))
1721 }
1722}
1723
1724pub struct ViewIterator {
1726 extent: Extent, pos: SliceIterator, }
1729
1730impl Iterator for ViewIterator {
1731 type Item = (Point, usize);
1732 fn next(&mut self) -> Option<Self::Item> {
1733 let rank = self.pos.next()?;
1735 let coords = self.pos.slice.coordinates(rank).unwrap();
1737 let point = coords.in_(&self.extent).unwrap();
1738 Some((point, rank))
1739 }
1740}
1741
1742pub trait ViewExt: View {
1744 #[allow(clippy::result_large_err)] fn range<R: Into<Range>>(&self, dim: &str, range: R) -> Result<Self::View, ViewError>;
1771
1772 #[allow(clippy::result_large_err)] fn group_by(&self, dim: &str) -> Result<impl Iterator<Item = Self::View>, ViewError>;
1809
1810 fn extent(&self) -> Extent;
1812
1813 fn iter(&self) -> impl Iterator<Item = (Point, Self::Item)> + '_;
1815
1816 fn values(&self) -> impl Iterator<Item = Self::Item> + '_;
1818}
1819
1820impl<T: View> ViewExt for T {
1821 fn range<R: Into<Range>>(&self, dim: &str, range: R) -> Result<Self::View, ViewError> {
1822 let (labels, slice) = self.region().into_inner();
1823 let range = range.into();
1824 let dim = labels
1825 .iter()
1826 .position(|l| dim == l)
1827 .ok_or_else(|| ViewError::InvalidDim(dim.to_string()))?;
1828 let (mut offset, mut sizes, mut strides) = slice.into_inner();
1829 let (begin, end, step) = range.resolve(sizes[dim]);
1830 if end <= begin {
1831 return Err(ViewError::EmptyRange {
1832 range,
1833 dim: dim.to_string(),
1834 size: sizes[dim],
1835 });
1836 }
1837
1838 offset += strides[dim] * begin;
1839 sizes[dim] = (end - begin).div_ceil(step);
1840 strides[dim] *= step;
1841 let slice = Slice::new(offset, sizes, strides).unwrap();
1842
1843 self.subset(Region { labels, slice })
1844 }
1845
1846 fn group_by(&self, dim: &str) -> Result<impl Iterator<Item = Self::View>, ViewError> {
1847 let (labels, slice) = self.region().into_inner();
1848
1849 let dim = labels
1850 .iter()
1851 .position(|l| dim == l)
1852 .ok_or_else(|| ViewError::InvalidDim(dim.to_string()))?;
1853
1854 let (offset, sizes, strides) = slice.into_inner();
1855 let mut ranks_iter = Slice::new(offset, sizes[..dim].to_vec(), strides[..dim].to_vec())
1856 .unwrap()
1857 .iter();
1858
1859 let labels = labels[dim..].to_vec();
1860 let sizes = sizes[dim..].to_vec();
1861 let strides = strides[dim..].to_vec();
1862
1863 Ok(std::iter::from_fn(move || {
1864 let rank = ranks_iter.next()?;
1865 let slice = Slice::new(rank, sizes.clone(), strides.clone()).unwrap();
1866 Some(
1868 self.subset(Region {
1869 labels: labels.clone(),
1870 slice,
1871 })
1872 .unwrap(),
1873 )
1874 }))
1875 }
1876
1877 fn extent(&self) -> Extent {
1878 let (labels, slice) = self.region().into_inner();
1879 Extent::new(labels, slice.sizes().to_vec()).unwrap()
1880 }
1881
1882 fn iter(&self) -> impl Iterator<Item = (Point, Self::Item)> + '_ {
1883 let points = ViewIterator {
1884 extent: self.extent(),
1885 pos: self.region().slice().iter(),
1886 };
1887
1888 points.map(|(point, _)| (point.clone(), self.get(point.rank()).unwrap()))
1889 }
1890
1891 fn values(&self) -> impl Iterator<Item = Self::Item> + '_ {
1892 (0usize..self.extent().num_ranks()).map(|rank| self.get(rank).unwrap())
1893 }
1894}
1895
1896#[macro_export]
1904macro_rules! extent {
1905 ( $( $label:ident = $size:expr ),* $(,)? ) => {
1906 {
1907 #[allow(unused_mut)]
1908 let mut labels = Vec::new();
1909 #[allow(unused_mut)]
1910 let mut sizes = Vec::new();
1911
1912 $(
1913 labels.push(stringify!($label).to_string());
1914 sizes.push($size);
1915 )*
1916
1917 $crate::view::Extent::new(labels, sizes).unwrap()
1918 }
1919 };
1920}
1921
1922#[cfg(test)]
1923mod test {
1924 use super::labels::*;
1925 use super::*;
1926 use crate::Shape;
1927 use crate::shape;
1928 use crate::slice::CartesianIterator;
1929
1930 #[test]
1931 fn region_to_rankspace_maps_ranks_in_order() {
1932 let region: Region = extent!(host = 2, gpu = 3).into();
1933 let space = RankSpace::from(region.clone());
1934
1935 assert_eq!(space.cardinality(), region.num_ranks());
1936 let space_ranks: Vec<usize> = space.iter_ranks().map(Rank::get).collect();
1937 let region_ranks: Vec<usize> = region.slice().iter().collect();
1938 assert_eq!(space_ranks, region_ranks);
1939
1940 for (index, base_rank) in region.slice().iter().enumerate() {
1941 assert_eq!(space.local_index_of(Rank(base_rank)), Some(index));
1942 }
1943 }
1944
1945 #[test]
1946 fn region_to_rankspace_preserves_base_ranks_when_sliced() {
1947 let region: Region = extent!(replica = 4, gpu = 2).into();
1948 let sub = region
1949 .range("replica", 1..3)
1950 .expect("slice should be valid");
1951 let space = RankSpace::from(sub.clone());
1952
1953 assert_eq!(space.cardinality(), sub.num_ranks());
1954 let space_ranks: Vec<usize> = space.iter_ranks().map(Rank::get).collect();
1955 let region_ranks: Vec<usize> = sub.slice().iter().collect();
1956 assert_eq!(space_ranks, region_ranks);
1957 assert!(!space.contains_rank(Rank(0)));
1958 }
1959
1960 #[test]
1961 fn test_is_safe_ident() {
1962 assert!(is_safe_ident("x"));
1963 assert!(is_safe_ident("gpu_0"));
1964 assert!(!is_safe_ident("dim/0"));
1965 assert!(!is_safe_ident("x y"));
1966 assert!(!is_safe_ident("x=y"));
1967 }
1968 #[test]
1969 fn test_fmt_label() {
1970 assert_eq!(fmt_label("x"), "x");
1971 assert_eq!(fmt_label("dim/0"), "\"dim/0\"");
1972 }
1973
1974 #[test]
1975 fn test_points_basic() {
1976 let extent = extent!(x = 4, y = 5, z = 6);
1977 let _p1 = extent.point(vec![1, 2, 3]).unwrap();
1978 let _p2 = vec![1, 2, 3].in_(&extent).unwrap();
1979
1980 assert_eq!(extent.num_ranks(), 4 * 5 * 6);
1981
1982 let p3 = extent.point_of_rank(0).unwrap();
1983 assert_eq!(p3.coords(), &[0, 0, 0]);
1984 assert_eq!(p3.rank(), 0);
1985
1986 let p4 = extent.point_of_rank(1).unwrap();
1987 assert_eq!(p4.coords(), &[0, 0, 1]);
1988 assert_eq!(p4.rank(), 1);
1989
1990 let p5 = extent.point_of_rank(2).unwrap();
1991 assert_eq!(p5.coords(), &[0, 0, 2]);
1992 assert_eq!(p5.rank(), 2);
1993
1994 let p6 = extent.point_of_rank(6 * 5 + 1).unwrap();
1995 assert_eq!(p6.coords(), &[1, 0, 1]);
1996 assert_eq!(p6.rank(), 6 * 5 + 1);
1997 assert_eq!(p6.coord(0), 1);
1998 assert_eq!(p6.coord(1), 0);
1999 assert_eq!(p6.coord(2), 1);
2000
2001 assert_eq!(extent.points().collect::<Vec<_>>().len(), 4 * 5 * 6);
2002 for (rank, point) in extent.points().enumerate() {
2003 let c = point.coords();
2004 let (x, y, z) = (c[0], c[1], c[2]);
2005 assert_eq!(z + y * 6 + x * 6 * 5, rank);
2006 assert_eq!(point.rank(), rank);
2007 }
2008 }
2009
2010 #[test]
2011 fn points_iterates_ranks_in_row_major_order() {
2012 let ext = extent!(x = 2, y = 3, z = 4); let mut it = ext.points();
2014
2015 for expected_rank in 0..ext.num_ranks() {
2016 let p = it.next().expect("expected another Point");
2017 assert_eq!(
2018 p.rank, expected_rank,
2019 "ranks must be consecutive in row-major order"
2020 );
2021 }
2022 assert!(
2023 it.next().is_none(),
2024 "iterator must be exhausted after num_ranks items"
2025 );
2026 }
2027
2028 #[test]
2029 fn points_iterates_single_point_for_0d_extent() {
2030 let ext = extent!();
2032 let mut it = ext.points();
2033
2034 let p = it
2035 .next()
2036 .expect("0-D extent should yield exactly one point");
2037 assert_eq!(p.rank, 0);
2038 assert_eq!(p.extent, ext);
2039
2040 assert!(
2041 it.next().is_none(),
2042 "no more points after the single 0-D point"
2043 );
2044 }
2045
2046 macro_rules! assert_view {
2047 ($view:expr, $extent:expr, $( $($coord:expr),+ => $rank:expr );* $(;)?) => {
2048 let view = $view;
2049 assert_eq!(view.extent(), $extent);
2050 let expected: Vec<_> = vec![$(($extent.point(vec![$($coord),+]).unwrap(), $rank)),*];
2051 let actual: Vec<_> = ViewExt::iter(&view).collect();
2052 assert_eq!(actual, expected);
2053 };
2054 }
2055
2056 #[test]
2057 fn test_view_basic() {
2058 let extent = extent!(x = 4, y = 4);
2059 assert_view!(
2060 extent.range("x", 0..2).unwrap(),
2061 extent!(x = 2, y = 4),
2062 0, 0 => 0;
2063 0, 1 => 1;
2064 0, 2 => 2;
2065 0, 3 => 3;
2066 1, 0 => 4;
2067 1, 1 => 5;
2068 1, 2 => 6;
2069 1, 3 => 7;
2070 );
2071 assert_view!(
2072 extent.range("x", 1).unwrap().range("y", 2..).unwrap(),
2073 extent!(x = 1, y = 2),
2074 0, 0 => 6;
2075 0, 1 => 7;
2076 );
2077 assert_view!(
2078 extent.range("y", Range(0, None, 2)).unwrap(),
2079 extent!(x = 4, y = 2),
2080 0, 0 => 0;
2081 0, 1 => 2;
2082 1, 0 => 4;
2083 1, 1 => 6;
2084 2, 0 => 8;
2085 2, 1 => 10;
2086 3, 0 => 12;
2087 3, 1 => 14;
2088 );
2089 assert_view!(
2090 extent.range("y", Range(0, None, 2)).unwrap().range("x", 2..).unwrap(),
2091 extent!(x = 2, y = 2),
2092 0, 0 => 8;
2093 0, 1 => 10;
2094 1, 0 => 12;
2095 1, 1 => 14;
2096 );
2097
2098 let extent = extent!(x = 10, y = 2);
2099 assert_view!(
2100 extent.range("x", Range(0, None, 2)).unwrap(),
2101 extent!(x = 5, y = 2),
2102 0, 0 => 0;
2103 0, 1 => 1;
2104 1, 0 => 4;
2105 1, 1 => 5;
2106 2, 0 => 8;
2107 2, 1 => 9;
2108 3, 0 => 12;
2109 3, 1 => 13;
2110 4, 0 => 16;
2111 4, 1 => 17;
2112 );
2113 assert_view!(
2114 extent.range("x", Range(0, None, 2)).unwrap().range("x", 2..).unwrap().range("y", 1).unwrap(),
2115 extent!(x = 3, y = 1),
2116 0, 0 => 9;
2117 1, 0 => 13;
2118 2, 0 => 17;
2119 );
2120
2121 let extent = extent!(zone = 4, host = 2, gpu = 8);
2122 assert_view!(
2123 extent.range("zone", 0).unwrap().range("gpu", Range(0, None, 2)).unwrap(),
2124 extent!(zone = 1, host = 2, gpu = 4),
2125 0, 0, 0 => 0;
2126 0, 0, 1 => 2;
2127 0, 0, 2 => 4;
2128 0, 0, 3 => 6;
2129 0, 1, 0 => 8;
2130 0, 1, 1 => 10;
2131 0, 1, 2 => 12;
2132 0, 1, 3 => 14;
2133 );
2134
2135 let extent = extent!(x = 3);
2136 assert_view!(
2137 extent.range("x", Range(0, None, 2)).unwrap(),
2138 extent!(x = 2),
2139 0 => 0;
2140 1 => 2;
2141 );
2142 }
2143
2144 #[test]
2145 fn test_point_indexing() {
2146 let extent = Extent::new(vec!["x".into(), "y".into(), "z".into()], vec![4, 5, 6]).unwrap();
2147 let point = extent.point(vec![1, 2, 3]).unwrap();
2148
2149 assert_eq!(point.coord(0), 1);
2150 assert_eq!(point.coord(1), 2);
2151 assert_eq!(point.coord(2), 3);
2152 }
2153
2154 #[test]
2155 #[should_panic]
2156 fn test_point_indexing_out_of_bounds() {
2157 let extent = Extent::new(vec!["x".into(), "y".into()], vec![4, 5]).unwrap();
2158 let point = extent.point(vec![1, 2]).unwrap();
2159
2160 let _ = point.coord(5); }
2162
2163 #[test]
2164 fn test_point_into_iter() {
2165 let extent = Extent::new(vec!["x".into(), "y".into(), "z".into()], vec![4, 5, 6]).unwrap();
2166 let point = extent.point(vec![1, 2, 3]).unwrap();
2167
2168 let coords: Vec<usize> = (&point).into_iter().collect();
2169 assert_eq!(coords, vec![1, 2, 3]);
2170
2171 let mut sum = 0;
2172 for coord in &point {
2173 sum += coord;
2174 }
2175 assert_eq!(sum, 6);
2176 }
2177
2178 #[test]
2179 fn test_extent_basic() {
2180 let extent = extent!(x = 10, y = 5, z = 1);
2181 assert_eq!(
2182 extent.iter().collect::<Vec<_>>(),
2183 vec![
2184 ("x".to_string(), 10),
2185 ("y".to_string(), 5),
2186 ("z".to_string(), 1)
2187 ]
2188 );
2189 }
2190
2191 #[test]
2192 fn test_extent_display() {
2193 let extent = Extent::new(vec!["x".into(), "y".into(), "z".into()], vec![4, 5, 6]).unwrap();
2194 assert_eq!(format!("{}", extent), "{x: 4, y: 5, z: 6}");
2195
2196 let extent = Extent::new(vec!["dim/0".into(), "dim/1".into()], vec![4, 5]).unwrap();
2197 assert_eq!(format!("{}", extent), "{\"dim/0\": 4, \"dim/1\": 5}");
2198
2199 let empty_extent = Extent::new(vec![], vec![]).unwrap();
2200 assert_eq!(format!("{}", empty_extent), "{}");
2201 }
2202
2203 #[test]
2204 fn extent_label_helpers() {
2205 let e = extent!(zone = 3, host = 2, gpu = 4);
2206 for (i, (lbl, sz)) in e.iter().enumerate() {
2207 assert_eq!(e.position(&lbl), Some(i));
2208 assert_eq!(e.size(&lbl), Some(sz));
2209 }
2210 assert_eq!(e.position("nope"), None);
2211 assert_eq!(e.size("nope"), None);
2212 }
2213
2214 #[test]
2215 fn test_extent_0d() {
2216 let e = Extent::new(vec![], vec![]).unwrap();
2217 assert_eq!(e.num_ranks(), 1);
2218
2219 let points: Vec<_> = e.points().collect();
2220 assert_eq!(points.len(), 1);
2221 assert_eq!(points[0].coords(), &[] as &[usize]);
2222 assert_eq!(points[0].rank(), 0);
2223
2224 let mut it = (&points[0]).into_iter();
2226 assert_eq!(it.len(), 0);
2227 assert!(it.next().is_none()); assert!(it.next().is_none()); }
2230
2231 #[test]
2232 fn test_extent_concat() {
2233 let extent1 = extent!(x = 2, y = 3);
2235 let extent2 = extent!(z = 4, w = 5);
2236
2237 let result = extent1.concat(&extent2).unwrap();
2238 assert_eq!(result.labels(), &["x", "y", "z", "w"]);
2239 assert_eq!(result.sizes(), &[2, 3, 4, 5]);
2240 assert_eq!(result.num_ranks(), 2 * 3 * 4 * 5);
2241
2242 let empty = extent!();
2244 let result = extent1.concat(&empty).unwrap();
2245 assert_eq!(result.labels(), &["x", "y"]);
2246 assert_eq!(result.sizes(), &[2, 3]);
2247
2248 let result = empty.concat(&extent1).unwrap();
2249 assert_eq!(result.labels(), &["x", "y"]);
2250 assert_eq!(result.sizes(), &[2, 3]);
2251
2252 let result = empty.concat(&empty).unwrap();
2254 assert_eq!(result.labels(), &[] as &[String]);
2255 assert_eq!(result.sizes(), &[] as &[usize]);
2256 assert_eq!(result.num_ranks(), 1); let result = extent1.concat(&extent1);
2260 assert!(
2261 result.is_err(),
2262 "Self-concatenation should error due to overlapping labels"
2263 );
2264 match result.unwrap_err() {
2265 ExtentError::OverlappingLabel { label } => {
2266 assert!(label == "x"); }
2268 other => panic!("Expected OverlappingLabel error, got {:?}", other),
2269 }
2270
2271 let result = extent1.concat(&extent2).unwrap();
2273 let point = result.point(vec![1, 2, 3, 4]).unwrap();
2274 assert_eq!(point.coords(), vec![1, 2, 3, 4]);
2275 assert_eq!(point.extent(), &result);
2276
2277 let extent_a = extent!(x = 2, y = 3);
2279 let extent_b = extent!(y = 3, z = 4); let result = extent_a.concat(&extent_b);
2281 assert!(
2282 result.is_err(),
2283 "Should error on overlapping labels even with same size"
2284 );
2285 match result.unwrap_err() {
2286 ExtentError::OverlappingLabel { label } => {
2287 assert_eq!(label, "y"); }
2289 other => panic!("Expected OverlappingLabel error, got {:?}", other),
2290 }
2291
2292 let extent_x = extent!(x = 2, y = 3);
2294 let extent_y = extent!(z = 4);
2295 assert_eq!(
2296 extent_x.concat(&extent_y).unwrap().labels(),
2297 &["x", "y", "z"]
2298 );
2299 assert_eq!(
2300 extent_y.concat(&extent_x).unwrap().labels(),
2301 &["z", "x", "y"]
2302 );
2303
2304 let extent_m = extent!(x = 2);
2306 let extent_n = extent!(y = 3);
2307 let extent_o = extent!(z = 4);
2308
2309 let left_assoc = extent_m
2310 .concat(&extent_n)
2311 .unwrap()
2312 .concat(&extent_o)
2313 .unwrap();
2314 let right_assoc = extent_m
2315 .concat(&extent_n.concat(&extent_o).unwrap())
2316 .unwrap();
2317
2318 assert_eq!(left_assoc, right_assoc);
2319 assert_eq!(left_assoc.labels(), &["x", "y", "z"]);
2320 assert_eq!(left_assoc.sizes(), &[2, 3, 4]);
2321 assert_eq!(left_assoc.num_ranks(), 2 * 3 * 4);
2322 }
2323
2324 #[test]
2325 fn extent_unity_equiv_to_0d() {
2326 let e = Extent::unity();
2327 assert!(e.is_empty());
2328 assert_eq!(e.num_ranks(), 1);
2329 let pts: Vec<_> = e.points().collect();
2330 assert_eq!(pts.len(), 1);
2331 assert_eq!(pts[0].rank(), 0);
2332 assert!(pts[0].coords().is_empty());
2333 }
2334
2335 #[test]
2336 fn test_point_display() {
2337 let extent = Extent::new(vec!["x".into(), "y".into(), "z".into()], vec![4, 5, 6]).unwrap();
2338 let point = extent.point(vec![1, 2, 3]).unwrap();
2339 assert_eq!(format!("{}", point), "x=1/4,y=2/5,z=3/6");
2340
2341 assert!(extent.point(vec![]).is_err());
2342
2343 let empty_extent = Extent::new(vec![], vec![]).unwrap();
2344 let empty_point = empty_extent.point(vec![]).unwrap();
2345 assert_eq!(format!("{}", empty_point), "");
2346 }
2347
2348 #[test]
2349 fn test_point_display_with_quoted_labels() {
2350 let ext = Extent::new(vec!["dim/0".into(), "dim,1".into()], vec![3, 5]).unwrap();
2352
2353 assert_eq!(format!("{}", ext), "{\"dim/0\": 3, \"dim,1\": 5}");
2355
2356 let p = ext.point(vec![1, 2]).unwrap();
2358 assert_eq!(format!("{}", p), "\"dim/0\"=1/3,\"dim,1\"=2/5");
2359 }
2360
2361 #[test]
2362 fn test_relative_point() {
2363 pub fn relative_point(rank_on_root_mesh: usize, shape: &Shape) -> anyhow::Result<Point> {
2366 let coords = shape.slice().coordinates(rank_on_root_mesh)?;
2367 let extent = Extent::new(shape.labels().to_vec(), shape.slice().sizes().to_vec())?;
2368 Ok(extent.point(coords)?)
2369 }
2370
2371 let root_shape = shape! { replicas = 4, hosts = 4, gpus = 4 };
2372 let sliced_shape = root_shape
2384 .select("replicas", crate::Range(0, Some(4), 3))
2385 .unwrap()
2386 .select("hosts", crate::Range(1, Some(4), 2))
2387 .unwrap()
2388 .select("gpus", crate::Range(0, Some(4), 2))
2389 .unwrap();
2390 let ranks_on_root_mesh = &[4, 6, 12, 14, 52, 54, 60, 62];
2391 assert_eq!(
2392 sliced_shape.slice().iter().collect::<Vec<_>>(),
2393 ranks_on_root_mesh,
2394 );
2395
2396 let ranks_on_sliced_mesh = ranks_on_root_mesh
2397 .iter()
2398 .map(|&r| relative_point(r, &sliced_shape).unwrap().rank());
2399 assert_eq!(
2400 ranks_on_sliced_mesh.collect::<Vec<_>>(),
2401 vec![0, 1, 2, 3, 4, 5, 6, 7]
2402 );
2403 }
2404
2405 #[test]
2406 fn test_iter_subviews() {
2407 let extent = extent!(zone = 4, host = 4, gpu = 8);
2408
2409 assert_eq!(extent.group_by("gpu").unwrap().count(), 16);
2410 assert_eq!(extent.group_by("zone").unwrap().count(), 1);
2411
2412 let mut parts = extent.group_by("gpu").unwrap();
2413 assert_view!(
2414 parts.next().unwrap(),
2415 extent!(gpu = 8),
2416 0 => 0;
2417 1 => 1;
2418 2 => 2;
2419 3 => 3;
2420 4 => 4;
2421 5 => 5;
2422 6 => 6;
2423 7 => 7;
2424 );
2425 assert_view!(
2426 parts.next().unwrap(),
2427 extent!(gpu = 8),
2428 0 => 8;
2429 1 => 9;
2430 2 => 10;
2431 3 => 11;
2432 4 => 12;
2433 5 => 13;
2434 6 => 14;
2435 7 => 15;
2436 );
2437 }
2438
2439 #[test]
2440 fn test_view_values() {
2441 let extent = extent!(x = 4, y = 4);
2442 assert_eq!(
2443 extent.values().collect::<Vec<_>>(),
2444 (0..16).collect::<Vec<_>>()
2445 );
2446 let region = extent.range("y", 1).unwrap();
2447 assert_eq!(region.values().collect::<Vec<_>>(), vec![1, 5, 9, 13]);
2448 }
2449
2450 #[test]
2451 fn region_is_subset_algebra() {
2452 let e = extent!(x = 5, y = 4);
2453 let a = e.range("x", 1..4).unwrap(); let b = a.range("y", 1..3).unwrap(); let c = e.range("x", 0..2).unwrap(); assert!(b.region().is_subset(&a.region()));
2458 assert!(b.region().is_subset(&e.region()));
2459 assert!(a.region().is_subset(&e.region()));
2460
2461 assert!(!c.region().is_subset(&a.region()));
2462 assert!(c.region().is_subset(&e.region()));
2463 }
2464
2465 #[test]
2466 fn test_remap() {
2467 let region: Region = extent!(x = 4, y = 4).into();
2468 assert_eq!(
2470 region.remap(®ion).unwrap().collect::<Vec<_>>(),
2471 (0..16).collect::<Vec<_>>()
2472 );
2473
2474 let subset = region.range("x", 2..).unwrap();
2475 assert_eq!(subset.num_ranks(), 8);
2476 assert_eq!(
2477 region.remap(&subset).unwrap().collect::<Vec<_>>(),
2478 vec![8, 9, 10, 11, 12, 13, 14, 15],
2479 );
2480
2481 let subset = subset.range("y", 1).unwrap();
2482 assert_eq!(subset.num_ranks(), 2);
2483 assert_eq!(
2484 region.remap(&subset).unwrap().collect::<Vec<_>>(),
2485 vec![9, 13],
2486 );
2487
2488 let ext = extent!(replica = 8, gpu = 4);
2491 let replica1 = ext.range("replica", 1).unwrap();
2492 assert_eq!(replica1.extent(), extent!(replica = 1, gpu = 4));
2493 let replica1_gpu12 = replica1.range("gpu", 1..3).unwrap();
2494 assert_eq!(replica1_gpu12.extent(), extent!(replica = 1, gpu = 2));
2495 assert_eq!(
2497 replica1.remap(&replica1_gpu12).unwrap().collect::<Vec<_>>(),
2498 vec![1, 2],
2499 );
2500 }
2501
2502 #[test]
2503 fn test_base_local_rank_conversion() {
2504 fn point(rank: usize, region: &Region) -> Point {
2505 region.extent().point_of_rank(rank).unwrap()
2506 }
2507
2508 let extent = extent!(replicas = 4, gpus = 2);
2509 let region = extent.range("replicas", 1..3).unwrap();
2510 assert!(
2518 region
2519 .base_rank_of_point(extent.point_of_rank(0).unwrap())
2520 .is_err()
2521 );
2522 assert_eq!(region.base_rank_of_point(point(0, ®ion)).unwrap(), 2);
2524 assert_eq!(region.base_rank_of_point(point(1, ®ion)).unwrap(), 3);
2525 assert_eq!(region.base_rank_of_point(point(2, ®ion)).unwrap(), 4);
2526 assert_eq!(region.base_rank_of_point(point(3, ®ion)).unwrap(), 5);
2527 assert_eq!(region.point_of_base_rank(2).unwrap(), point(0, ®ion));
2529 assert_eq!(region.point_of_base_rank(3).unwrap(), point(1, ®ion));
2530 assert_eq!(region.point_of_base_rank(4).unwrap(), point(2, ®ion));
2531 assert_eq!(region.point_of_base_rank(5).unwrap(), point(3, ®ion));
2532 assert!(region.point_of_base_rank(1).is_err());
2534 assert!(region.point_of_base_rank(6).is_err());
2535
2536 let subset = region
2538 .range("replicas", 1..2)
2539 .unwrap()
2540 .range("gpus", 1..2)
2541 .unwrap();
2542 assert_eq!(subset.base_rank_of_point(point(0, &subset)).unwrap(), 5);
2549 assert_eq!(subset.point_of_base_rank(5).unwrap(), point(0, &subset));
2550 assert!(subset.point_of_base_rank(4).is_err());
2552 assert!(subset.point_of_base_rank(6).is_err());
2553 }
2554
2555 use proptest::prelude::*;
2556
2557 use crate::strategy::gen_extent;
2558 use crate::strategy::gen_region;
2559 use crate::strategy::gen_region_strided;
2560
2561 proptest! {
2562 #[test]
2563 fn test_region_parser(region in gen_region(1..=5, 1024)) {
2564 assert_eq!(
2566 region,
2567 region.to_string().parse::<Region>().unwrap(),
2568 "failed to roundtrip region {}", region
2569 );
2570 }
2571 }
2572
2573 proptest! {
2585 #[test]
2586 fn region_parser_with_offset_roundtrips(region in gen_region(1..=4, 8)) {
2587 let (labels, slice) = region.clone().into_inner();
2588 let region_off = Region {
2589 labels,
2590 slice: Slice::new(8, slice.sizes().to_vec(), slice.strides().to_vec()).unwrap(),
2591 };
2592 let s = region_off.to_string();
2593 let parsed: Region = s.parse().unwrap();
2594 prop_assert_eq!(parsed, region_off);
2595 }
2596 }
2597
2598 proptest! {
2609 #[test]
2610 fn region_strided_display_parse_roundtrips(
2611 region in gen_region_strided(1..=4, 6, 3, 16)
2612 ) {
2613 let s = region.to_string();
2624 let parsed: Region = s.parse().unwrap();
2625 prop_assert_eq!(parsed, region);
2626 }
2627 }
2628
2629 proptest! {
2640 #[test]
2641 fn region_strided_display_matches_slice(
2642 region in gen_region_strided(1..=4, 6, 3, 16)
2643 ) {
2644 let s = region.to_string();
2645 let slice = region.slice();
2646
2647 if slice.offset() != 0 {
2649 let prefix: Vec<_> = s.split('+').collect();
2650 prop_assert!(prefix.len() > 1, "expected offset+ form in {}", s);
2651 let offset_str = prefix[0];
2652 let offset_val: usize = offset_str.parse().unwrap();
2653 prop_assert_eq!(offset_val, slice.offset(), "offset mismatch in {}", s);
2654 } else {
2655 prop_assert!(!s.contains('+'), "unexpected +offset in {}", s);
2656 }
2657
2658 let body = s.split('+').next_back().unwrap(); let parts: Vec<_> = body.split(',').collect();
2661 prop_assert_eq!(parts.len(), slice.sizes().len());
2662
2663 for (i, part) in parts.iter().enumerate() {
2664 let rhs = part.split('=').nth(1).unwrap();
2666 let mut nums = rhs.split('/');
2667 let size_val: usize = nums.next().unwrap().parse().unwrap();
2668 let stride_val: usize = nums.next().unwrap().parse().unwrap();
2669
2670 prop_assert_eq!(size_val, slice.sizes()[i], "size mismatch at dim {} in {}", i, s);
2671 prop_assert_eq!(stride_val, slice.strides()[i], "stride mismatch at dim {} in {}", i, s);
2672 }
2673 }
2674 }
2675
2676 #[test]
2677 fn test_point_from_str_round_trip() {
2678 let points = vec![
2679 extent!(x = 4, y = 5, z = 6).point(vec![1, 2, 3]).unwrap(),
2680 extent!(host = 2, gpu = 8).point(vec![0, 7]).unwrap(),
2681 extent!().point(vec![]).unwrap(),
2682 extent!(x = 10).point(vec![5]).unwrap(),
2683 ];
2684
2685 for point in points {
2686 assert_eq!(point, point.to_string().parse().unwrap());
2687 }
2688 }
2689
2690 #[test]
2691 fn test_point_from_str_basic() {
2692 let cases = vec![
2693 ("x=1/4,y=2/5", extent!(x = 4, y = 5), vec![1, 2]),
2694 ("host=0/2,gpu=7/8", extent!(host = 2, gpu = 8), vec![0, 7]),
2695 ("z=3/6", extent!(z = 6), vec![3]),
2696 ("", extent!(), vec![]), (" x = 1 / 4 , y = 2 / 5 ", extent!(x = 4, y = 5), vec![1, 2]),
2699 ];
2700
2701 for (input, expected_extent, expected_coords) in cases {
2702 let parsed: Point = input.parse().unwrap();
2703 let expected = expected_extent.point(expected_coords).unwrap();
2704 assert_eq!(parsed, expected, "failed to parse: {}", input);
2705 }
2706 }
2707
2708 #[test]
2709 fn test_point_from_str_quoted() {
2710 let extent = Extent::new(vec!["dim/0".into(), "dim,1".into()], vec![3, 5]).unwrap();
2712 let point = extent.point(vec![1, 2]).unwrap();
2713
2714 let display_str = point.to_string();
2715 assert_eq!(display_str, "\"dim/0\"=1/3,\"dim,1\"=2/5");
2716
2717 let parsed: Point = display_str.parse().unwrap();
2718 assert_eq!(parsed, point);
2719
2720 let parsed: Point = "\"dim/0\"=1/3,\"dim,1\"=2/5".parse().unwrap();
2721 assert_eq!(parsed, point);
2722 }
2723
2724 #[test]
2725 fn test_point_from_str_error_cases() {
2726 let error_cases = vec![
2728 "x=1,y=2/5", "x=1/4,y=2", "x=1/4,y=/5", "x=/4,y=2/5", "x=1/4,y=2/", "x=1/,y=2/5", "x=1/4=5,y=2/5", "x=1/4/6,y=2/5", "x=abc/4,y=2/5", "x=1/abc,y=2/5", "=1/4,y=2/5", "x=1/4,", "x=1/4,=2/5", "x=1/4,y", "x", "x=", "x=1/4,y=10/5", ];
2746
2747 for input in error_cases {
2748 let result: Result<Point, PointError> = input.parse();
2749 assert!(result.is_err(), "Expected error for input: '{}'", input);
2750 }
2751 }
2752
2753 #[test]
2754 fn test_point_from_str_coordinate_validation() {
2755 let input = "x=5/4,y=2/5"; let result: Result<Point, PointError> = input.parse();
2758 assert!(
2759 result.is_err(),
2760 "Expected error for out-of-bounds coordinate"
2761 );
2762
2763 match result.unwrap_err() {
2764 PointError::OutOfRangeIndex { size, index } => {
2765 assert_eq!(size, 4);
2766 assert_eq!(index, 5);
2767 }
2768 _ => panic!("Expected OutOfRangeIndex error"),
2769 }
2770 }
2771
2772 #[test]
2773 fn test_point_from_str_consistency_validation() {
2774 let input = "x=1/4,y=2/5,z=3/6";
2779 let parsed: Point = input.parse().unwrap();
2780
2781 assert_eq!(parsed.extent().labels(), &["x", "y", "z"]);
2782 assert_eq!(parsed.extent().sizes(), &[4, 5, 6]);
2783 assert_eq!(parsed.coords(), vec![1, 2, 3]);
2784 }
2785
2786 proptest! {
2787 #[test]
2790 fn point_coord_and_iter_agree(extent in gen_extent(0..=4, 8)) {
2791 for p in extent.points() {
2792 let via_coords = p.coords();
2793 let via_into_iter: Vec<_> = (&p).into_iter().collect();
2794 prop_assert_eq!(via_into_iter, via_coords.clone(), "coord_iter mismatch for {}", p);
2795
2796 for (i, &coord) in via_coords.iter().enumerate() {
2797 prop_assert_eq!(p.coord(i), coord, "coord(i) mismatch at axis {} for {}", i, p);
2798 }
2799 }
2800 }
2801
2802 #[test]
2804 fn points_count_matches_num_ranks(extent in gen_extent(0..=4, 8)) {
2805 let c = extent.points().count();
2806 prop_assert_eq!(c, extent.num_ranks(), "count {} != num_ranks {}", c, extent.num_ranks());
2807 }
2808
2809 #[test]
2813 fn coord_iter_exact_size_invariants(extent in gen_extent(0..=4, 8)) {
2814 for p in extent.points() {
2815 let mut it = (&p).into_iter();
2816
2817 let mut remaining = p.len();
2820 prop_assert_eq!(it.len(), remaining);
2821 prop_assert_eq!(it.size_hint(), (remaining, Some(remaining)));
2822
2823 let mut yielded = Vec::with_capacity(remaining);
2825
2826 while let Some(v) = it.next() {
2829 yielded.push(v);
2830 remaining -= 1;
2831 prop_assert_eq!(it.len(), remaining);
2832 prop_assert_eq!(it.size_hint(), (remaining, Some(remaining)));
2833 }
2834
2835 prop_assert_eq!(remaining, 0);
2838 prop_assert!(it.next().is_none());
2839 prop_assert!(it.next().is_none());
2840
2841 prop_assert_eq!(yielded, p.coords());
2843 }
2844 }
2845
2846 #[test]
2850 fn rank_of_coords_dim_mismatch(extent in gen_extent(0..=4, 8)) {
2851 let want = extent.len();
2852 let wrong = if want == 0 { 1 } else { want - 1 };
2854 let bad = vec![0usize; wrong];
2855
2856 match extent.rank_of_coords(&bad).unwrap_err() {
2857 PointError::DimMismatch { expected, actual } => {
2858 prop_assert_eq!(expected, want, "expected len mismatch");
2859 prop_assert_eq!(actual, wrong, "actual len mismatch");
2860 }
2861 other => prop_assert!(false, "expected DimMismatch, got {:?}", other),
2862 }
2863 }
2864
2865 #[test]
2870 fn rank_of_coords_out_of_range_index(extent in gen_extent(1..=4, 8)) {
2871 let sizes = extent.sizes().to_vec();
2873 let mut coords = vec![0usize; sizes.len()];
2875 let axis = 0usize;
2877 coords[axis] = sizes[axis];
2878
2879 match extent.rank_of_coords(&coords).unwrap_err() {
2880 PointError::OutOfRangeIndex { size, index } => {
2881 prop_assert_eq!(size, sizes[axis], "reported size mismatch");
2882 prop_assert_eq!(index, sizes[axis], "reported index mismatch");
2883 }
2884 other => prop_assert!(false, "expected OutOfRangeIndex, got {:?}", other),
2885 }
2886 }
2887
2888 #[test]
2891 fn point_of_rank_out_of_range(extent in gen_extent(0..=4, 8)) {
2892 let total = extent.num_ranks(); match extent.point_of_rank(total).unwrap_err() {
2894 PointError::OutOfRangeRank { total: t, rank: r } => {
2895 prop_assert_eq!(t, total, "reported total mismatch");
2896 prop_assert_eq!(r, total, "reported rank mismatch");
2897 }
2898 other => prop_assert!(false, "expected OutOfRangeRank, got {:?}", other),
2899 }
2900 }
2901
2902 #[test]
2904 fn point_display_parse_round_trip(extent in gen_extent(0..=4, 8)) {
2905 for point in extent.points() {
2906 let display = point.to_string();
2907 let parsed: Point = display.parse().unwrap();
2908 prop_assert_eq!(parsed, point, "round-trip failed for point: {}", display);
2909 }
2910 }
2911 }
2912
2913 proptest! {
2914 #[test]
2917 fn points_iterator_equivalent_to_legacy_cartesian(extent in gen_extent(0..=4, 8)) {
2918 let sizes = extent.sizes().to_vec();
2919
2920 let legacy_len = CartesianIterator::new(sizes.clone()).count();
2922 prop_assert_eq!(legacy_len, extent.num_ranks());
2923
2924 let legacy_coords = CartesianIterator::new(sizes);
2927 for (step, (p, coords)) in extent.points().zip(legacy_coords).enumerate() {
2928 let old_rank = extent.rank_of_coords(&coords).expect("valid legacy coords");
2929 prop_assert_eq!(p.rank(), old_rank, "rank mismatch at step {} for coords {:?}", step, coords);
2930 prop_assert_eq!(p.coords(), coords, "coords mismatch at step {}", step);
2931 }
2932 }
2933 }
2934}