cubecl_ir/
branch.rs

1use alloc::{boxed::Box, format, vec::Vec};
2use core::fmt::Display;
3
4use crate::{OperationArgs, OperationReflect};
5
6use super::{OperationCode, Scope, Value};
7use crate::TypeHash;
8
9/// All branching types.
10#[allow(clippy::large_enum_variant)]
11#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
12#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationCode)]
13#[operation(opcode_name = BranchOpCode)]
14pub enum Branch {
15    /// An if statement.
16    If(Box<If>),
17    /// An if else statement.
18    IfElse(Box<IfElse>),
19    /// A switch statement
20    Switch(Box<Switch>),
21    /// A range loop.
22    RangeLoop(Box<RangeLoop>),
23    /// A loop.
24    Loop(Box<Loop>),
25    /// A return statement.
26    Return,
27    /// A break statement.
28    Break,
29    /// Unreachable block end (equivalent to `unreachable_unchecked()`)
30    Unreachable,
31}
32
33impl OperationReflect for Branch {
34    type OpCode = BranchOpCode;
35
36    fn op_code(&self) -> Self::OpCode {
37        self.__match_opcode()
38    }
39
40    fn sanitize_args(&mut self, scope: &Scope) {
41        match self {
42            Branch::If(if_) => if_.cond.sanitize_args_ptr(scope),
43            Branch::IfElse(if_else) => if_else.cond.sanitize_args_ptr(scope),
44            Branch::Switch(switch) => switch.value.sanitize_args_ptr(scope),
45            Branch::RangeLoop(range_loop) => range_loop.sanitize_args_ptr(scope),
46            Branch::Loop(_) => {}
47            Branch::Return | Branch::Break | Branch::Unreachable => {}
48        }
49    }
50}
51
52impl Display for Branch {
53    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
54        match self {
55            Branch::If(if_) => write!(f, "if({}) {}", if_.cond, if_.scope),
56            Branch::IfElse(if_else) => write!(
57                f,
58                "if({}) {} else {}",
59                if_else.cond, if_else.scope_if, if_else.scope_else
60            ),
61            Branch::Switch(switch) => write!(
62                f,
63                "switch({}) {:?}",
64                switch.value,
65                switch
66                    .cases
67                    .iter()
68                    .map(|case| format!("{}", case.0))
69                    .collect::<Vec<_>>(),
70            ),
71            Branch::RangeLoop(range_loop) => write!(
72                f,
73                "for({} in {}{}{}) {}",
74                range_loop.i,
75                range_loop.start,
76                if range_loop.inclusive { "..=" } else { ".." },
77                range_loop.end,
78                range_loop.scope
79            ),
80            Branch::Loop(loop_) => write!(f, "loop {}", loop_.scope),
81            Branch::Return => write!(f, "return"),
82            Branch::Break => write!(f, "break"),
83            Branch::Unreachable => write!(f, "unreachable"),
84        }
85    }
86}
87
88#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
89#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
90#[allow(missing_docs)]
91pub struct If {
92    pub cond: Value,
93    pub scope: Scope,
94}
95
96#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
97#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
98#[allow(missing_docs)]
99pub struct IfElse {
100    pub cond: Value,
101    pub scope_if: Scope,
102    pub scope_else: Scope,
103}
104
105#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
106#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
107#[allow(missing_docs)]
108pub struct Switch {
109    pub value: Value,
110    pub scope_default: Scope,
111    pub cases: Vec<(Value, Scope)>,
112}
113
114#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
115#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
116#[allow(missing_docs)]
117pub struct RangeLoop {
118    pub i: Value,
119    pub start: Value,
120    pub end: Value,
121    pub step: Option<Value>,
122    pub inclusive: bool,
123    pub scope: Scope,
124}
125
126impl RangeLoop {
127    pub fn sanitize_args_ptr(&mut self, scope: &Scope) {
128        self.start.sanitize_args_ptr(scope);
129        self.end.sanitize_args_ptr(scope);
130        self.step.sanitize_args_ptr(scope);
131    }
132}
133
134#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
135#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash)]
136#[allow(missing_docs)]
137pub struct Loop {
138    pub scope: Scope,
139}
cubecl_ir/branch.rs

cubecl_ir/
branch.rs
