nxpu_ir/
func.rs

1//! Functions, entry points, and local variables.
2
3use std::collections::HashMap;
4
5use crate::arena::{Arena, Handle};
6use crate::expr::Expression;
7use crate::global::Binding;
8use crate::stmt::Block;
9use crate::types::Type;
10
11/// A function argument declaration.
12#[derive(Clone, Debug)]
13pub struct FunctionArgument {
14    /// Optional argument name.
15    pub name: Option<String>,
16    /// The type of this argument.
17    pub ty: Handle<Type>,
18    /// Optional binding (e.g. built-in or location).
19    pub binding: Option<Binding>,
20}
21
22/// The return type and optional binding of a function.
23#[derive(Clone, Debug)]
24pub struct FunctionResult {
25    /// The return type.
26    pub ty: Handle<Type>,
27    /// Optional binding for the return value.
28    pub binding: Option<Binding>,
29}
30
31/// A function-local variable.
32#[derive(Clone, Debug)]
33pub struct LocalVariable {
34    /// Optional variable name.
35    pub name: Option<String>,
36    /// The type of this variable.
37    pub ty: Handle<Type>,
38    /// Optional initializer expression.
39    pub init: Option<Handle<Expression>>,
40}
41
42/// An IR function.
43#[derive(Clone, Debug)]
44pub struct Function {
45    /// Optional function name.
46    pub name: Option<String>,
47    /// Formal parameters.
48    pub arguments: Vec<FunctionArgument>,
49    /// Return type and optional binding.
50    pub result: Option<FunctionResult>,
51    /// Function-local variable declarations.
52    pub local_variables: Arena<LocalVariable>,
53    /// Expression arena for this function.
54    pub expressions: Arena<Expression>,
55    /// Map from expression handles to user-defined names.
56    pub named_expressions: HashMap<Handle<Expression>, String>,
57    /// The function body.
58    pub body: Block,
59}
60
61impl Function {
62    /// Creates an empty function with the given name.
63    pub fn new(name: impl Into<String>) -> Self {
64        Self {
65            name: Some(name.into()),
66            arguments: Vec::new(),
67            result: None,
68            local_variables: Arena::new(),
69            expressions: Arena::new(),
70            named_expressions: HashMap::new(),
71            body: Vec::new(),
72        }
73    }
74}
75
76/// A compute shader entry point.
77#[derive(Clone, Debug)]
78pub struct EntryPoint {
79    /// Entry point name (matches the WGSL function name).
80    pub name: String,
81    /// Workgroup dimensions `[x, y, z]`.
82    pub workgroup_size: [u32; 3],
83    /// The entry point function body.
84    pub function: Function,
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use crate::expr::Literal;
91    use crate::types::{Scalar, TypeInner};
92
93    #[test]
94    fn function_new() {
95        let f = Function::new("test");
96        assert_eq!(f.name.as_deref(), Some("test"));
97        assert_eq!(f.arguments.len(), 0);
98        assert!(f.result.is_none());
99        assert_eq!(f.body.len(), 0);
100        assert_eq!(f.expressions.len(), 0);
101    }
102
103    #[test]
104    fn function_with_local_vars() {
105        let mut types = crate::arena::UniqueArena::new();
106        let f32_ty = types.insert(Type {
107            name: None,
108            inner: TypeInner::Scalar(Scalar::F32),
109        });
110
111        let mut f = Function::new("test");
112        let init = f.expressions.append(Expression::Literal(Literal::F32(0.0)));
113        let _var = f.local_variables.append(LocalVariable {
114            name: Some("sum".into()),
115            ty: f32_ty,
116            init: Some(init),
117        });
118        assert_eq!(f.local_variables.len(), 1);
119    }
120
121    #[test]
122    fn entry_point() {
123        let ep = EntryPoint {
124            name: "main".into(),
125            workgroup_size: [256, 1, 1],
126            function: Function::new("main"),
127        };
128        assert_eq!(ep.workgroup_size, [256, 1, 1]);
129    }
130}
nxpu_ir/func.rs

nxpu_ir/
func.rs
