Skip to main content

ndslice/
view.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! The purpose of this module is to provide data structures to efficiently
10//! describe subsets of nonnegative integers, called *ranks*, represented by
11//! `usize` values. Ranks are organized into multi-dimensional spaces in which
12//! each discrete point is a rank, mapped in row-major order.
13//!
14//! These dimensions represent a space that carry semantic meaning, such that
15//! ranks are usually sub-set along some dimension of the space. For example,
16//! ranks may be organized into a space with dimensions "replica", "host", "gpu",
17//! and we'd expect to subset along these dimensions, for example to select all
18//! GPUs in a given replica.
19//!
20//! This alignment helps provide a simple and efficient representation, internally
21//! in the form of [`crate::Slice`], comprising an offset, sizes, and strides that
22//! index into the space.
23//!
24//! - [`Extent`]: the *shape* of the space, naming each dimension and specifying
25//!               their sizes.
26//! - [`Point`]: a specific coordinate in an extent, together with its linearized rank.
27//! - [`Region`]: a (possibly sparse) hyper-rectangle of ranks within a larger extent.
28//!               Since it is always rectangular, it also defines its own extent.
29//! - [`View`]: a collection of items indexed by [`Region`]. Views provide standard
30//!             manipulation operations and use ranks as an efficient indexing scheme.
31
32use 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; // exclusively for `impl From<Shape> for Extent`
45use crate::Slice;
46use crate::SliceError;
47use crate::SliceIterator;
48use crate::parse::Parser;
49use crate::parse::ParserError;
50
51/// Errors that can occur when constructing or validating an `Extent`.
52#[derive(Debug, thiserror::Error)]
53pub enum ExtentError {
54    /// The number of labels does not match the number of sizes.
55    ///
56    /// This occurs when constructing an `Extent` from parallel
57    /// `Vec<String>` and `Vec<usize>` inputs that are not the same
58    /// length.
59    #[error("label/sizes dimension mismatch: {num_labels} != {num_sizes}")]
60    DimMismatch {
61        /// Number of dimension labels provided.
62        num_labels: usize,
63        /// Number of dimension sizes provided.
64        num_sizes: usize,
65    },
66
67    /// An overlapping label was found.
68    ///
69    /// This occurs when attempting to combine extents that
70    /// share one or more dimension labels, which is not allowed.
71    #[error("overlapping label found: {label}")]
72    OverlappingLabel {
73        /// The label that appears in both extents.
74        label: String,
75    },
76}
77
78/// `Extent` defines the logical shape of a multidimensional space by
79/// assigning a size to each named dimension. It abstracts away memory
80/// layout and focuses solely on structure — what dimensions exist and
81/// how many elements each contains.
82///
83/// Conceptually, it corresponds to a coordinate space in the
84/// mathematical sense.
85#[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// `ExtentData` is represented as:
97// - `labels`: dimension names like `"zone"`, `"host"`, `"gpu"`
98// - `sizes`: number of elements in each dimension, independent of
99//   stride or storage layout
100#[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        // Safe: Shape guarantees labels.len() == sizes.len().
109        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    /// Creates a new `Extent` from the given labels and sizes.
121    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    /// Returns the ordered list of dimension labels in this extent.
139    pub fn labels(&self) -> &[String] {
140        &self.inner.labels
141    }
142
143    /// Returns the dimension sizes, ordered to match the labels.
144    pub fn sizes(&self) -> &[usize] {
145        &self.inner.sizes
146    }
147
148    /// Returns the size of the dimension with the given label, if it
149    /// exists.
150    pub fn size(&self, label: &str) -> Option<usize> {
151        self.position(label).map(|pos| self.sizes()[pos])
152    }
153
154    /// Returns the position of the dimension with the given label, if
155    /// it exists exists.
156    pub fn position(&self, label: &str) -> Option<usize> {
157        self.labels().iter().position(|l| l == label)
158    }
159
160    // Computes the row-major logical rank of the given coordinates
161    // in this extent.
162    //
163    // ```text
164    // Σ (coord[i] × ∏(sizes[j] for j > i))
165    // ```
166    //
167    // where 'coord' is the point's coordinate and 'sizes' is the
168    // extent's dimension sizes.
169    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    /// Creates a [`Point`] in this extent from the given coordinate
190    /// vector.
191    ///
192    /// The coordinates are interpreted in **row-major** order against
193    /// `self.sizes()`. This constructor does not store the
194    /// coordinates; it computes the linear **rank** and returns a
195    /// `Point` that stores `{ rank, extent }`.
196    ///
197    /// # Errors
198    ///
199    /// Returns:
200    /// - [`PointError::DimMismatch`] if `coords.len() != self.len()`.
201    /// - [`PointError::OutOfRangeIndex`] if any coordinate `coords[i]
202    ///   >= self.sizes()[i]`.
203    ///
204    /// # Examples
205    ///
206    /// ```
207    /// use ndslice::extent;
208    ///
209    /// let ext = extent!(x = 2, y = 3, z = 4);
210    /// let p = ext.point(vec![1, 2, 3]).unwrap();
211    /// assert_eq!(p.rank(), 1 * (3 * 4) + 2 * 4 + 3); // row-major
212    /// assert_eq!(p.coords(), vec![1, 2, 3]);
213    /// ```
214    ///
215    /// Dimension mismatch:
216    /// ```
217    /// use ndslice::PointError;
218    /// use ndslice::extent;
219    ///
220    /// let ext = extent!(x = 2, y = 3);
221    /// let err = ext.point(vec![1]).unwrap_err();
222    /// matches!(err, PointError::DimMismatch { .. });
223    /// ```
224    ///
225    /// Coordinate out of range:
226    /// ```
227    /// use ndslice::PointError;
228    /// use ndslice::extent;
229    ///
230    /// let ext = extent!(x = 2, y = 3);
231    /// let err = ext.point(vec![1, 3]).unwrap_err(); // y size is 3, max index is 2
232    /// matches!(err, PointError::OutOfRangeIndex { .. });
233    /// ```
234    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    /// Returns the [`Point`] corresponding to the given linearized
242    /// `rank` within this extent, using row-major order.
243    ///
244    /// # Errors
245    ///
246    /// Returns [`PointError::OutOfRangeRank`] if `rank >=
247    /// self.num_ranks()`, i.e. when the requested rank lies outside
248    /// the bounds of this extent.
249    ///
250    /// # Examples
251    ///
252    /// ```
253    /// use ndslice::extent;
254    ///
255    /// let ext = extent!(x = 2, y = 3);
256    /// assert_eq!(ext.num_ranks(), 6);
257    ///
258    /// let p = ext.point_of_rank(4).unwrap();
259    /// assert_eq!(p.coords(), vec![1, 1]); // row-major: x=1, y=1
260    /// assert_eq!(p.rank(), 4);
261    ///
262    /// assert!(ext.point_of_rank(6).is_err()); // out of range
263    /// ```
264    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    /// Returns the number of dimensions in this extent.
276    ///
277    /// For example, an extent defined as `(x=2, y=3, z=4)` has
278    /// dimensionality 3.
279    pub fn len(&self) -> usize {
280        self.sizes().len()
281    }
282
283    /// Returns true if this extent has zero dimensions.
284    ///
285    /// A 0-dimensional extent corresponds to the scalar case: a
286    /// coordinate space with exactly one rank (the empty tuple `[]`).
287    pub fn is_empty(&self) -> bool {
288        self.sizes().is_empty()
289    }
290
291    /// Returns the total number of ranks (points) in this extent.
292    ///
293    /// This is the product of all dimension sizes, i.e. the number of
294    /// distinct coordinates in row-major order.
295    pub fn num_ranks(&self) -> usize {
296        self.sizes().iter().product()
297    }
298
299    /// Convert this extent into its labels and sizes.
300    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    /// Creates a slice representing the full extent.
308    pub fn to_slice(&self) -> Slice {
309        Slice::new_row_major(self.sizes())
310    }
311
312    /// Iterate over this extent's labels and sizes.
313    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    /// Iterate points in this extent.
321    pub fn points(&self) -> ExtentPointsIterator<'_> {
322        ExtentPointsIterator::new(self)
323    }
324
325    /// Append the dimensions of `other` to this extent, preserving order.
326    ///
327    /// Duplicate labels are not allowed: if any label in `other` already appears
328    /// in `self`, this returns `ExtentError::OverlappingLabel`.
329    ///
330    /// This operation is not commutative: `a.concat(&b)` may differ from
331    /// `b.concat(&a)`.
332    pub fn concat(&self, other: &Extent) -> Result<Self, ExtentError> {
333        use std::collections::HashSet;
334        // Check for any overlapping labels in linear time using hash set
335        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        // Combine labels and sizes from both extents with pre-allocated memory
340        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
350/// Label formatting utilities shared across `Extent`, `Region`, and
351/// `Point`.
352///
353/// - [`is_safe_ident`] determines whether a label can be printed
354///   bare, i.e. consists only of `[A-Za-z0-9_]+`.
355/// - [`fmt_label`] returns the label unchanged if it is safe,
356///   otherwise quotes it using Rust string-literal syntax (via
357///   `format!("{:?}", s)`).
358///
359/// This ensures a consistent, unambiguous display format across all
360/// types.
361mod labels {
362    /// A "safe" identifier consists only of ASCII alphanumeric chars
363    /// or underscores (`[A-Za-z0-9_]+`). These can be displayed
364    /// without quotes.
365    pub(super) fn is_safe_ident(s: &str) -> bool {
366        s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
367    }
368
369    /// Render a label according to the quoting rule:
370    /// - Safe identifiers are returned as-is.
371    /// - Otherwise the label is quoted using Rust string literal
372    ///   syntax (via `format!("{:?}", s)`).
373    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
382/// Formats an `Extent` as a compact map‐literal:
383/// ```text
384/// {label: size, label: size, ...}
385/// ```
386/// # Grammar
387///
388/// ```text
389/// Extent   := "{" [ Pair ( "," Pair )* ] "}"
390/// Pair     := Label ": " Size
391/// Label    := SafeIdent | Quoted
392/// SafeIdent:= [A-Za-z0-9_]+
393/// Quoted   := "\"" ( [^"\\] | "\\" . )* "\""
394/// Size     := [0-9]+
395/// ```
396///
397/// # Quoting rules
398///
399/// - Labels that are **not** `SafeIdent` are rendered using Rust
400///   string literal syntax (via `format!("{:?}", label)`), e.g.
401///   `"dim/0"` or `"x y"`.
402/// - "Safe" means ASCII alphanumeric or underscore (`[A-Za-z0-9_]+`).
403///   Everything else is quoted. This keeps common identifiers
404///   unquoted and unambiguous.
405///
406/// # Examples
407///
408/// ```text
409/// {x: 4, y: 5, z: 6}
410/// {"dim/0": 3, "dim,1": 5}
411/// {}
412/// ```
413///
414/// Implementation note: label rendering goes through `fmt_label`,
415/// which emits the label as-is for safe idents, otherwise as a Rust
416/// string literal.
417impl 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
435/// An iterator for points in an extent.
436pub 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    /// Advances the iterator and returns the next [`Point`] in
454    /// row-major order.
455    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/// Errors that can occur when constructing or evaluating a `Point`.
470#[derive(Debug, Error)]
471pub enum PointError {
472    /// The number of coordinates does not match the number of
473    /// dimensions defined by the associated extent.
474    ///
475    /// This occurs when creating a `Point` with a coordinate vector
476    /// of incorrect length relative to the dimensionality of the
477    /// extent.
478    #[error("dimension mismatch: expected {expected}, got {actual}")]
479    DimMismatch {
480        /// Number of dimensions expected from the extent.
481        expected: usize,
482        /// Number of coordinates actually provided.
483        actual: usize,
484    },
485
486    /// The provided rank is outside the valid range for the extent.
487    ///
488    /// Ranks are the linearized row-major indices of all points in
489    /// the extent, spanning the half-open interval `[0, total)`. This
490    /// error occurs when a rank greater than or equal to `total` is
491    /// requested.
492    #[error("out of range: total ranks {total}; does not contain rank {rank}")]
493    OutOfRangeRank {
494        /// The total number of valid ranks in the extent.
495        total: usize,
496        /// The rank that was requested but not valid.
497        rank: usize,
498    },
499
500    /// A coordinate index is outside the valid range for its
501    /// dimension.
502    ///
503    /// Each dimension of an extent has a size `size`, with valid
504    /// indices spanning the half-open interval `[0, size)`. This
505    /// error occurs when a coordinate `index` is greater than or
506    /// equal to `size`.
507    #[error("out of range: dim size {size}; does not contain index {index}")]
508    OutOfRangeIndex {
509        /// The size of the offending dimension.
510        size: usize,
511        /// The invalid coordinate index that was requested.
512        index: usize,
513    },
514
515    /// Failed to parse a point from a string.
516    #[error("failed to parse point: {reason}")]
517    ParseError { reason: String },
518}
519
520/// `Point` represents a specific coordinate within the
521/// multi-dimensional space defined by an [`Extent`].
522///
523/// A `Point` can be viewed in two equivalent ways:
524/// - **Coordinates**: a tuple of indices, one per dimension,
525///   retrievable with [`Point::coord`] and [`Point::coords`].
526/// - **Rank**: a single linearized index into the extent's row-major
527///   ordering, retrievable with [`Point::rank`].
528///
529/// Internally, a `Point` stores:
530/// - A `rank`: the row-major linearized index of this point.
531/// - An `extent`: the extent that defines its dimensionality and
532///   sizes.
533///
534/// These fields are private; use the accessor methods instead.
535///
536/// # Examples
537///
538/// ```
539/// use ndslice::extent;
540///
541/// let ext = extent!(zone = 2, host = 4, gpu = 8);
542/// let point = ext.point(vec![1, 2, 3]).unwrap();
543///
544/// // Coordinate-based access
545/// assert_eq!(point.coord(0), 1);
546/// assert_eq!(point.coord(1), 2);
547/// assert_eq!(point.coord(2), 3);
548///
549/// // Rank-based access
550/// assert_eq!(point.rank(), 1 * (4 * 8) + 2 * 8 + 3);
551/// ```
552#[derive(Clone, Deserialize, Serialize, PartialEq, Eq, Hash, Debug)]
553pub struct Point {
554    rank: usize,
555    extent: Extent,
556}
557
558/// An iterator over the coordinates of a [`Point`] in row-major
559/// order.
560///
561/// Yields each coordinate component one at a time, without allocating
562/// a full coordinate vector.
563///
564/// The iteration is deterministic: the `i`-th call to `next()`
565/// returns the coordinate along axis `i`.
566///
567/// # Examples
568/// ```
569/// use ndslice::extent;
570///
571/// let ext = extent!(x = 2, y = 3);
572/// let point = ext.point(vec![1, 2]).unwrap();
573///
574/// let coords: Vec<_> = point.coords_iter().collect();
575/// assert_eq!(coords, vec![1, 2]);
576/// ```
577pub 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    /// Computes and returns the coordinate for the current axis, then
588    /// advances the iterator.
589    ///
590    /// Returns `None` once all dimensions have been exhausted.
591    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    /// Returns the exact number of coordinates remaining.
603    ///
604    /// Since the dimensionality of the [`Point`] is known up front,
605    /// this always returns `(n, Some(n))` where `n` is the number of
606    /// axes not yet yielded.
607    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    /// Iterate over the coordinate values of a [`Point`] (without
620    /// allocating).
621    ///
622    /// This allows using a `Point` directly in a `for` loop:
623    ///
624    /// ```
625    /// use ndslice::extent;
626    ///
627    /// let ext = extent!(x = 2, y = 3);
628    /// let point = ext.point(vec![1, 2]).unwrap();
629    ///
630    /// let coords: Vec<_> = (&point).into_iter().collect();
631    /// assert_eq!(coords, vec![1, 2]);
632    ///
633    /// for coord in &point {
634    ///     println!("{}", coord);
635    /// }
636    /// ```
637    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
648/// Extension trait for creating a `Point` from a coordinate vector
649/// and an `Extent`.
650///
651/// This trait provides the `.in_(&extent)` method, which constructs a
652/// `Point` using the caller as the coordinate vector and the given
653/// extent as the shape context.
654///
655/// # Example
656/// ```
657/// use ndslice::Extent;
658/// use ndslice::view::InExtent;
659/// let extent = Extent::new(vec!["x".into(), "y".into()], vec![3, 4]).unwrap();
660/// let point = vec![1, 2].in_(&extent).unwrap();
661/// assert_eq!(point.rank(), 1 * 4 + 2);
662/// ```
663pub trait InExtent {
664    fn in_(self, extent: &Extent) -> Result<Point, PointError>;
665}
666
667impl InExtent for Vec<usize> {
668    /// Creates a `Point` with the provided coordinates in the given
669    /// extent.
670    ///
671    /// Delegates to `Extent::point`.
672    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    /// Returns the coordinate of this [`Point`] along the given axis.
688    ///
689    /// The axis index `i` must be less than the number of dimensions
690    /// in the [`Extent`], otherwise this function will panic.
691    /// Computes only the `i`-th coordinate from the point's row-major
692    /// `rank`, avoiding materialization of the full coordinate
693    /// vector.
694    ///
695    /// # Examples
696    /// ```
697    /// use ndslice::extent;
698    ///
699    /// let ext = extent!(x = 2, y = 3);
700    /// let point = ext.point(vec![1, 2]).unwrap();
701    /// assert_eq!(point.coord(0), 1); // x
702    /// assert_eq!(point.coord(1), 2); // y
703    /// ```
704    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    /// Returns the full coordinate vector for this [`Point`]
711    /// (allocates).
712    ///
713    /// The vector contains one coordinate per dimension of the
714    /// [`Extent`], reconstructed from the point's row-major `rank`.
715    ///
716    /// # Examples
717    /// ```
718    /// use ndslice::extent;
719    ///
720    /// let ext = extent!(x = 2, y = 3);
721    /// let point = ext.point(vec![1, 2]).unwrap();
722    /// assert_eq!(point.coords(), vec![1, 2]);
723    /// ```
724    pub fn coords(&self) -> Vec<usize> {
725        self.coords_iter().collect()
726    }
727
728    /// Returns the linearized row-major rank of this [`Point`] within
729    /// its [`Extent`].
730    pub fn rank(&self) -> usize {
731        self.rank
732    }
733
734    /// Returns the [`Extent`] that defines the coordinate space of
735    /// this [`Point`].
736    pub fn extent(&self) -> &Extent {
737        &self.extent
738    }
739
740    /// Returns the number of dimensions in this [`Point`]'s
741    /// [`Extent`].
742    ///
743    /// This corresponds to the dimensionality of the coordinate
744    /// space, i.e. how many separate axes (labels) are present.
745    ///
746    /// # Examples
747    /// ```
748    /// use ndslice::extent;
749    ///
750    /// let ext = extent!(x = 2, y = 3, z = 4);
751    /// let point = ext.point(vec![1, 2, 3]).unwrap();
752    /// assert_eq!(point.len(), 3); // x, y, z
753    /// ```
754    pub fn len(&self) -> usize {
755        self.extent.len()
756    }
757
758    /// Returns `true` if this [`Point`] lies in a 0-dimensional
759    /// [`Extent`].
760    ///
761    /// A 0-D extent has no coordinate axes and exactly one valid
762    /// point (the empty tuple `[]`).
763    ///
764    /// # Examples
765    /// ```
766    /// use ndslice::extent;
767    ///
768    /// let ext = extent!();
769    /// let point = ext.point(vec![]).unwrap();
770    /// assert!(point.is_empty());
771    /// assert_eq!(point.len(), 0);
772    /// ```
773    pub fn is_empty(&self) -> bool {
774        self.extent.len() == 0
775    }
776
777    /// Formats the coordinates of this Point as a string suitable for display names.
778    /// Returns a string in the format: "{'label': coord/size, 'label': coord/size, ...}"
779    ///
780    /// # Examples
781    /// ```
782    /// use ndslice::extent;
783    ///
784    /// let ext = extent!(x = 2, y = 3);
785    /// let point = ext.point(vec![1, 2]).unwrap();
786    /// assert_eq!(point.format_as_dict(), "{'x': 1/2, 'y': 2/3}");
787    /// ```
788    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
803/// Formats a `Point` as a comma-separated list of per-axis
804/// coordinates against the point’s extent:
805/// ```text
806/// label=coord/size[,label=coord/size,...]
807/// ```
808///
809/// # Grammar
810/// ```text
811/// Point    := Pair ( "," Pair )*
812/// Pair     := Label "=" Coord "/" Size
813/// Label    := SafeIdent | Quoted
814/// SafeIdent:= [A-Za-z0-9_]+
815/// Quoted   := "\"" ( [^"\\] | "\\" . )* "\""   // Rust string-literal style
816/// Coord    := [0-9]+
817/// Size     := [0-9]+
818/// ```
819///
820/// # Quoting rules
821/// - Labels that are **not** `SafeIdent` are rendered using Rust
822///   string-literal syntax (via `labels::fmt_label`), e.g. `"dim/0"` or
823///   `"x y"`.
824/// - "Safe" means ASCII alphanumeric or underscore (`[A-Za-z0-9_]+`).
825///   Everything else is quoted.
826/// - Coordinates are shown in row-major order and each is paired with
827///   that axis’s size from the point’s extent.
828///
829/// # Examples
830/// ```text
831/// x=1/4,y=2/5,z=3/6
832/// "dim/0"=1/3,"dim,1"=2/5
833/// ```
834///
835/// Implementation note: label rendering is delegated to
836/// `labels::fmt_label` to keep quoting behavior consistent with
837/// `Extent` and `Region`.
838impl 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(); // quote
888                let mut label = String::new();
889                let mut escaped = false;
890
891                // Consume label until closing quote
892                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(); // consume comma
1014                while chars.peek() == Some(&' ') {
1015                    chars.next();
1016                }
1017                // After consuming a comma, there must be another dimension
1018                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/// Errors that occur while operating on views.
1043#[derive(Debug, Error)]
1044pub enum ViewError {
1045    /// The provided dimension does not exist in the relevant extent.
1046    #[error("no such dimension: {0}")]
1047    InvalidDim(String),
1048
1049    /// A view was attempted to be constructed from an empty (resolved) range.
1050    #[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/// Errors that occur while operating on Region.
1068#[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/// `Region` describes a region of a possibly-larger space of ranks, organized into
1078/// a hyperrect.
1079///
1080/// Internally, region consist of a set of labels and a [`Slice`], as it allows for
1081/// a compact but useful representation of the ranks. However, this representation
1082/// may change in the future.
1083#[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    /// Constructor to build arbitrary regions (incl. non-contiguous /
1115    /// offset).
1116    #[allow(dead_code)]
1117    pub fn new(labels: Vec<String>, slice: Slice) -> Self {
1118        Self { labels, slice }
1119    }
1120
1121    /// The labels of the dimensions of this region.
1122    pub fn labels(&self) -> &[String] {
1123        &self.labels
1124    }
1125
1126    /// The slice representing this region.
1127    /// Note: this representation may change.
1128    pub fn slice(&self) -> &Slice {
1129        &self.slice
1130    }
1131
1132    /// Convert this region into its constituent labels and slice.
1133    pub fn into_inner(self) -> (Vec<String>, Slice) {
1134        (self.labels, self.slice)
1135    }
1136
1137    /// Returns the extent of the region.
1138    pub fn extent(&self) -> Extent {
1139        Extent::new(self.labels.clone(), self.slice.sizes().to_vec()).unwrap()
1140    }
1141
1142    /// Returns `true` if this region is a subset of `other`, i.e., if `other`
1143    /// contains at least all of the ranks in this region.
1144    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                        // r < l
1158                        right.next();
1159                    }
1160                }
1161                (Some(_), None) => return false,
1162                (None, _) => return true,
1163            }
1164        }
1165    }
1166
1167    /// Remap the target to ranks in this region. The returned iterator iterates
1168    /// over each rank in `target`, providing the corresponding rank in `self`.
1169    /// This is useful when mapping between different subspaces.
1170    ///
1171    /// ```
1172    /// # use ndslice::Region;
1173    /// # use ndslice::ViewExt;
1174    /// # use ndslice::extent;
1175    /// let ext = extent!(replica = 8, gpu = 4);
1176    /// let replica1 = ext.range("replica", 1).unwrap();
1177    /// assert_eq!(replica1.extent(), extent!(replica = 1, gpu = 4));
1178    /// let replica1_gpu12 = replica1.range("gpu", 1..3).unwrap();
1179    /// assert_eq!(replica1_gpu12.extent(), extent!(replica = 1, gpu = 2));
1180    /// // The first rank in `replica1_gpu12` is the second rank in `replica1`.
1181    /// assert_eq!(
1182    ///     replica1.remap(&replica1_gpu12).unwrap().collect::<Vec<_>>(),
1183    ///     vec![1, 2],
1184    /// );
1185    /// ```
1186    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    /// Returns the total number of ranks in the region.
1206    pub fn num_ranks(&self) -> usize {
1207        self.slice.len()
1208    }
1209
1210    /// Convert a rank in this region's extent into its corresponding rank in
1211    /// the base space defined by the region's `Slice`.
1212    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    /// Convert a rank in the base space into the corresponding `Point` in this
1226    /// region's extent (if the base rank lies within the region's `Slice`).
1227    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
1239// We would make this impl<T: Viewable> From<T> for View,
1240// except this conflicts with the blanket impl for From<&T> for View.
1241impl 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
1265/// Formats a `Region` in a compact rectangular syntax:
1266/// ```text
1267/// [offset+]label=size/stride[,label=size/stride,...]
1268/// ```
1269/// # Grammar
1270///
1271/// ```text
1272/// Region   := [ Offset "+" ] Pair ( "," Pair )*
1273/// Offset   := [0-9]+
1274/// Pair     := Label "=" Size "/" Stride
1275/// Label    := SafeIdent | Quoted
1276/// SafeIdent:= [A-Za-z0-9_]+
1277/// Quoted   := "\"" ( [^"\\] | "\\" . )* "\""   // Rust string literal style
1278/// Size     := [0-9]+
1279/// Stride   := [0-9]+
1280/// ```
1281///
1282/// # Quoting rules
1283///
1284/// - Labels that are **not** `SafeIdent` are rendered using Rust
1285///   string-literal syntax (via `format!("{:?}", label)`).
1286/// - "Safe" means ASCII alphanumeric or underscore
1287///   (`[A-Za-z0-9_]+`). Everything else is quoted.
1288/// - The optional `Offset+` prefix appears only when the slice offset
1289///   is nonzero.
1290///
1291/// # Examples
1292///
1293/// ```text
1294/// x=2/1,y=3/2
1295/// 8+"dim/0"=4/1,"dim,1"=5/4
1296/// ```
1297///
1298/// # Notes
1299///
1300/// This format is both human-readable and machine-parsable. The
1301/// corresponding [`FromStr`] implementation accepts exactly this
1302/// grammar, including quoted labels. The quoting rule makes
1303/// round-trip unambiguous.
1304impl 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
1334/// Parses a `Region` from the textual form emitted by
1335/// [`Display`](Self::fmt).
1336///
1337/// The accepted syntax and quoting rules are exactly those documented
1338/// on `Display`: comma-separated `label=size/stride` pairs with an
1339/// optional `offset+` prefix, and labels that are either safe
1340/// identifiers or Rust string literals.
1341///
1342/// Returns a `RegionParseError` on malformed input.
1343///
1344/// # Examples
1345/// ```
1346/// use ndslice::view::Region;
1347///
1348/// let r: Region = "x=2/1,y=3/2".parse().unwrap();
1349/// assert_eq!(r.labels(), &["x", "y"]);
1350///
1351/// let q: Region = "8+\"dim/0\"=4/1".parse().unwrap();
1352/// assert_eq!(q.labels(), &["dim/0"]);
1353/// ```
1354impl 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            // Accept either a quoted label output or a bare token.
1377            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
1397/// Dense builder: constructs a mesh from a complete sequence of
1398/// values in the canonical order for `region`.
1399///
1400/// # Semantics
1401/// - Input must contain exactly `region.num_ranks()` items.
1402/// - Values must be in the same order as `region.slice().iter()`
1403///   (i.e., the order observed by `ViewExt::values()`).
1404///
1405/// # Errors
1406/// Returns [`Self::Error`] if `values.len() != region.num_ranks()`.
1407///
1408/// See also: [`BuildFromRegionIndexed`]
1409pub trait BuildFromRegion<T>: Sized {
1410    type Error;
1411
1412    /// Validates cardinality/shape and constructs a mesh.
1413    fn build_dense(region: Region, values: Vec<T>) -> Result<Self, Self::Error>;
1414
1415    /// Caller guarantees correct cardinality/order; no validation.
1416    fn build_dense_unchecked(region: Region, values: Vec<T>) -> Self;
1417}
1418
1419/// Indexed builder: constructs a mesh from sparse `(rank, value)`
1420/// pairs.
1421///
1422/// # Semantics
1423/// - Every rank in `0..region.num_ranks()` must be provided at least
1424///   once.
1425/// - Out-of-bounds ranks (`rank >= num_ranks()`) cause an error.
1426/// - Missing ranks cause an error.
1427/// - Duplicate ranks are allowed; the **last write wins**.
1428///
1429/// # Errors
1430/// Returns [`Self::Error`] if coverage is incomplete or a rank is out
1431/// of bounds.
1432pub trait BuildFromRegionIndexed<T>: Sized {
1433    type Error;
1434
1435    /// Validates coverage and bounds, and constructs a mesh from
1436    /// `(rank, value)` pairs.
1437    fn build_indexed(
1438        region: Region,
1439        pairs: impl IntoIterator<Item = (usize, T)>,
1440    ) -> Result<Self, Self::Error>;
1441}
1442
1443/// Mesh-aware collecting adapter.
1444///
1445/// Unlike `FromIterator`, this trait takes both an iterator *and* a
1446/// [`Region`] to build a mesh:
1447///
1448/// - Meshes always require a shape (`Region`) supplied externally.
1449/// - Cardinality must match: the iterator must yield exactly
1450///   `region.num_ranks()` items, or an error is returned.
1451/// - Construction goes through [`BuildFromRegion`], which validates
1452///   and builds the concrete mesh type.
1453///
1454/// In short: `collect_mesh` does for meshes what `collect` does for
1455/// ordinary collections, but with shape-awareness and validation.
1456pub 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
1462/// Blanket impl: enables `.collect_mesh(region)` on any iterator of
1463/// `T`.
1464impl<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/// A canonical cardinality mismatch descriptor.
1477#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1478pub struct InvalidCardinality {
1479    pub expected: usize,
1480    pub actual: usize,
1481}
1482
1483/// Exact-size, mesh-aware collecting adapter.
1484///
1485/// Like [`CollectMeshExt`], but for `ExactSizeIterator`. Performs a
1486/// `len()` pre-check to fail fast (no allocation) when `len() !=
1487/// region.num_ranks()`. On success, constructs `M` without
1488/// re-validating cardinality.
1489pub 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
1496/// Blanket impl: enables `.collect_exact_mesh(region)` on any
1497/// `ExactSizeIterator` of `T`.
1498impl<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
1516/// Indexed collecting adapter.
1517///
1518/// Consume `(rank, value)` pairs plus a [`Region`] and build a mesh
1519/// via [`BuildFromRegionIndexed`].
1520///
1521/// # Semantics (recommended contract)
1522/// Implementations of [`BuildFromRegionIndexed`] are expected to
1523/// enforce:
1524/// - **Coverage:** every rank in `0..region.num_ranks()` is provided
1525///   at least once.
1526/// - **Bounds:** any out-of-bounds rank (`rank >= num_ranks`) is an
1527///   error.
1528/// - **Duplicates:** allowed; **last write wins**.
1529/// - **Missing ranks:** an error (incomplete coverage).
1530///
1531/// # Errors
1532/// Propagates whatever [`BuildFromRegionIndexed::build_indexed`]
1533/// returns (e.g., a cardinality/bounds error) from the target mesh
1534/// type.
1535///
1536/// See also: [`BuildFromRegionIndexed`] for the authoritative policy.
1537pub 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
1543/// Blanket impl: enables `.collect_indexed(region)` on any iterator
1544/// of `(usize, T)` pairs.
1545impl<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
1558/// Map into any mesh `M`.
1559pub 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
1586/// Blanket impl: enables `.map_into(...)` and `.try_map_into`` on any
1587/// `Ranked`.
1588impl<T: Ranked> MapIntoExt for T {}
1589
1590/// A View is a collection of items in a space indexed by a [`Region`].
1591pub trait View: Sized {
1592    /// The type of item in this view.
1593    type Item;
1594
1595    /// The type of sub-view produced by manipulating (e.g., slicing) this view.
1596    type View: View;
1597
1598    /// The ranks contained in this view.
1599    fn region(&self) -> Region;
1600
1601    /// Retrieve the item corresponding to the given `rank` in the [`Region`]
1602    /// of this view. An implementation *MUST* return a value for all ranks
1603    /// defined in this view.
1604    fn get(&self, rank: usize) -> Option<Self::Item>;
1605
1606    /// Subsets this view with the provided ranks. This is mainly used
1607    /// by combinators on Views themselves. The set of ranks passed in
1608    /// must be a subset of the ranks of the base view.
1609    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `ViewError`.
1610    fn subset(&self, region: Region) -> Result<Self::View, ViewError>;
1611}
1612
1613/// A [`Region`] is also a View.
1614impl View for Region {
1615    /// The type of item is the rank in the underlying space.
1616    type Item = usize;
1617
1618    /// The type of sub-view is also a [`Region`].
1619    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
1641/// An [`Extent`] is also a View.
1642impl View for Extent {
1643    /// The type of item is the rank itself.
1644    type Item = usize;
1645
1646    /// The type of sub-view can be a [`Region`], since
1647    /// [`Extent`] can only describe a complete space.
1648    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
1670/// Ranked is a helper trait to implement `View` on a ranked collection
1671/// of items.
1672pub trait Ranked: Sized {
1673    /// The type of item in this view.
1674    type Item: 'static;
1675
1676    /// The ranks contained in this view.
1677    fn region(&self) -> &Region;
1678
1679    /// Return the item at `rank`
1680    fn get(&self, rank: usize) -> Option<&Self::Item>;
1681}
1682
1683/// Extension of [`Ranked`] for types that can materialize a new owned
1684/// view.
1685///
1686/// `sliced` constructs a new instance containing only the items in
1687/// the given `region`. Unlike [`Ranked::get`], which only provides
1688/// borrowed access, `sliced` requires that the implementor can
1689/// produce owned values for the specified sub-region.
1690pub trait RankedSliceable: Ranked {
1691    /// Construct a new Ranked containing the ranks in this view that
1692    /// are part of region. The caller guarantees that
1693    /// `region.is_subset(self.region())`.
1694    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
1724/// An iterator over views.
1725pub struct ViewIterator {
1726    extent: Extent,     // Note that `extent` and...
1727    pos: SliceIterator, // ... `pos` share the same `Slice`.
1728}
1729
1730impl Iterator for ViewIterator {
1731    type Item = (Point, usize);
1732    fn next(&mut self) -> Option<Self::Item> {
1733        // This is a rank in the base space.
1734        let rank = self.pos.next()?;
1735        // Here, we convert to view space.
1736        let coords = self.pos.slice.coordinates(rank).unwrap();
1737        let point = coords.in_(&self.extent).unwrap();
1738        Some((point, rank))
1739    }
1740}
1741
1742/// Extension methods for view construction.
1743pub trait ViewExt: View {
1744    /// Construct a view comprising the range of points along the provided dimension.
1745    ///
1746    /// ## Examples
1747    ///
1748    /// ```
1749    /// use ndslice::Range;
1750    /// use ndslice::ViewExt;
1751    /// use ndslice::extent;
1752    ///
1753    /// let ext = extent!(zone = 4, host = 2, gpu = 8);
1754    ///
1755    /// // Subselect zone index 0.
1756    /// assert_eq!(ext.range("zone", 0).unwrap().iter().count(), 16);
1757    ///
1758    /// // Even GPUs within zone 0
1759    /// assert_eq!(
1760    ///     ext.range("zone", 0)
1761    ///         .unwrap()
1762    ///         .range("gpu", Range(0, None, 2))
1763    ///         .unwrap()
1764    ///         .iter()
1765    ///         .count(),
1766    ///     8
1767    /// );
1768    /// ```
1769    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `ViewError`.
1770    fn range<R: Into<Range>>(&self, dim: &str, range: R) -> Result<Self::View, ViewError>;
1771
1772    /// Group by view on `dim`. The returned iterator enumerates all groups
1773    /// as views in the extent of `dim` to the last dimension of the view.
1774    ///
1775    /// ## Examples
1776    ///
1777    /// ```
1778    /// use ndslice::ViewExt;
1779    /// use ndslice::extent;
1780    ///
1781    /// let ext = extent!(zone = 4, host = 2, gpu = 8);
1782    ///
1783    /// // We generate one view for each zone.
1784    /// assert_eq!(ext.group_by("host").unwrap().count(), 4);
1785    ///
1786    /// let mut parts = ext.group_by("host").unwrap();
1787    ///
1788    /// let zone0 = parts.next().unwrap();
1789    /// let mut zone0_points = zone0.iter();
1790    /// assert_eq!(zone0.extent(), extent!(host = 2, gpu = 8));
1791    /// assert_eq!(
1792    ///     zone0_points.next().unwrap(),
1793    ///     (extent!(host = 2, gpu = 8).point(vec![0, 0]).unwrap(), 0)
1794    /// );
1795    /// assert_eq!(
1796    ///     zone0_points.next().unwrap(),
1797    ///     (extent!(host = 2, gpu = 8).point(vec![0, 1]).unwrap(), 1)
1798    /// );
1799    ///
1800    /// let zone1 = parts.next().unwrap();
1801    /// assert_eq!(zone1.extent(), extent!(host = 2, gpu = 8));
1802    /// assert_eq!(
1803    ///     zone1.iter().next().unwrap(),
1804    ///     (extent!(host = 2, gpu = 8).point(vec![0, 0]).unwrap(), 16)
1805    /// );
1806    /// ```
1807    #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `ViewError`.
1808    fn group_by(&self, dim: &str) -> Result<impl Iterator<Item = Self::View>, ViewError>;
1809
1810    /// The extent of this view. Every point in this space is defined.
1811    fn extent(&self) -> Extent;
1812
1813    /// Iterate over all points in this region.
1814    fn iter(&self) -> impl Iterator<Item = (Point, Self::Item)> + '_;
1815
1816    /// Iterate over the values in the region.
1817    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            // These are always valid sub-views.
1867            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/// Construct a new extent with the given set of dimension-size pairs.
1897///
1898/// ```
1899/// let s = ndslice::extent!(host = 2, gpu = 8);
1900/// assert_eq!(s.labels(), &["host".to_string(), "gpu".to_string()]);
1901/// assert_eq!(s.sizes(), &[2, 8]);
1902/// ```
1903#[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); // 24 ranks
2013        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        // 0-D extent has exactly one rank (the single point).
2031        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); // Should panic
2161    }
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        // Iterator invariants for 0-D point.
2225        let mut it = (&points[0]).into_iter();
2226        assert_eq!(it.len(), 0);
2227        assert!(it.next().is_none()); // no items
2228        assert!(it.next().is_none()); // fused
2229    }
2230
2231    #[test]
2232    fn test_extent_concat() {
2233        // Test basic concatenation of two extents with preserved order of labels
2234        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        // Test concatenating with empty extent
2243        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        // Test concatenating two empty extents
2253        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); // 0-dimensional extent has 1 rank
2257
2258        // Test self-concatenation (overlapping labels should cause error)
2259        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"); // Overlapping label should be "x"
2267            }
2268            other => panic!("Expected OverlappingLabel error, got {:?}", other),
2269        }
2270
2271        // Test concatenation creates valid points
2272        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        // Test error case: overlapping labels with same size (should error)
2278        let extent_a = extent!(x = 2, y = 3);
2279        let extent_b = extent!(y = 3, z = 4); // y overlaps with same size
2280        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"); // the overlapping label
2288            }
2289            other => panic!("Expected OverlappingLabel error, got {:?}", other),
2290        }
2291
2292        // Test that Extent::concat preserves order and is not commutative
2293        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        // Test associativity: (a ⊕ b) ⊕ c == a ⊕ (b ⊕ c) for disjoint labels
2305        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        // Labels include characters ("/", ",") that force quoting.
2351        let ext = Extent::new(vec!["dim/0".into(), "dim,1".into()], vec![3, 5]).unwrap();
2352
2353        // Extent::Display should quote both labels.
2354        assert_eq!(format!("{}", ext), "{\"dim/0\": 3, \"dim,1\": 5}");
2355
2356        // Point::Display should also quote labels consistently.
2357        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        // Given a rank in the root shape, return the corresponding point in the
2364        // provided shape, which is a view of the root shape.
2365        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        // rows are `hosts`, cols are gpus
2373        // replicas = 0
2374        //     0,    1,  2,    3,
2375        //     (4),  5,  (6),  7,
2376        //     8,    9,  10,   11,
2377        //     (12), 13, (14), 15,
2378        // replicas = 3, which is [replicas=0] + 48
2379        //     48,   49, 50,   51,
2380        //     (52), 53, (54), 55,
2381        //     56,   57, 58,   59,
2382        //     (60), 61, (62), 63,
2383        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(); // 3×4
2454        let b = a.range("y", 1..3).unwrap(); // 3×2 (subset of a)
2455        let c = e.range("x", 0..2).unwrap(); // 2×4 (overlaps, not subset of a)
2456
2457        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        // Self-remap
2469        assert_eq!(
2470            region.remap(&region).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        // Test double subsetting:
2489
2490        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        // The first rank in `replica1_gpu12` is the second rank in `replica1`.
2496        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        // region is a 2x2 region of extent, with the ranks in the region and
2511        // and ranks in the extent are mapped as follows:
2512        //  0,        1
2513        // [2] -> 0, [3] -> 1
2514        // [4] -> 2, [5] -> 3
2515        //  6,        7
2516        // Use a point from an extent different from this region should fail:
2517        assert!(
2518            region
2519                .base_rank_of_point(extent.point_of_rank(0).unwrap())
2520                .is_err()
2521        );
2522        // Convert ranks in the extent to ranks in the region:
2523        assert_eq!(region.base_rank_of_point(point(0, &region)).unwrap(), 2);
2524        assert_eq!(region.base_rank_of_point(point(1, &region)).unwrap(), 3);
2525        assert_eq!(region.base_rank_of_point(point(2, &region)).unwrap(), 4);
2526        assert_eq!(region.base_rank_of_point(point(3, &region)).unwrap(), 5);
2527        // Convert ranks in the region to ranks in the extent:
2528        assert_eq!(region.point_of_base_rank(2).unwrap(), point(0, &region));
2529        assert_eq!(region.point_of_base_rank(3).unwrap(), point(1, &region));
2530        assert_eq!(region.point_of_base_rank(4).unwrap(), point(2, &region));
2531        assert_eq!(region.point_of_base_rank(5).unwrap(), point(3, &region));
2532        // Coverting ranks outside the region should fail:
2533        assert!(region.point_of_base_rank(1).is_err());
2534        assert!(region.point_of_base_rank(6).is_err());
2535
2536        // Slice region to give a subregion:
2537        let subset = region
2538            .range("replicas", 1..2)
2539            .unwrap()
2540            .range("gpus", 1..2)
2541            .unwrap();
2542        // subset is a 1x1 region of extent, with ranksmapped as follows:
2543        // 0,  1
2544        // 2,  3
2545        // 4, [5] -> 0
2546        // 6,  7
2547        // Convert ranks in the extent to ranks in the subset:
2548        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        // or fail if the rank is not in the subset:
2551        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            // Roundtrip display->parse correctly and preserves equality.
2565            assert_eq!(
2566                region,
2567                region.to_string().parse::<Region>().unwrap(),
2568                "failed to roundtrip region {}", region
2569            );
2570        }
2571    }
2572
2573    // Property: `Region::Display` and `FromStr` remain a lossless
2574    // round-trip even when the slice has a nonzero offset.
2575    //
2576    // - Construct a region, then force its slice to have `offset =
2577    //   8`.
2578    // - Convert that region to a string via `Display`.
2579    // - Parse it back via `FromStr`.
2580    //
2581    // The parsed region must equal the original, showing that
2582    // offsets are encoded and decoded consistently in the textual
2583    // format.
2584    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    // Property: For any randomly generated strided `Region`,
2599    // converting it to a string with `Display` and parsing it back
2600    // with `FromStr` yields the same region.
2601    //
2602    // This ensures that:
2603    // - Strided layouts (with arbitrary steps and begins) are
2604    //   faithfully represented by the textual format.
2605    // - Offsets, sizes, and strides survive the round-trip without
2606    //   loss.
2607    // - Quoting rules for labels remain consistent.
2608    proptest! {
2609        #[test]
2610        fn region_strided_display_parse_roundtrips(
2611            region in gen_region_strided(1..=4, 6, 3, 16)
2612        ) {
2613            // Example: 122+"d/0"=1/30,"d/1"=5/6,"d/2"=4/1
2614            //
2615            // Decoding:
2616            // - `122+` is the slice offset.
2617            // - Each "label=size/stride" shows the post-slice size
2618            //   and stride.
2619            //   Sizes are reduced by `ceil((base_size - begin) /
2620            //   step)`.
2621            //   Strides are the base row-major strides, each
2622            //   multiplied by its step.
2623            let s = region.to_string();
2624            let parsed: Region = s.parse().unwrap();
2625            prop_assert_eq!(parsed, region);
2626        }
2627    }
2628
2629    // Property: `Region::Display` faithfully reflects its underlying
2630    // `Slice`.
2631    //
2632    // - The offset printed as `offset+` must equal
2633    //   `region.slice().offset()`.
2634    // - Each `label=size/stride` entry must show the size and stride
2635    //   from the underlying slice.
2636    //
2637    // This ensures that the textual representation is consistent
2638    // with the region’s internal geometry.
2639    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            // Check offset if present
2648            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            // Collect all size/stride pairs from the string
2659            let body = s.split('+').next_back().unwrap(); // after offset if any
2660            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                // part looks like label=size/stride
2665                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![]), // empty point
2697            // Test with spaces
2698            (" 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        // Test parsing points with quoted labels
2711        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        // Test various error cases
2727        let error_cases = vec![
2728            "x=1,y=2/5",     // missing size for x
2729            "x=1/4,y=2",     // missing size for y
2730            "x=1/4,y=/5",    // missing coord for y
2731            "x=/4,y=2/5",    // missing coord for x
2732            "x=1/4,y=2/",    // missing size after /
2733            "x=1/,y=2/5",    // missing size after /
2734            "x=1/4=5,y=2/5", // extra equals
2735            "x=1/4/6,y=2/5", // extra slash
2736            "x=abc/4,y=2/5", // invalid coord
2737            "x=1/abc,y=2/5", // invalid size
2738            "=1/4,y=2/5",    // missing label
2739            "x=1/4,",        // trailing comma with empty part
2740            "x=1/4,=2/5",    // missing label after comma
2741            "x=1/4,y",       // incomplete dimension
2742            "x",             // just a label
2743            "x=",            // label with equals but no coord/size
2744            "x=1/4,y=10/5",  // coord out of bounds (y has size 5, max coord is 4)
2745        ];
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        // Test that coordinates are validated against sizes
2756        let input = "x=5/4,y=2/5"; // x coord is 5 but size is 4 (max valid coord is 3)
2757        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        // Test that all dimension must be consistent (same labels and sizes for given extent)
2775        // This is implicitly tested by the round-trip tests, but let's be explicit
2776
2777        // Valid consistent point
2778        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        /// `Point.coord(i)` and `(&Point).into_iter()` must agree with
2788        /// `coords()`.
2789        #[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        // `points().count()` must equal `num_ranks()`.
2803        #[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        // `CoordIter` must report an exact size, decreasing by one on
2810        // each iteration, ending at zero, and yield the same sequence
2811        // as `coords()`.
2812        #[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                // Initial length matches dimensionality; size_hint
2818                // agrees.
2819                let mut remaining = p.len();
2820                prop_assert_eq!(it.len(), remaining);
2821                prop_assert_eq!(it.size_hint(), (remaining, Some(remaining)));
2822
2823                // Track yielded coords to compare with p.coords()
2824                let mut yielded = Vec::with_capacity(remaining);
2825
2826                // len() decreases by 1 per step; size_hint stays
2827                // consistent.
2828                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                // Exhausted: zero remaining, fused behavior (keeps
2836                // returning None).
2837                prop_assert_eq!(remaining, 0);
2838                prop_assert!(it.next().is_none());
2839                prop_assert!(it.next().is_none());
2840
2841                // Sequence equals full coords() reconstruction.
2842                prop_assert_eq!(yielded, p.coords());
2843            }
2844        }
2845
2846        // `rank_of_coords` must reject coordinate vectors of the
2847        // wrong length with a `PointError::DimMismatch` that reports
2848        // both the expected and actual dimensionality.
2849        #[test]
2850        fn rank_of_coords_dim_mismatch(extent in gen_extent(0..=4, 8)) {
2851            let want = extent.len();
2852            // Pick a wrong coords length for the extent.
2853            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        // `rank_of_coords` must reject coordinates with an index ≥
2866        // the dimension's size, producing
2867        // `PointError::OutOfRangeIndex` that reports both the
2868        // dimension size and the offending index.
2869        #[test]
2870        fn rank_of_coords_out_of_range_index(extent in gen_extent(1..=4, 8)) {
2871            // `extent` has at least 1 dim here.
2872            let sizes = extent.sizes().to_vec();
2873            // Start with a valid zero vector.
2874            let mut coords = vec![0usize; sizes.len()];
2875            // Bump one axis out of range.
2876            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        /// `point_of_rank` must reject `rank == num_ranks()` (first OOB),
2889        /// returning `OutOfRangeRank` with the correct fields.
2890        #[test]
2891        fn point_of_rank_out_of_range(extent in gen_extent(0..=4, 8)) {
2892            let total = extent.num_ranks(); // first invalid rank
2893            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        // Property test: Point display/parse round-trip should always work
2903        #[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        /// New rank-walk iterator must match legacy CartesianIterator
2915        /// + rank_of_coords.
2916        #[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            // Compare lengths without cloning the iterator.
2921            let legacy_len = CartesianIterator::new(sizes.clone()).count();
2922            prop_assert_eq!(legacy_len, extent.num_ranks());
2923
2924            // Compare element-by-element: (Point from new iter) vs
2925            // (coords from legacy)
2926            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}