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