tenferro_ops/
shape_extent.rs

1//! Shape extent metadata with exactness guarantees.
2
3/// A dimension expression plus the guarantee it provides.
4///
5/// `Exact` means the expression is the runtime size. `UpperBound` means the
6/// runtime size is no larger than the expression. `Unknown` preserves rank
7/// when no useful bound is available.
8///
9/// # Examples
10///
11/// ```
12/// use tenferro_ops::shape_extent::ShapeExtent;
13///
14/// let extent = ShapeExtent::exact(4usize);
15/// assert_eq!(extent.as_exact(), Some(&4));
16/// ```
17#[derive(Clone, Debug, PartialEq, Eq, Hash)]
18pub enum ShapeExtent<D> {
19    /// The dimension expression is exact.
20    Exact(D),
21    /// The dimension expression is an upper bound.
22    UpperBound(D),
23    /// The dimension is rank-known but otherwise unknown.
24    Unknown,
25}
26
27impl<D> ShapeExtent<D> {
28    /// Construct an exact extent.
29    ///
30    /// # Examples
31    ///
32    /// ```
33    /// use tenferro_ops::shape_extent::ShapeExtent;
34    ///
35    /// let extent = ShapeExtent::exact(3usize);
36    /// assert!(extent.is_exact());
37    /// ```
38    pub fn exact(dim: D) -> Self {
39        Self::Exact(dim)
40    }
41
42    /// Construct an upper-bound extent.
43    ///
44    /// # Examples
45    ///
46    /// ```
47    /// use tenferro_ops::shape_extent::ShapeExtent;
48    ///
49    /// let extent = ShapeExtent::upper_bound(3usize);
50    /// assert!(!extent.is_exact());
51    /// ```
52    pub fn upper_bound(dim: D) -> Self {
53        Self::UpperBound(dim)
54    }
55
56    /// Construct an unknown extent.
57    ///
58    /// # Examples
59    ///
60    /// ```
61    /// use tenferro_ops::shape_extent::ShapeExtent;
62    ///
63    /// let extent: ShapeExtent<usize> = ShapeExtent::unknown();
64    /// assert_eq!(extent.bound_expr(), None);
65    /// ```
66    pub fn unknown() -> Self {
67        Self::Unknown
68    }
69
70    /// Return true when this extent is exact.
71    ///
72    /// # Examples
73    ///
74    /// ```
75    /// use tenferro_ops::shape_extent::ShapeExtent;
76    ///
77    /// assert!(ShapeExtent::exact(2usize).is_exact());
78    /// assert!(!ShapeExtent::upper_bound(2usize).is_exact());
79    /// ```
80    pub fn is_exact(&self) -> bool {
81        matches!(self, Self::Exact(_))
82    }
83
84    /// Return the exact dimension expression, if this extent is exact.
85    ///
86    /// # Examples
87    ///
88    /// ```
89    /// use tenferro_ops::shape_extent::ShapeExtent;
90    ///
91    /// assert_eq!(ShapeExtent::exact(5usize).as_exact(), Some(&5));
92    /// assert_eq!(ShapeExtent::upper_bound(5usize).as_exact(), None);
93    /// ```
94    pub fn as_exact(&self) -> Option<&D> {
95        match self {
96            Self::Exact(dim) => Some(dim),
97            Self::UpperBound(_) | Self::Unknown => None,
98        }
99    }
100
101    /// Return the known bound expression, if any.
102    ///
103    /// # Examples
104    ///
105    /// ```
106    /// use tenferro_ops::shape_extent::ShapeExtent;
107    ///
108    /// assert_eq!(ShapeExtent::upper_bound(5usize).bound_expr(), Some(&5));
109    /// ```
110    pub fn bound_expr(&self) -> Option<&D> {
111        match self {
112            Self::Exact(dim) | Self::UpperBound(dim) => Some(dim),
113            Self::Unknown => None,
114        }
115    }
116
117    /// Map the contained dimension expression while preserving exactness.
118    ///
119    /// # Examples
120    ///
121    /// ```
122    /// use tenferro_ops::shape_extent::ShapeExtent;
123    ///
124    /// let extent = ShapeExtent::upper_bound(5usize).map(|dim| dim + 1);
125    /// assert_eq!(extent.bound_expr(), Some(&6));
126    /// ```
127    pub fn map<E>(self, f: impl FnOnce(D) -> E) -> ShapeExtent<E> {
128        match self {
129            Self::Exact(dim) => ShapeExtent::Exact(f(dim)),
130            Self::UpperBound(dim) => ShapeExtent::UpperBound(f(dim)),
131            Self::Unknown => ShapeExtent::Unknown,
132        }
133    }
134}
tenferro_ops/shape_extent.rs

tenferro_ops/
shape_extent.rs
