cubecl_core/frontend/comptime_option.rs
1use crate::{self as cubecl};
2use cubecl::prelude::*;
3
4#[derive(Default, Clone, Copy, CubeType)]
5pub enum ComptimeOption<T: CubeType> {
6 #[default]
7 None,
8 Some(T),
9}
10
11impl<T: CubeType<ExpandType: Clone>> Clone for ComptimeOptionExpand<T> {
12 fn clone(&self) -> Self {
13 self.clone_unchecked()
14 }
15}
16
17impl<T: CubeType<ExpandType: Copy>> Copy for ComptimeOptionExpand<T> {}
18
19#[allow(clippy::derivable_impls)]
20impl<T: CubeType> Default for ComptimeOptionExpand<T> {
21 fn default() -> Self {
22 Self::None
23 }
24}
25
26#[allow(non_snake_case)]
27impl<T: CubeType> ComptimeOption<T> {
28 pub fn __expand_Some(scope: &Scope, value: T::ExpandType) -> ComptimeOptionExpand<T> {
29 Self::__expand_new_Some(scope, value)
30 }
31}
32
33impl<T: CubeType> ComptimeOptionExpand<T> {
34 pub fn is_some(&self) -> bool {
35 match self {
36 ComptimeOptionExpand::Some(_) => true,
37 ComptimeOptionExpand::None => false,
38 }
39 }
40
41 pub fn unwrap(self) -> T::ExpandType {
42 match self {
43 Self::Some(val) => val,
44 Self::None => panic!("Unwrap on a None CubeOption"),
45 }
46 }
47
48 pub fn is_none(&self) -> bool {
49 !self.is_some()
50 }
51
52 pub fn unwrap_or(self, fallback: T::ExpandType) -> T::ExpandType {
53 match self {
54 ComptimeOptionExpand::Some(val) => val,
55 ComptimeOptionExpand::None => fallback,
56 }
57 }
58}
59
60pub enum ComptimeOptionArgs<T: LaunchArg, R: Runtime> {
61 Some(<T as LaunchArg>::RuntimeArg<R>),
62 None,
63}
64
65impl<T: LaunchArg, R: Runtime> From<Option<<T as LaunchArg>::RuntimeArg<R>>>
66 for ComptimeOptionArgs<T, R>
67{
68 fn from(value: Option<<T as LaunchArg>::RuntimeArg<R>>) -> Self {
69 match value {
70 Some(arg) => Self::Some(arg),
71 None => Self::None,
72 }
73 }
74}
75
76impl<T: LaunchArg + 'static + CubeType> LaunchArg for ComptimeOption<T> {
77 type RuntimeArg<R: Runtime> = ComptimeOptionArgs<T, R>;
78 type CompilationArg = ComptimeOptionCompilationArg<T>;
79
80 fn register<R: Runtime>(
81 arg: Self::RuntimeArg<R>,
82 launcher: &mut KernelLauncher<R>,
83 ) -> Self::CompilationArg {
84 match arg {
85 ComptimeOptionArgs::Some(arg) => {
86 ComptimeOptionCompilationArg::Some(T::register(arg, launcher))
87 }
88 ComptimeOptionArgs::None => ComptimeOptionCompilationArg::None,
89 }
90 }
91
92 fn expand(
93 arg: &Self::CompilationArg,
94 builder: &mut KernelBuilder,
95 ) -> <Self as CubeType>::ExpandType {
96 match arg {
97 ComptimeOptionCompilationArg::Some(arg) => {
98 ComptimeOptionExpand::Some(T::expand(arg, builder))
99 }
100 ComptimeOptionCompilationArg::None => ComptimeOptionExpand::None,
101 }
102 }
103}
104
105pub enum ComptimeOptionCompilationArg<T: LaunchArg> {
106 Some(<T as LaunchArg>::CompilationArg),
107 None,
108}
109
110impl<T: LaunchArg> Clone for ComptimeOptionCompilationArg<T> {
111 fn clone(&self) -> Self {
112 match self {
113 ComptimeOptionCompilationArg::Some(arg) => {
114 ComptimeOptionCompilationArg::Some(arg.clone())
115 }
116 ComptimeOptionCompilationArg::None => ComptimeOptionCompilationArg::None,
117 }
118 }
119}
120
121impl<T: LaunchArg> PartialEq for ComptimeOptionCompilationArg<T> {
122 fn eq(&self, other: &Self) -> bool {
123 match (self, other) {
124 (
125 ComptimeOptionCompilationArg::Some(arg_0),
126 ComptimeOptionCompilationArg::Some(arg_1),
127 ) => arg_0 == arg_1,
128 (ComptimeOptionCompilationArg::None, ComptimeOptionCompilationArg::None) => true,
129 _ => false,
130 }
131 }
132}
133
134impl<T: LaunchArg> Eq for ComptimeOptionCompilationArg<T> {}
135
136impl<T: LaunchArg> core::hash::Hash for ComptimeOptionCompilationArg<T> {
137 fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
138 match self {
139 ComptimeOptionCompilationArg::Some(arg) => {
140 arg.hash(state);
141 }
142 ComptimeOptionCompilationArg::None => {}
143 };
144 }
145}
146
147impl<T: LaunchArg> core::fmt::Debug for ComptimeOptionCompilationArg<T> {
148 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
149 match self {
150 ComptimeOptionCompilationArg::Some(arg) => f.debug_tuple("Some").field(arg).finish(),
151 ComptimeOptionCompilationArg::None => write!(f, "None"),
152 }
153 }
154}
155
156mod impls {
157 use core::ops::{Deref, DerefMut};
158
159 use super::*;
160 use ComptimeOption::Some;
161 type Option<T> = ComptimeOption<T>;
162 type OptionExpand<T> = ComptimeOptionExpand<T>;
163
164 /////////////////////////////////////////////////////////////////////////////
165 // Type implementation
166 /////////////////////////////////////////////////////////////////////////////
167
168 mod base {
169 use super::*;
170 use ComptimeOption::{None, Some};
171
172 impl<T: CubeType> ComptimeOption<T> {
173 /// Returns `true` if the option is a [`Some`] value.
174 ///
175 /// # Examples
176 ///
177 /// ```
178 /// let x: Option<u32> = Some(2);
179 /// assert_eq!(x.is_some(), true);
180 ///
181 /// let x: Option<u32> = None;
182 /// assert_eq!(x.is_some(), false);
183 /// ```
184 #[must_use = "if you intended to assert that this has a value, consider `.unwrap()` instead"]
185 pub fn is_some(&self) -> bool {
186 matches!(*self, Some(_))
187 }
188
189 /// Returns `true` if the option is a [`Some`] and the value inside of it matches a predicate.
190 ///
191 /// # Examples
192 ///
193 /// ```
194 /// let x: Option<u32> = Some(2);
195 /// assert_eq!(x.is_some_and(|x| x > 1), true);
196 ///
197 /// let x: Option<u32> = Some(0);
198 /// assert_eq!(x.is_some_and(|x| x > 1), false);
199 ///
200 /// let x: Option<u32> = None;
201 /// assert_eq!(x.is_some_and(|x| x > 1), false);
202 ///
203 /// let x: Option<String> = Some("ownership".to_string());
204 /// assert_eq!(x.as_ref().is_some_and(|x| x.len() > 1), true);
205 /// println!("still alive {:?}", x);
206 /// ```
207 #[must_use]
208 pub fn is_some_and(self, f: impl FnOnce(T) -> bool) -> bool {
209 match self {
210 None => false,
211 Some(x) => f(x),
212 }
213 }
214
215 /// Returns `true` if the option is a [`None`] or the value inside of it matches a predicate.
216 ///
217 /// # Examples
218 ///
219 /// ```
220 /// let x: Option<u32> = Some(2);
221 /// assert_eq!(x.is_none_or(|x| x > 1), true);
222 ///
223 /// let x: Option<u32> = Some(0);
224 /// assert_eq!(x.is_none_or(|x| x > 1), false);
225 ///
226 /// let x: Option<u32> = None;
227 /// assert_eq!(x.is_none_or(|x| x > 1), true);
228 ///
229 /// let x: Option<String> = Some("ownership".to_string());
230 /// assert_eq!(x.as_ref().is_none_or(|x| x.len() > 1), true);
231 /// println!("still alive {:?}", x);
232 /// ```
233 #[must_use]
234 pub fn is_none_or(self, f: impl FnOnce(T) -> bool) -> bool {
235 match self {
236 None => true,
237 Some(x) => f(x),
238 }
239 }
240
241 /// Converts from `&Option<T>` to `Option<&T>`.
242 ///
243 /// # Examples
244 ///
245 /// Calculates the length of an <code>Option<[String]></code> as an <code>Option<[usize]></code>
246 /// without moving the [`String`]. The [`map`] method takes the `self` argument by value,
247 /// consuming the original, so this technique uses `as_ref` to first take an `Option` to a
248 /// reference to the value inside the original.
249 ///
250 /// [`map`]: Option::map
251 /// [String]: ../../std/string/struct.String.html "String"
252 /// [`String`]: ../../std/string/struct.String.html "String"
253 ///
254 /// ```
255 /// let text: Option<String> = Some("Hello, world!".to_string());
256 /// // First, cast `Option<String>` to `Option<&String>` with `as_ref`,
257 /// // then consume *that* with `map`, leaving `text` on the stack.
258 /// let text_length: Option<usize> = text.as_ref().map(|s| s.len());
259 /// println!("still can print text: {text:?}");
260 /// ```
261 pub fn as_ref(&self) -> Option<&T> {
262 match *self {
263 Some(ref x) => Some(x),
264 None => None,
265 }
266 }
267
268 /// Converts from `&mut Option<T>` to `Option<&mut T>`.
269 ///
270 /// # Examples
271 ///
272 /// ```
273 /// let mut x = Some(2);
274 /// match x.as_mut() {
275 /// Some(v) => *v = 42,
276 /// None => {},
277 /// }
278 /// assert_eq!(x, Some(42));
279 /// ```
280 pub fn as_mut(&mut self) -> Option<&mut T> {
281 match *self {
282 Some(ref mut x) => Some(x),
283 None => None,
284 }
285 }
286
287 /// Returns the contained [`Some`] value, consuming the `self` value.
288 ///
289 /// # Panics
290 ///
291 /// Panics if the value is a [`None`] with a custom panic message provided by
292 /// `msg`.
293 ///
294 /// # Examples
295 ///
296 /// ```
297 /// let x = Some("value");
298 /// assert_eq!(x.expect("fruits are healthy"), "value");
299 /// ```
300 ///
301 /// ```should_panic
302 /// let x: Option<&str> = None;
303 /// x.expect("fruits are healthy"); // panics with `fruits are healthy`
304 /// ```
305 ///
306 /// # Recommended Message Style
307 ///
308 /// We recommend that `expect` messages are used to describe the reason you
309 /// _expect_ the `Option` should be `Some`.
310 ///
311 /// ```should_panic
312 /// # let slice: &[u8] = &[];
313 /// let item = slice.get(0)
314 /// .expect("slice should not be empty");
315 /// ```
316 ///
317 /// **Hint**: If you're having trouble remembering how to phrase expect
318 /// error messages remember to focus on the word "should" as in "env
319 /// variable should be set by blah" or "the given binary should be available
320 /// and executable by the current user".
321 ///
322 /// For more detail on expect message styles and the reasoning behind our
323 /// recommendation please refer to the section on ["Common Message
324 /// Styles"](../../std/error/index.html#common-message-styles) in the [`std::error`](../../std/error/index.html) module docs.
325 #[track_caller]
326 pub fn expect(self, msg: &str) -> T {
327 match self {
328 Some(val) => val,
329 None => panic!("{msg}"),
330 }
331 }
332
333 /// Returns the contained [`Some`] value, consuming the `self` value.
334 ///
335 /// Because this function may panic, its use is generally discouraged.
336 /// Panics are meant for unrecoverable errors, and
337 /// [may abort the entire program][panic-abort].
338 ///
339 /// Instead, prefer to use pattern matching and handle the [`None`]
340 /// case explicitly, or call [`unwrap_or`], [`unwrap_or_else`], or
341 /// [`unwrap_or_default`]. In functions returning `Option`, you can use
342 /// [the `?` (try) operator][try-option].
343 ///
344 /// [panic-abort]: https://doc.rust-lang.org/book/ch09-01-unrecoverable-errors-with-panic.html
345 /// [try-option]: https://doc.rust-lang.org/book/ch09-02-recoverable-errors-with-result.html#where-the--operator-can-be-used
346 /// [`unwrap_or`]: Option::unwrap_or
347 /// [`unwrap_or_else`]: Option::unwrap_or_else
348 /// [`unwrap_or_default`]: Option::unwrap_or_default
349 ///
350 /// # Panics
351 ///
352 /// Panics if the self value equals [`None`].
353 ///
354 /// # Examples
355 ///
356 /// ```
357 /// let x = Some("air");
358 /// assert_eq!(x.unwrap(), "air");
359 /// ```
360 ///
361 /// ```should_panic
362 /// let x: Option<&str> = None;
363 /// assert_eq!(x.unwrap(), "air"); // fails
364 /// ```
365 pub fn unwrap(self) -> T {
366 match self {
367 Some(val) => val,
368 None => panic!("called `Option::unwrap()` on a `None` value"),
369 }
370 }
371
372 /// Returns the contained [`Some`] value or computes it from a closure.
373 ///
374 /// # Examples
375 ///
376 /// ```
377 /// let k = 10;
378 /// assert_eq!(Some(4).unwrap_or_else(|| 2 * k), 4);
379 /// assert_eq!(None.unwrap_or_else(|| 2 * k), 20);
380 /// ```
381 pub fn unwrap_or_else<F>(self, f: F) -> T
382 where
383 F: FnOnce() -> T,
384 {
385 match self {
386 Some(x) => x,
387 None => f(),
388 }
389 }
390
391 /// Maps an `Option<T>` to `Option<U>` by applying a function to a contained value (if `Some`) or returns `None` (if `None`).
392 ///
393 /// # Examples
394 ///
395 /// Calculates the length of an <code>Option<[String]></code> as an
396 /// <code>Option<[usize]></code>, consuming the original:
397 ///
398 /// [String]: ../../std/string/struct.String.html "String"
399 /// ```
400 /// let maybe_some_string = Some(String::from("Hello, World!"));
401 /// // `Option::map` takes self *by value*, consuming `maybe_some_string`
402 /// let maybe_some_len = maybe_some_string.map(|s| s.len());
403 /// assert_eq!(maybe_some_len, Some(13));
404 ///
405 /// let x: Option<&str> = None;
406 /// assert_eq!(x.map(|s| s.len()), None);
407 /// ```
408 pub fn map<U, F>(self, f: F) -> Option<U>
409 where
410 F: FnOnce(T) -> U,
411 U: CubeType,
412 {
413 match self {
414 Some(x) => Some(f(x)),
415 None => None,
416 }
417 }
418
419 /// Calls a function with a reference to the contained value if [`Some`].
420 ///
421 /// Returns the original option.
422 ///
423 /// # Examples
424 ///
425 /// ```
426 /// let list = vec![1, 2, 3];
427 ///
428 /// // prints "got: 2"
429 /// let x = list
430 /// .get(1)
431 /// .inspect(|x| println!("got: {x}"))
432 /// .expect("list should be long enough");
433 ///
434 /// // prints nothing
435 /// list.get(5).inspect(|x| println!("got: {x}"));
436 /// ```
437 pub fn inspect<F>(self, f: F) -> Self
438 where
439 F: FnOnce(&T),
440 {
441 if let Some(ref x) = self {
442 f(x);
443 }
444
445 self
446 }
447
448 /// Returns the provided default result (if none),
449 /// or applies a function to the contained value (if any).
450 ///
451 /// Arguments passed to `map_or` are eagerly evaluated; if you are passing
452 /// the result of a function call, it is recommended to use [`map_or_else`],
453 /// which is lazily evaluated.
454 ///
455 /// [`map_or_else`]: Option::map_or_else
456 ///
457 /// # Examples
458 ///
459 /// ```
460 /// let x = Some("foo");
461 /// assert_eq!(x.map_or(42, |v| v.len()), 3);
462 ///
463 /// let x: Option<&str> = None;
464 /// assert_eq!(x.map_or(42, |v| v.len()), 42);
465 /// ```
466 pub fn map_or<U, F>(self, default: U, f: F) -> U
467 where
468 F: FnOnce(T) -> U,
469 {
470 match self {
471 Some(t) => f(t),
472 None => default,
473 }
474 }
475 /// Computes a default function result (if none), or
476 /// applies a different function to the contained value (if any).
477 ///
478 /// # Basic examples
479 ///
480 /// ```
481 /// let k = 21;
482 ///
483 /// let x = Some("foo");
484 /// assert_eq!(x.map_or_else(|| 2 * k, |v| v.len()), 3);
485 ///
486 /// let x: Option<&str> = None;
487 /// assert_eq!(x.map_or_else(|| 2 * k, |v| v.len()), 42);
488 /// ```
489 ///
490 /// # Handling a Result-based fallback
491 ///
492 /// A somewhat common occurrence when dealing with optional values
493 /// in combination with [`Result<T, E>`] is the case where one wants to invoke
494 /// a fallible fallback if the option is not present. This example
495 /// parses a command line argument (if present), or the contents of a file to
496 /// an integer. However, unlike accessing the command line argument, reading
497 /// the file is fallible, so it must be wrapped with `Ok`.
498 ///
499 /// ```no_run
500 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
501 /// let v: u64 = std::env::args()
502 /// .nth(1)
503 /// .map_or_else(|| std::fs::read_to_string("/etc/someconfig.conf"), Ok)?
504 /// .parse()?;
505 /// # Ok(())
506 /// # }
507 /// ```
508 pub fn map_or_else<U, D, F>(self, default: D, f: F) -> U
509 where
510 D: FnOnce() -> U,
511 F: FnOnce(T) -> U,
512 {
513 match self {
514 Some(t) => f(t),
515 None => default(),
516 }
517 }
518
519 /// Maps an `Option<T>` to a `U` by applying function `f` to the contained
520 /// value if the option is [`Some`], otherwise if [`None`], returns the
521 /// [default value] for the type `U`.
522 ///
523 /// # Examples
524 ///
525 /// ```ignore
526 ///
527 /// let x: Option<&str> = Some("hi");
528 /// let y: Option<&str> = None;
529 ///
530 /// assert_eq!(x.map_or_default(|x| x.len()), 2);
531 /// assert_eq!(y.map_or_default(|y| y.len()), 0);
532 /// ```
533 ///
534 /// [default value]: Default::default
535 pub fn map_or_default<U, F>(self, f: F) -> U
536 where
537 U: Default,
538 F: FnOnce(T) -> U,
539 {
540 match self {
541 Some(t) => f(t),
542 None => U::default(),
543 }
544 }
545
546 /// Converts from `Option<T>` (or `&Option<T>`) to `Option<&T::Target>`.
547 ///
548 /// Leaves the original Option in-place, creating a new one with a reference
549 /// to the original one, additionally coercing the contents via [`Deref`].
550 ///
551 /// # Examples
552 ///
553 /// ```
554 /// let x: Option<String> = Some("hey".to_owned());
555 /// assert_eq!(x.as_deref(), Some("hey"));
556 ///
557 /// let x: Option<String> = None;
558 /// assert_eq!(x.as_deref(), None);
559 /// ```
560 pub fn as_deref(&self) -> Option<&T::Target>
561 where
562 T: Deref,
563 T::Target: CubeType,
564 {
565 self.as_ref().map(Deref::deref)
566 }
567
568 /// Converts from `Option<T>` (or `&mut Option<T>`) to `Option<&mut T::Target>`.
569 ///
570 /// Leaves the original `Option` in-place, creating a new one containing a mutable reference to
571 /// the inner type's [`Deref::Target`] type.
572 ///
573 /// # Examples
574 ///
575 /// ```
576 /// let mut x: Option<String> = Some("hey".to_owned());
577 /// assert_eq!(x.as_deref_mut().map(|x| {
578 /// x.make_ascii_uppercase();
579 /// x
580 /// }), Some("HEY".to_owned().as_mut_str()));
581 /// ```
582 pub fn as_deref_mut(&mut self) -> Option<&mut T::Target>
583 where
584 T: DerefMut,
585 T::Target: CubeType,
586 {
587 self.as_mut().map(DerefMut::deref_mut)
588 }
589
590 /// Returns [`None`] if the option is [`None`], otherwise calls `f` with the
591 /// wrapped value and returns the result.
592 ///
593 /// Some languages call this operation flatmap.
594 ///
595 /// # Examples
596 ///
597 /// ```
598 /// fn sq_then_to_string(x: u32) -> Option<String> {
599 /// x.checked_mul(x).map(|sq| sq.to_string())
600 /// }
601 ///
602 /// assert_eq!(Some(2).and_then(sq_then_to_string), Some(4.to_string()));
603 /// assert_eq!(Some(1_000_000).and_then(sq_then_to_string), None); // overflowed!
604 /// assert_eq!(None.and_then(sq_then_to_string), None);
605 /// ```
606 ///
607 /// Often used to chain fallible operations that may return [`None`].
608 ///
609 /// ```
610 /// let arr_2d = [["A0", "A1"], ["B0", "B1"]];
611 ///
612 /// let item_0_1 = arr_2d.get(0).and_then(|row| row.get(1));
613 /// assert_eq!(item_0_1, Some(&"A1"));
614 ///
615 /// let item_2_0 = arr_2d.get(2).and_then(|row| row.get(0));
616 /// assert_eq!(item_2_0, None);
617 /// ```
618 pub fn and_then<U, F>(self, f: F) -> Option<U>
619 where
620 F: FnOnce(T) -> Option<U>,
621 U: CubeType,
622 {
623 match self {
624 Some(x) => f(x),
625 None => None,
626 }
627 }
628
629 /// Returns [`None`] if the option is [`None`], otherwise calls `predicate`
630 /// with the wrapped value and returns:
631 ///
632 /// - [`Some(t)`] if `predicate` returns `true` (where `t` is the wrapped
633 /// value), and
634 /// - [`None`] if `predicate` returns `false`.
635 ///
636 /// This function works similar to [`Iterator::filter()`]. You can imagine
637 /// the `Option<T>` being an iterator over one or zero elements. `filter()`
638 /// lets you decide which elements to keep.
639 ///
640 /// # Examples
641 ///
642 /// ```rust
643 /// fn is_even(n: &i32) -> bool {
644 /// n % 2 == 0
645 /// }
646 ///
647 /// assert_eq!(None.filter(is_even), None);
648 /// assert_eq!(Some(3).filter(is_even), None);
649 /// assert_eq!(Some(4).filter(is_even), Some(4));
650 /// ```
651 ///
652 /// [`Some(t)`]: Some
653 pub fn filter<P>(self, predicate: P) -> Self
654 where
655 P: FnOnce(&T) -> bool,
656 {
657 if let Some(x) = self
658 && predicate(&x)
659 {
660 return Some(x);
661 }
662 None
663 }
664
665 /// Returns the option if it contains a value, otherwise calls `f` and
666 /// returns the result.
667 ///
668 /// # Examples
669 ///
670 /// ```
671 /// fn nobody() -> Option<&'static str> { None }
672 /// fn vikings() -> Option<&'static str> { Some("vikings") }
673 ///
674 /// assert_eq!(Some("barbarians").or_else(vikings), Some("barbarians"));
675 /// assert_eq!(None.or_else(vikings), Some("vikings"));
676 /// assert_eq!(None.or_else(nobody), None);
677 /// ```
678 pub fn or_else<F>(self, f: F) -> Option<T>
679 where
680 F: FnOnce() -> Option<T>,
681 {
682 match self {
683 x @ Some(_) => x,
684 None => f(),
685 }
686 }
687
688 /// Zips `self` and another `Option` with function `f`.
689 ///
690 /// If `self` is `Some(s)` and `other` is `Some(o)`, this method returns `Some(f(s, o))`.
691 /// Otherwise, `None` is returned.
692 ///
693 /// # Examples
694 ///
695 /// ```ignore
696 ///
697 /// #[derive(Debug, PartialEq)]
698 /// struct Point {
699 /// x: f64,
700 /// y: f64,
701 /// }
702 ///
703 /// impl Point {
704 /// fn new(x: f64, y: f64) -> Self {
705 /// Self { x, y }
706 /// }
707 /// }
708 ///
709 /// let x = Some(17.5);
710 /// let y = Some(42.7);
711 ///
712 /// assert_eq!(x.zip_with(y, Point::new), Some(Point { x: 17.5, y: 42.7 }));
713 /// assert_eq!(x.zip_with(None, Point::new), None);
714 /// ```
715 pub fn zip_with<U, F, R>(self, other: Option<U>, f: F) -> Option<R>
716 where
717 F: FnOnce(T, U) -> R,
718 U: CubeType,
719 R: CubeType,
720 {
721 match (self, other) {
722 (Some(a), Some(b)) => Some(f(a, b)),
723 _ => None,
724 }
725 }
726
727 /// Reduces two options into one, using the provided function if both are `Some`.
728 ///
729 /// If `self` is `Some(s)` and `other` is `Some(o)`, this method returns `Some(f(s, o))`.
730 /// Otherwise, if only one of `self` and `other` is `Some`, that one is returned.
731 /// If both `self` and `other` are `None`, `None` is returned.
732 ///
733 /// # Examples
734 ///
735 /// ```ignore
736 ///
737 /// let s12 = Some(12);
738 /// let s17 = Some(17);
739 /// let n = None;
740 /// let f = |a, b| a + b;
741 ///
742 /// assert_eq!(s12.reduce(s17, f), Some(29));
743 /// assert_eq!(s12.reduce(n, f), Some(12));
744 /// assert_eq!(n.reduce(s17, f), Some(17));
745 /// assert_eq!(n.reduce(n, f), None);
746 /// ```
747 pub fn reduce<U, R, F>(self, other: Option<U>, f: F) -> Option<R>
748 where
749 T: Into<R>,
750 U: CubeType + Into<R>,
751 F: FnOnce(T, U) -> R,
752 R: CubeType,
753 {
754 match (self, other) {
755 (Some(a), Some(b)) => Some(f(a, b)),
756 (Some(a), _) => Some(a.into()),
757 (_, Some(b)) => Some(b.into()),
758 _ => None,
759 }
760 }
761 }
762
763 impl<T: CubeType> ComptimeOption<T> {
764 /////////////////////////////////////////////////////////////////////////
765 // Querying the contained values
766 /////////////////////////////////////////////////////////////////////////
767
768 /// Returns `true` if the option is a [`None`] value.
769 ///
770 /// # Examples
771 ///
772 /// ```
773 /// let x: Option<u32> = Some(2);
774 /// assert_eq!(x.is_none(), false);
775 ///
776 /// let x: Option<u32> = None;
777 /// assert_eq!(x.is_none(), true);
778 /// ```
779 #[must_use = "if you intended to assert that this doesn't have a value, consider \
780 wrapping this in an `assert!()` instead"]
781 pub fn is_none(&self) -> bool {
782 !self.is_some()
783 }
784
785 /////////////////////////////////////////////////////////////////////////
786 // Getting to contained values
787 /////////////////////////////////////////////////////////////////////////
788
789 /// Returns the contained [`Some`] value or a provided default.
790 ///
791 /// Arguments passed to `unwrap_or` are eagerly evaluated; if you are passing
792 /// the result of a function call, it is recommended to use [`unwrap_or_else`],
793 /// which is lazily evaluated.
794 ///
795 /// [`unwrap_or_else`]: Option::unwrap_or_else
796 ///
797 /// # Examples
798 ///
799 /// ```
800 /// assert_eq!(Some("car").unwrap_or("bike"), "car");
801 /// assert_eq!(None.unwrap_or("bike"), "bike");
802 /// ```
803 pub fn unwrap_or(self, default: T) -> T {
804 match self {
805 Some(x) => x,
806 None => default,
807 }
808 }
809
810 /// Returns the contained [`Some`] value or a default.
811 ///
812 /// Consumes the `self` argument then, if [`Some`], returns the contained
813 /// value, otherwise if [`None`], returns the [default value] for that
814 /// type.
815 ///
816 /// # Examples
817 ///
818 /// ```
819 /// let x: Option<u32> = None;
820 /// let y: Option<u32> = Some(12);
821 ///
822 /// assert_eq!(x.unwrap_or_default(), 0);
823 /// assert_eq!(y.unwrap_or_default(), 12);
824 /// ```
825 ///
826 /// [default value]: Default::default
827 /// [`parse`]: str::parse
828 /// [`FromStr`]: crate::str::FromStr
829 pub fn unwrap_or_default(self) -> T
830 where
831 T: Default + IntoRuntime,
832 {
833 match self {
834 Some(x) => x,
835 None => comptime![T::default()].runtime(),
836 }
837 }
838
839 /// Returns the contained [`Some`] value, consuming the `self` value,
840 /// without checking that the value is not [`None`].
841 ///
842 /// # Safety
843 ///
844 /// Calling this method on [`None`] is *[undefined behavior]*.
845 ///
846 /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
847 ///
848 /// # Examples
849 ///
850 /// ```
851 /// let x = Some("air");
852 /// assert_eq!(unsafe { x.unwrap_unchecked() }, "air");
853 /// ```
854 ///
855 /// ```no_run
856 /// let x: Option<&str> = None;
857 /// assert_eq!(unsafe { x.unwrap_unchecked() }, "air"); // Undefined behavior!
858 /// ```
859 pub unsafe fn unwrap_unchecked(self) -> T {
860 match self {
861 Some(val) => val,
862 // SAFETY: the safety contract must be upheld by the caller.
863 None => comptime![unsafe { core::hint::unreachable_unchecked() }],
864 }
865 }
866
867 /////////////////////////////////////////////////////////////////////////
868 // Boolean operations on the values, eager and lazy
869 /////////////////////////////////////////////////////////////////////////
870
871 /// Returns [`None`] if the option is [`None`], otherwise returns `optb`.
872 ///
873 /// Arguments passed to `and` are eagerly evaluated; if you are passing the
874 /// result of a function call, it is recommended to use [`and_then`], which is
875 /// lazily evaluated.
876 ///
877 /// [`and_then`]: Option::and_then
878 ///
879 /// # Examples
880 ///
881 /// ```
882 /// let x = Some(2);
883 /// let y: Option<&str> = None;
884 /// assert_eq!(x.and(y), None);
885 ///
886 /// let x: Option<u32> = None;
887 /// let y = Some("foo");
888 /// assert_eq!(x.and(y), None);
889 ///
890 /// let x = Some(2);
891 /// let y = Some("foo");
892 /// assert_eq!(x.and(y), Some("foo"));
893 ///
894 /// let x: Option<u32> = None;
895 /// let y: Option<&str> = None;
896 /// assert_eq!(x.and(y), None);
897 /// ```
898 pub fn and<U>(self, optb: Option<U>) -> Option<U>
899 where
900 U: CubeType,
901 {
902 match self {
903 Some(_) => optb,
904 Option::None => Option::new_None(),
905 }
906 }
907
908 /// Returns the option if it contains a value, otherwise returns `optb`.
909 ///
910 /// Arguments passed to `or` are eagerly evaluated; if you are passing the
911 /// result of a function call, it is recommended to use [`or_else`], which is
912 /// lazily evaluated.
913 ///
914 /// [`or_else`]: Option::or_else
915 ///
916 /// # Examples
917 ///
918 /// ```
919 /// let x = Some(2);
920 /// let y = None;
921 /// assert_eq!(x.or(y), Some(2));
922 ///
923 /// let x = None;
924 /// let y = Some(100);
925 /// assert_eq!(x.or(y), Some(100));
926 ///
927 /// let x = Some(2);
928 /// let y = Some(100);
929 /// assert_eq!(x.or(y), Some(2));
930 ///
931 /// let x: Option<u32> = None;
932 /// let y = None;
933 /// assert_eq!(x.or(y), None);
934 /// ```
935 pub fn or(self, optb: Option<T>) -> Option<T> {
936 match self {
937 x @ Some(_) => x,
938 None => optb,
939 }
940 }
941
942 /// Returns [`Some`] if exactly one of `self`, `optb` is [`Some`], otherwise returns [`None`].
943 ///
944 /// # Examples
945 ///
946 /// ```
947 /// let x = Some(2);
948 /// let y: Option<u32> = None;
949 /// assert_eq!(x.xor(y), Some(2));
950 ///
951 /// let x: Option<u32> = None;
952 /// let y = Some(2);
953 /// assert_eq!(x.xor(y), Some(2));
954 ///
955 /// let x = Some(2);
956 /// let y = Some(2);
957 /// assert_eq!(x.xor(y), None);
958 ///
959 /// let x: Option<u32> = None;
960 /// let y: Option<u32> = None;
961 /// assert_eq!(x.xor(y), None);
962 /// ```
963 pub fn xor(self, optb: Option<T>) -> Option<T> {
964 match (self, optb) {
965 (a @ Some(_), None) => a,
966 (None, b @ Some(_)) => b,
967 _ => Option::None,
968 }
969 }
970
971 /////////////////////////////////////////////////////////////////////////
972 // Misc
973 /////////////////////////////////////////////////////////////////////////
974
975 /// Zips `self` with another `Option`.
976 ///
977 /// If `self` is `Some(s)` and `other` is `Some(o)`, this method returns `Some((s, o))`.
978 /// Otherwise, `None` is returned.
979 ///
980 /// # Examples
981 ///
982 /// ```
983 /// let x = Some(1);
984 /// let y = Some("hi");
985 /// let z = None::<u8>;
986 ///
987 /// assert_eq!(x.zip(y), Some((1, "hi")));
988 /// assert_eq!(x.zip(z), None);
989 /// ```
990 pub fn zip<U>(self, other: Option<U>) -> Option<(T, U)>
991 where
992 U: CubeType,
993 {
994 match (self, other) {
995 (Some(a), Some(b)) => Option::Some((a, b)),
996 _ => Option::None,
997 }
998 }
999 }
1000 }
1001
1002 mod expand {
1003 use super::*;
1004 use ComptimeOptionExpand::{None, Some};
1005
1006 #[doc(hidden)]
1007 impl<T: CubeType> ComptimeOptionExpand<T> {
1008 pub fn __expand_is_some_method(&self, _scope: &Scope) -> bool {
1009 matches!(*self, Some(_))
1010 }
1011
1012 pub fn __expand_is_some_and_method(
1013 self,
1014 scope: &Scope,
1015 f: impl FnOnce(&Scope, T::ExpandType) -> bool,
1016 ) -> bool {
1017 match self {
1018 None => false,
1019 Some(x) => f(scope, x),
1020 }
1021 }
1022
1023 pub fn __expand_is_none_or_method(
1024 self,
1025 scope: &Scope,
1026 f: impl FnOnce(&Scope, T::ExpandType) -> bool,
1027 ) -> bool {
1028 match self {
1029 None => true,
1030 Some(x) => f(scope, x),
1031 }
1032 }
1033
1034 fn __expand_len_method(&self, _scope: &Scope) -> usize {
1035 match self {
1036 Some(_) => 1,
1037 None => 0,
1038 }
1039 }
1040
1041 pub fn __expand_expect_method(self, _scope: &Scope, msg: &str) -> T::ExpandType {
1042 match self {
1043 Some(val) => val,
1044 None => panic!("{msg}"),
1045 }
1046 }
1047
1048 #[allow(clippy::unnecessary_literal_unwrap)]
1049 pub fn __expand_unwrap_method(self, _scope: &Scope) -> T::ExpandType {
1050 match self {
1051 Some(val) => val,
1052 None => core::option::Option::None.unwrap(),
1053 }
1054 }
1055
1056 pub fn __expand_unwrap_or_else_method<F>(self, scope: &Scope, f: F) -> T::ExpandType
1057 where
1058 F: FnOnce(&Scope) -> T::ExpandType,
1059 {
1060 match self {
1061 Some(x) => x,
1062 None => f(scope),
1063 }
1064 }
1065
1066 pub fn __expand_map_method<U, F>(self, scope: &Scope, f: F) -> ComptimeOptionExpand<U>
1067 where
1068 U: CubeType,
1069 F: FnOnce(&Scope, T::ExpandType) -> U::ExpandType,
1070 {
1071 match self {
1072 Some(x) => Some(f(scope, x)),
1073 None => None,
1074 }
1075 }
1076
1077 pub fn __expand_as_ref_method(&self, _scope: &Scope) -> ComptimeOptionExpand<&T> {
1078 match self {
1079 Some(x) => Some(x),
1080 None => None,
1081 }
1082 }
1083
1084 pub fn __expand_as_mut_method(
1085 &mut self,
1086 _scope: &Scope,
1087 ) -> ComptimeOptionExpand<&mut T> {
1088 match self {
1089 Some(x) => Some(x),
1090 None => None,
1091 }
1092 }
1093
1094 pub fn __expand_inspect_method<F>(self, scope: &Scope, f: F) -> Self
1095 where
1096 F: FnOnce(&Scope, &T::ExpandType),
1097 {
1098 if let Some(x) = &self {
1099 f(scope, x);
1100 }
1101
1102 self
1103 }
1104
1105 pub fn __expand_map_or_method<U, F>(
1106 self,
1107 scope: &Scope,
1108 default: U::ExpandType,
1109 f: F,
1110 ) -> U::ExpandType
1111 where
1112 F: FnOnce(&Scope, T::ExpandType) -> U::ExpandType,
1113 U: CubeType,
1114 {
1115 match self {
1116 Some(t) => f(scope, t),
1117 None => default,
1118 }
1119 }
1120
1121 pub fn __expand_map_or_else_method<U, D, F>(
1122 self,
1123 scope: &Scope,
1124 default: D,
1125 f: F,
1126 ) -> U::ExpandType
1127 where
1128 U: CubeType,
1129 D: FnOnce(&Scope) -> U::ExpandType,
1130 F: FnOnce(&Scope, T::ExpandType) -> U::ExpandType,
1131 {
1132 match self {
1133 Some(t) => f(scope, t),
1134 None => default(scope),
1135 }
1136 }
1137
1138 pub fn __expand_map_or_default_method<U, F>(self, scope: &Scope, f: F) -> U::ExpandType
1139 where
1140 U: CubeType + Default + Into<U::ExpandType>,
1141 F: FnOnce(&Scope, T::ExpandType) -> U::ExpandType,
1142 {
1143 match self {
1144 Some(t) => f(scope, t),
1145 None => U::default().into(),
1146 }
1147 }
1148
1149 pub fn __expand_as_deref_method(self, scope: &Scope) -> ComptimeOptionExpand<T::Target>
1150 where
1151 T: Deref<Target: CubeType + Sized>,
1152 T::ExpandType: DerefExpand<Target = <T::Target as CubeType>::ExpandType>,
1153 {
1154 self.__expand_map_method(scope, |scope, it| it.__expand_deref_method(scope))
1155 }
1156
1157 pub fn __expand_as_deref_mut_method(
1158 self,
1159 scope: &Scope,
1160 ) -> ComptimeOptionExpand<T::Target>
1161 where
1162 T: DerefMut<Target: CubeType + Sized>,
1163 T::ExpandType: DerefExpand<Target = <T::Target as CubeType>::ExpandType>,
1164 {
1165 self.__expand_map_method(scope, |scope, it| it.__expand_deref_method(scope))
1166 }
1167
1168 pub fn __expand_and_then_method<U, F>(
1169 self,
1170 scope: &Scope,
1171 f: F,
1172 ) -> ComptimeOptionExpand<U>
1173 where
1174 U: CubeType,
1175 F: FnOnce(&Scope, T::ExpandType) -> ComptimeOptionExpand<U>,
1176 {
1177 match self {
1178 Some(x) => f(scope, x),
1179 None => None,
1180 }
1181 }
1182
1183 pub fn __expand_filter_method<P>(self, scope: &Scope, predicate: P) -> Self
1184 where
1185 P: FnOnce(&Scope, &T::ExpandType) -> bool,
1186 {
1187 if let Some(x) = self
1188 && predicate(scope, &x)
1189 {
1190 Some(x)
1191 } else {
1192 None
1193 }
1194 }
1195
1196 pub fn __expand_or_else_method<F>(self, scope: &Scope, f: F) -> ComptimeOptionExpand<T>
1197 where
1198 F: FnOnce(&Scope) -> ComptimeOptionExpand<T>,
1199 {
1200 match self {
1201 x @ Some(_) => x,
1202 None => f(scope),
1203 }
1204 }
1205
1206 // Entry methods that return &mut T excluded for now
1207
1208 pub fn __expand_take_method(&mut self, _scope: &Scope) -> ComptimeOptionExpand<T> {
1209 core::mem::take(self)
1210 }
1211
1212 pub fn __expand_take_if_method<P>(
1213 &mut self,
1214 scope: &Scope,
1215 predicate: P,
1216 ) -> ComptimeOptionExpand<T>
1217 where
1218 P: FnOnce(&Scope, &mut T::ExpandType) -> bool,
1219 {
1220 match self {
1221 Some(value) => {
1222 if predicate(scope, value) {
1223 self.__expand_take_method(scope)
1224 } else {
1225 None
1226 }
1227 }
1228 _ => None,
1229 }
1230 }
1231
1232 pub fn __expand_replace_method(
1233 &mut self,
1234 _scope: &Scope,
1235 value: T::ExpandType,
1236 ) -> ComptimeOptionExpand<T> {
1237 core::mem::replace(self, Some(value))
1238 }
1239
1240 pub fn __expand_zip_with_method<U, F, R>(
1241 self,
1242 scope: &Scope,
1243 other: ComptimeOptionExpand<U>,
1244 f: F,
1245 ) -> ComptimeOptionExpand<R>
1246 where
1247 F: FnOnce(&Scope, T::ExpandType, U::ExpandType) -> R::ExpandType,
1248 R: CubeType,
1249 U: CubeType,
1250 {
1251 match (self, other) {
1252 (Some(a), Some(b)) => Some(f(scope, a, b)),
1253 _ => None,
1254 }
1255 }
1256
1257 pub fn __expand_reduce_method<U, R, F>(
1258 self,
1259 scope: &Scope,
1260 other: ComptimeOptionExpand<U>,
1261 f: F,
1262 ) -> ComptimeOptionExpand<R>
1263 where
1264 U: CubeType,
1265 R: CubeType,
1266 T::ExpandType: Into<R::ExpandType>,
1267 U::ExpandType: Into<R::ExpandType>,
1268 F: FnOnce(&Scope, T::ExpandType, U::ExpandType) -> R::ExpandType,
1269 {
1270 match (self, other) {
1271 (Some(a), Some(b)) => Some(f(scope, a, b)),
1272 (Some(a), _) => Some(a.into()),
1273 (_, Some(b)) => Some(b.into()),
1274 _ => None,
1275 }
1276 }
1277 }
1278
1279 impl<T: CubeType> ComptimeOptionExpand<T> {
1280 pub fn __expand_is_none_method(self, scope: &Scope) -> bool {
1281 !self.__expand_is_some_method(scope)
1282 }
1283 pub fn __expand_unwrap_or_method(
1284 self,
1285 _scope: &Scope,
1286 default: <T as cubecl::prelude::CubeType>::ExpandType,
1287 ) -> <T as cubecl::prelude::CubeType>::ExpandType {
1288 {
1289 match self {
1290 OptionExpand::Some(x) => x,
1291 OptionExpand::None => default,
1292 }
1293 }
1294 }
1295 pub fn __expand_unwrap_or_default_method(
1296 self,
1297 scope: &Scope,
1298 ) -> <T as cubecl::prelude::CubeType>::ExpandType
1299 where
1300 T: Default + IntoRuntime,
1301 {
1302 {
1303 match self {
1304 OptionExpand::Some(x) => x,
1305 OptionExpand::None => { T::default() }.__expand_runtime_method(scope),
1306 }
1307 }
1308 }
1309 pub fn __expand_unwrap_unchecked_method(
1310 self,
1311 _scope: &Scope,
1312 ) -> <T as cubecl::prelude::CubeType>::ExpandType {
1313 {
1314 match self {
1315 OptionExpand::Some(val) => val,
1316 OptionExpand::None => unsafe { core::hint::unreachable_unchecked() },
1317 }
1318 }
1319 }
1320 pub fn __expand_and_method<U>(
1321 self,
1322 scope: &Scope,
1323 optb: <Option<U> as cubecl::prelude::CubeType>::ExpandType,
1324 ) -> <Option<U> as cubecl::prelude::CubeType>::ExpandType
1325 where
1326 U: CubeType,
1327 {
1328 {
1329 match self {
1330 OptionExpand::Some(_) => optb,
1331 OptionExpand::None => Option::__expand_new_None(scope),
1332 }
1333 }
1334 }
1335 pub fn __expand_or_method(
1336 self,
1337 _scope: &Scope,
1338 optb: <Option<T> as cubecl::prelude::CubeType>::ExpandType,
1339 ) -> <Option<T> as cubecl::prelude::CubeType>::ExpandType {
1340 {
1341 match self {
1342 x @ OptionExpand::Some(_) => x,
1343 OptionExpand::None => optb,
1344 }
1345 }
1346 }
1347 pub fn __expand_xor_method(
1348 self,
1349 scope: &Scope,
1350 optb: <Option<T> as cubecl::prelude::CubeType>::ExpandType,
1351 ) -> <Option<T> as cubecl::prelude::CubeType>::ExpandType {
1352 {
1353 match (self, optb) {
1354 (a @ OptionExpand::Some(_), OptionExpand::None) => a,
1355 (OptionExpand::None, b @ OptionExpand::Some(_)) => b,
1356 _ => Option::__expand_new_None(scope),
1357 }
1358 }
1359 }
1360 pub fn __expand_zip_method<U>(
1361 self,
1362 scope: &Scope,
1363 other: <Option<U> as cubecl::prelude::CubeType>::ExpandType,
1364 ) -> <Option<(T, U)> as cubecl::prelude::CubeType>::ExpandType
1365 where
1366 U: CubeType,
1367 {
1368 {
1369 match (self, other) {
1370 (OptionExpand::Some(a), OptionExpand::Some(b)) => {
1371 let _arg_0 = (a, b);
1372 Option::__expand_Some(scope, _arg_0)
1373 }
1374 _ => Option::__expand_new_None(scope),
1375 }
1376 }
1377 }
1378 }
1379 }
1380
1381 impl<T: CubeType, U: CubeType> ComptimeOption<(T, U)> {
1382 /// Unzips an option containing a tuple of two options.
1383 ///
1384 /// If `self` is `Some((a, b))` this method returns `(Some(a), Some(b))`.
1385 /// Otherwise, `(None, None)` is returned.
1386 ///
1387 /// # Examples
1388 ///
1389 /// ```
1390 /// let x = Some((1, "hi"));
1391 /// let y = None::<(u8, u32)>;
1392 ///
1393 /// assert_eq!(x.unzip(), (Some(1), Some("hi")));
1394 /// assert_eq!(y.unzip(), (None, None));
1395 /// ```
1396 pub fn unzip(self) -> (Option<T>, Option<U>) {
1397 match self {
1398 Some((a, b)) => (Option::Some(a), Option::Some(b)),
1399 Option::None => (Option::None, Option::None),
1400 }
1401 }
1402 }
1403
1404 impl<T: CubeType, U: CubeType> ComptimeOptionExpand<(T, U)> {
1405 pub fn __expand_unzip_method(
1406 self,
1407 scope: &Scope,
1408 ) -> <(Option<T>, Option<U>) as cubecl::prelude::CubeType>::ExpandType {
1409 {
1410 match self {
1411 OptionExpand::Some((a, b)) => (
1412 {
1413 let _arg_0 = a;
1414 Option::__expand_Some(scope, _arg_0)
1415 },
1416 {
1417 let _arg_0 = b;
1418 Option::__expand_Some(scope, _arg_0)
1419 },
1420 ),
1421 OptionExpand::None => ({ Option::__expand_new_None(scope) }, {
1422 Option::__expand_new_None(scope)
1423 }),
1424 }
1425 }
1426 }
1427 }
1428}