1#![allow(missing_docs)]
2
3use std::any::{Any, TypeId};
6use std::cmp;
7use std::fmt::Debug;
8use std::ops::{Add, Div, Mul, Sub};
9use typenum::{self, Diff, Max, Maximum, Min, Minimum, Prod, Quot, Sum, Unsigned};
10
11#[cfg(feature = "serde-serialize-no-std")]
12use serde::{Deserialize, Deserializer, Serialize, Serializer};
13
14#[derive(Clone, Copy, Eq, PartialEq, Debug)]
16#[cfg_attr(
17 all(not(target_os = "cuda"), feature = "cuda"),
18 derive(cust::DeviceCopy)
19)]
20pub struct Dynamic {
21 value: usize,
22}
23
24impl Dynamic {
25 #[inline]
27 pub const fn new(value: usize) -> Self {
28 Self { value }
29 }
30}
31
32#[cfg(feature = "serde-serialize-no-std")]
33impl Serialize for Dynamic {
34 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
35 where
36 S: Serializer,
37 {
38 self.value.serialize(serializer)
39 }
40}
41
42#[cfg(feature = "serde-serialize-no-std")]
43impl<'de> Deserialize<'de> for Dynamic {
44 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
45 where
46 D: Deserializer<'de>,
47 {
48 usize::deserialize(deserializer).map(|x| Dynamic { value: x })
49 }
50}
51
52pub trait IsDynamic {}
54pub trait IsNotStaticOne {}
56
57impl IsDynamic for Dynamic {}
58impl IsNotStaticOne for Dynamic {}
59
60pub unsafe trait Dim: Any + Debug + Copy + PartialEq + Send + Sync {
63 #[inline(always)]
64 fn is<D: Dim>() -> bool {
65 TypeId::of::<Self>() == TypeId::of::<D>()
66 }
67
68 fn try_to_usize() -> Option<usize>;
71
72 fn value(&self) -> usize;
75
76 fn from_usize(dim: usize) -> Self;
79}
80
81unsafe impl Dim for Dynamic {
82 #[inline]
83 fn try_to_usize() -> Option<usize> {
84 None
85 }
86
87 #[inline]
88 fn from_usize(dim: usize) -> Self {
89 Self::new(dim)
90 }
91
92 #[inline]
93 fn value(&self) -> usize {
94 self.value
95 }
96}
97
98impl Add<usize> for Dynamic {
99 type Output = Dynamic;
100
101 #[inline]
102 fn add(self, rhs: usize) -> Self {
103 Self::new(self.value + rhs)
104 }
105}
106
107impl Sub<usize> for Dynamic {
108 type Output = Dynamic;
109
110 #[inline]
111 fn sub(self, rhs: usize) -> Self {
112 Self::new(self.value - rhs)
113 }
114}
115
116macro_rules! dim_ops(
123 ($($DimOp: ident, $DimNameOp: ident,
124 $Op: ident, $op: ident, $op_path: path,
125 $DimResOp: ident, $DimNameResOp: ident,
126 $ResOp: ident);* $(;)*) => {$(
127 pub type $DimResOp<D1, D2> = <D1 as $DimOp<D2>>::Output;
128
129 pub trait $DimOp<D: Dim>: Dim {
130 type Output: Dim;
131
132 fn $op(self, other: D) -> Self::Output;
133 }
134
135 impl<const A: usize, const B: usize> $DimOp<Const<B>> for Const<A>
136 where
137 Const<A>: ToTypenum,
138 Const<B>: ToTypenum,
139 <Const<A> as ToTypenum>::Typenum: $Op<<Const<B> as ToTypenum>::Typenum>,
140 $ResOp<<Const<A> as ToTypenum>::Typenum, <Const<B> as ToTypenum>::Typenum>: ToConst,
141 {
142 type Output =
143 <$ResOp<<Const<A> as ToTypenum>::Typenum, <Const<B> as ToTypenum>::Typenum> as ToConst>::Const;
144
145 fn $op(self, _: Const<B>) -> Self::Output {
146 Self::Output::name()
147 }
148 }
149
150 impl<D: Dim> $DimOp<D> for Dynamic {
151 type Output = Dynamic;
152
153 #[inline]
154 fn $op(self, other: D) -> Dynamic {
155 Dynamic::new($op_path(self.value, other.value()))
156 }
157 }
158
159 impl<D: DimName> $DimOp<Dynamic> for D {
161 type Output = Dynamic;
162
163 #[inline]
164 fn $op(self, other: Dynamic) -> Dynamic {
165 Dynamic::new($op_path(self.value(), other.value))
166 }
167 }
168
169 pub type $DimNameResOp<D1, D2> = <D1 as $DimNameOp<D2>>::Output;
170
171 pub trait $DimNameOp<D: DimName>: DimName {
172 type Output: DimName;
173
174 fn $op(self, other: D) -> Self::Output;
175 }
176
177 impl<const A: usize, const B: usize> $DimNameOp<Const<B>> for Const<A>
178 where
179 Const<A>: ToTypenum,
180 Const<B>: ToTypenum,
181 <Const<A> as ToTypenum>::Typenum: $Op<<Const<B> as ToTypenum>::Typenum>,
182 $ResOp<<Const<A> as ToTypenum>::Typenum, <Const<B> as ToTypenum>::Typenum>: ToConst,
183 {
184 type Output =
185 <$ResOp<<Const<A> as ToTypenum>::Typenum, <Const<B> as ToTypenum>::Typenum> as ToConst>::Const;
186
187 fn $op(self, _: Const<B>) -> Self::Output {
188 Self::Output::name()
189 }
190 }
191 )*}
192);
193
194dim_ops!(
195 DimAdd, DimNameAdd, Add, add, Add::add, DimSum, DimNameSum, Sum;
196 DimMul, DimNameMul, Mul, mul, Mul::mul, DimProd, DimNameProd, Prod;
197 DimSub, DimNameSub, Sub, sub, Sub::sub, DimDiff, DimNameDiff, Diff;
198 DimDiv, DimNameDiv, Div, div, Div::div, DimQuot, DimNameQuot, Quot;
199 DimMin, DimNameMin, Min, min, cmp::min, DimMinimum, DimNameMinimum, Minimum;
200 DimMax, DimNameMax, Max, max, cmp::max, DimMaximum, DimNameMaximum, Maximum;
201);
202
203#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
204#[cfg_attr(
205 all(not(target_os = "cuda"), feature = "cuda"),
206 derive(cust::DeviceCopy)
207)]
208pub struct Const<const R: usize>;
209
210pub trait DimName: Dim {
212 const USIZE: usize;
213
214 fn name() -> Self;
216
217 fn dim() -> usize;
220}
221
222#[cfg(feature = "serde-serialize-no-std")]
223impl<const D: usize> Serialize for Const<D> {
224 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
225 where
226 S: Serializer,
227 {
228 ().serialize(serializer)
229 }
230}
231
232#[cfg(feature = "serde-serialize-no-std")]
233impl<'de, const D: usize> Deserialize<'de> for Const<D> {
234 fn deserialize<Des>(deserializer: Des) -> Result<Self, Des::Error>
235 where
236 Des: Deserializer<'de>,
237 {
238 <()>::deserialize(deserializer).map(|_| Const::<D>)
239 }
240}
241
242#[cfg(feature = "rkyv-serialize-no-std")]
243mod rkyv_impl {
244 use super::Const;
245 use rkyv::{Archive, Deserialize, Fallible, Serialize};
246
247 impl<const R: usize> Archive for Const<R> {
248 type Archived = Self;
249 type Resolver = ();
250
251 fn resolve(
252 &self,
253 _: usize,
254 _: Self::Resolver,
255 _: &mut core::mem::MaybeUninit<Self::Archived>,
256 ) {
257 }
258 }
259
260 impl<S: Fallible + ?Sized, const R: usize> Serialize<S> for Const<R> {
261 fn serialize(&self, _: &mut S) -> Result<Self::Resolver, S::Error> {
262 Ok(())
263 }
264 }
265
266 impl<D: Fallible + ?Sized, const R: usize> Deserialize<Self, D> for Const<R> {
267 fn deserialize(&self, _: &mut D) -> Result<Self, D::Error> {
268 Ok(Const)
269 }
270 }
271}
272
273pub trait ToConst {
274 type Const: DimName;
275}
276
277pub trait ToTypenum {
278 type Typenum: Unsigned;
279}
280
281unsafe impl<const T: usize> Dim for Const<T> {
282 fn try_to_usize() -> Option<usize> {
283 Some(T)
284 }
285
286 fn value(&self) -> usize {
287 T
288 }
289
290 fn from_usize(dim: usize) -> Self {
291 assert_eq!(dim, T);
292 Self
293 }
294}
295
296impl<const T: usize> DimName for Const<T> {
297 const USIZE: usize = T;
298
299 #[inline]
300 fn name() -> Self {
301 Self
302 }
303
304 #[inline]
305 fn dim() -> usize {
306 T
307 }
308}
309
310pub type U1 = Const<1>;
311
312impl ToTypenum for Const<{ typenum::U1::USIZE }> {
313 type Typenum = typenum::U1;
314}
315
316impl ToConst for typenum::U1 {
317 type Const = Const<{ typenum::U1::USIZE }>;
318}
319
320macro_rules! from_to_typenum (
321 ($($D: ident),* $(,)*) => {$(
322 pub type $D = Const<{ typenum::$D::USIZE }>;
323
324 impl ToTypenum for Const<{ typenum::$D::USIZE }> {
325 type Typenum = typenum::$D;
326 }
327
328 impl ToConst for typenum::$D {
329 type Const = Const<{ typenum::$D::USIZE }>;
330 }
331
332 impl IsNotStaticOne for $D { }
333 )*}
334);
335
336from_to_typenum!(
337 U0, U2, U3, U4, U5, U6, U7, U8, U9, U10, U11, U12, U13, U14, U15, U16, U17, U18,
338 U19, U20, U21, U22, U23, U24, U25, U26, U27, U28, U29, U30, U31, U32, U33, U34, U35, U36, U37,
339 U38, U39, U40, U41, U42, U43, U44, U45, U46, U47, U48, U49, U50, U51, U52, U53, U54, U55, U56,
340 U57, U58, U59, U60, U61, U62, U63, U64, U65, U66, U67, U68, U69, U70, U71, U72, U73, U74, U75,
341 U76, U77, U78, U79, U80, U81, U82, U83, U84, U85, U86, U87, U88, U89, U90, U91, U92, U93, U94,
342 U95, U96, U97, U98, U99, U100, U101, U102, U103, U104, U105, U106, U107, U108, U109, U110,
343 U111, U112, U113, U114, U115, U116, U117, U118, U119, U120, U121, U122, U123, U124, U125, U126,
344 U127
345);