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