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