Skip to main content

nxpu_ir/
stmt.rs

1//! Statements — operations with side effects and control flow.
2
3use crate::arena::{Handle, Range};
4use crate::expr::{AtomicFunction, Expression};
5
6/// A block of statements.
7pub type Block = Vec<Statement>;
8
9/// Bitflags for synchronization barriers.
10#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
11pub struct Barrier(u32);
12
13impl Barrier {
14    /// Empty barrier (no flags set).
15    pub const EMPTY: Self = Self(0);
16    /// Storage buffer barrier.
17    pub const STORAGE: Self = Self(1);
18    /// Workgroup memory barrier.
19    pub const WORKGROUP: Self = Self(2);
20
21    /// Returns `true` if `self` contains all flags in `other`.
22    pub fn contains(self, other: Self) -> bool {
23        self.0 & other.0 == other.0
24    }
25
26    /// Returns `true` if no flags are set.
27    pub fn is_empty(self) -> bool {
28        self.0 == 0
29    }
30}
31
32impl std::ops::BitOr for Barrier {
33    type Output = Self;
34    fn bitor(self, rhs: Self) -> Self {
35        Self(self.0 | rhs.0)
36    }
37}
38
39impl std::ops::BitOrAssign for Barrier {
40    fn bitor_assign(&mut self, rhs: Self) {
41        self.0 |= rhs.0;
42    }
43}
44
45/// A statement in the IR.
46///
47/// Statements have side effects and/or control flow.
48/// They operate on expressions referenced by handles.
49#[derive(Clone, Debug)]
50pub enum Statement {
51    /// Mark a range of expressions as producing live values.
52    Emit(Range<Expression>),
53    /// Write a value through a pointer.
54    Store {
55        pointer: Handle<Expression>,
56        value: Handle<Expression>,
57    },
58    /// Conditional branch.
59    If {
60        condition: Handle<Expression>,
61        accept: Block,
62        reject: Block,
63    },
64    /// Unified loop construct (handles for/while/loop).
65    Loop {
66        body: Block,
67        continuing: Block,
68        break_if: Option<Handle<Expression>>,
69    },
70    /// Call a function.
71    Call {
72        function: Handle<crate::Function>,
73        arguments: Vec<Handle<Expression>>,
74        result: Option<Handle<Expression>>,
75    },
76    /// Perform an atomic operation.
77    Atomic {
78        pointer: Handle<Expression>,
79        fun: AtomicFunction,
80        value: Handle<Expression>,
81        result: Option<Handle<Expression>>,
82    },
83    /// Break out of the innermost loop.
84    Break,
85    /// Continue to the next iteration of the innermost loop.
86    Continue,
87    /// Return from the function.
88    Return { value: Option<Handle<Expression>> },
89    /// Synchronization barrier.
90    Barrier(Barrier),
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use crate::arena::Arena;
97    use crate::expr::Literal;
98
99    #[test]
100    fn barrier_flags() {
101        let storage = Barrier::STORAGE;
102        let workgroup = Barrier::WORKGROUP;
103        let both = storage | workgroup;
104        assert!(both.contains(storage));
105        assert!(both.contains(workgroup));
106        assert!(!storage.contains(workgroup));
107    }
108
109    #[test]
110    fn build_if_statement() {
111        let mut exprs = Arena::new();
112        let cond = exprs.append(Expression::Literal(Literal::Bool(true)));
113        let stmt = Statement::If {
114            condition: cond,
115            accept: vec![Statement::Break],
116            reject: vec![],
117        };
118        if let Statement::If { accept, reject, .. } = &stmt {
119            assert_eq!(accept.len(), 1);
120            assert_eq!(reject.len(), 0);
121        } else {
122            panic!("expected If");
123        }
124    }
125
126    #[test]
127    fn build_loop_statement() {
128        let stmt = Statement::Loop {
129            body: vec![Statement::Continue],
130            continuing: vec![],
131            break_if: None,
132        };
133        if let Statement::Loop {
134            body, continuing, ..
135        } = &stmt
136        {
137            assert_eq!(body.len(), 1);
138            assert_eq!(continuing.len(), 0);
139        } else {
140            panic!("expected Loop");
141        }
142    }
143}