nxpu_ir/
display.rs

1//! Display implementations and text dump for debugging.
2
3use std::fmt;
4
5use crate::Module;
6use crate::arena::{Handle, UniqueArena};
7use crate::expr::{
8    AtomicFunction, BinaryOp, Expression, Literal, MathFunction, SwizzleComponent, UnaryOp,
9};
10use crate::global::{AddressSpace, Binding, BuiltIn, ResourceBinding, StorageAccess};
11use crate::stmt::{Barrier, Statement};
12use crate::types::{ArraySize, Scalar, ScalarKind, Type, TypeInner, VectorSize};
13
14impl fmt::Display for ScalarKind {
15    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
16        match self {
17            Self::Bool => write!(f, "bool"),
18            Self::Sint => write!(f, "sint"),
19            Self::Uint => write!(f, "uint"),
20            Self::Float => write!(f, "float"),
21            Self::BFloat => write!(f, "bfloat"),
22        }
23    }
24}
25
26impl fmt::Display for Scalar {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        match self.kind {
29            ScalarKind::Bool => write!(f, "bool"),
30            ScalarKind::Sint => write!(f, "i{}", self.width * 8),
31            ScalarKind::Uint => write!(f, "u{}", self.width * 8),
32            ScalarKind::Float => write!(f, "f{}", self.width * 8),
33            ScalarKind::BFloat => write!(f, "bf{}", self.width * 8),
34        }
35    }
36}
37
38impl fmt::Display for VectorSize {
39    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
40        write!(f, "{}", *self as u32)
41    }
42}
43
44impl fmt::Display for StorageAccess {
45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46        let has_load = self.contains(StorageAccess::LOAD);
47        let has_store = self.contains(StorageAccess::STORE);
48        match (has_load, has_store) {
49            (true, true) => write!(f, "read_write"),
50            (true, false) => write!(f, "read"),
51            (false, true) => write!(f, "write"),
52            (false, false) => write!(f, "none"),
53        }
54    }
55}
56
57impl fmt::Display for AddressSpace {
58    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59        match self {
60            Self::Function => write!(f, "function"),
61            Self::Private => write!(f, "private"),
62            Self::Workgroup => write!(f, "workgroup"),
63            Self::Uniform => write!(f, "uniform"),
64            Self::Storage { access } => write!(f, "storage, {access}"),
65        }
66    }
67}
68
69impl fmt::Display for BuiltIn {
70    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71        match self {
72            Self::GlobalInvocationId => write!(f, "global_invocation_id"),
73            Self::LocalInvocationId => write!(f, "local_invocation_id"),
74            Self::LocalInvocationIndex => write!(f, "local_invocation_index"),
75            Self::WorkgroupId => write!(f, "workgroup_id"),
76            Self::NumWorkgroups => write!(f, "num_workgroups"),
77        }
78    }
79}
80
81impl fmt::Display for Binding {
82    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
83        match self {
84            Self::BuiltIn(b) => write!(f, "@builtin({b})"),
85            Self::Location { location } => write!(f, "@location({location})"),
86        }
87    }
88}
89
90impl fmt::Display for ResourceBinding {
91    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92        write!(f, "@group({}) @binding({})", self.group, self.binding)
93    }
94}
95
96impl fmt::Display for Literal {
97    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
98        match self {
99            Self::Bool(v) => write!(f, "{v}"),
100            Self::I32(v) => write!(f, "{v}i"),
101            Self::U32(v) => write!(f, "{v}u"),
102            Self::F32(v) => write!(f, "{v}f"),
103            Self::F64(v) => write!(f, "{v}lf"),
104            Self::AbstractInt(v) => write!(f, "{v}"),
105            Self::AbstractFloat(v) => write!(f, "{v}"),
106        }
107    }
108}
109
110impl fmt::Display for UnaryOp {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        match self {
113            Self::Negate => write!(f, "-"),
114            Self::LogicalNot => write!(f, "!"),
115            Self::BitwiseNot => write!(f, "~"),
116        }
117    }
118}
119
120impl fmt::Display for BinaryOp {
121    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
122        match self {
123            Self::Add => write!(f, "+"),
124            Self::Subtract => write!(f, "-"),
125            Self::Multiply => write!(f, "*"),
126            Self::Divide => write!(f, "/"),
127            Self::Modulo => write!(f, "%"),
128            Self::Equal => write!(f, "=="),
129            Self::NotEqual => write!(f, "!="),
130            Self::Less => write!(f, "<"),
131            Self::LessEqual => write!(f, "<="),
132            Self::Greater => write!(f, ">"),
133            Self::GreaterEqual => write!(f, ">="),
134            Self::LogicalAnd => write!(f, "&&"),
135            Self::LogicalOr => write!(f, "||"),
136            Self::BitwiseAnd => write!(f, "&"),
137            Self::BitwiseOr => write!(f, "|"),
138            Self::BitwiseXor => write!(f, "^"),
139            Self::ShiftLeft => write!(f, "<<"),
140            Self::ShiftRight => write!(f, ">>"),
141        }
142    }
143}
144
145impl fmt::Display for MathFunction {
146    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
147        let name = match self {
148            Self::Abs => "abs",
149            Self::Min => "min",
150            Self::Max => "max",
151            Self::Clamp => "clamp",
152            Self::Saturate => "saturate",
153            Self::Floor => "floor",
154            Self::Ceil => "ceil",
155            Self::Round => "round",
156            Self::Fract => "fract",
157            Self::Trunc => "trunc",
158            Self::Sin => "sin",
159            Self::Cos => "cos",
160            Self::Tan => "tan",
161            Self::Asin => "asin",
162            Self::Acos => "acos",
163            Self::Atan => "atan",
164            Self::Atan2 => "atan2",
165            Self::Sinh => "sinh",
166            Self::Cosh => "cosh",
167            Self::Tanh => "tanh",
168            Self::Sqrt => "sqrt",
169            Self::InverseSqrt => "inverseSqrt",
170            Self::Log => "log",
171            Self::Log2 => "log2",
172            Self::Exp => "exp",
173            Self::Exp2 => "exp2",
174            Self::Pow => "pow",
175            Self::Dot => "dot",
176            Self::Cross => "cross",
177            Self::Normalize => "normalize",
178            Self::Length => "length",
179            Self::Distance => "distance",
180            Self::Mix => "mix",
181            Self::Step => "step",
182            Self::SmoothStep => "smoothStep",
183            Self::Fma => "fma",
184            Self::ExtractBits => "extractBits",
185            Self::InsertBits => "insertBits",
186        };
187        write!(f, "{name}")
188    }
189}
190
191impl fmt::Display for Barrier {
192    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
193        let storage = self.contains(Barrier::STORAGE);
194        let workgroup = self.contains(Barrier::WORKGROUP);
195        match (storage, workgroup) {
196            (true, true) => write!(f, "storageBarrier | workgroupBarrier"),
197            (true, false) => write!(f, "storageBarrier"),
198            (false, true) => write!(f, "workgroupBarrier"),
199            (false, false) => write!(f, "<no barrier>"),
200        }
201    }
202}
203
204impl fmt::Display for SwizzleComponent {
205    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
206        match self {
207            Self::X => write!(f, "x"),
208            Self::Y => write!(f, "y"),
209            Self::Z => write!(f, "z"),
210            Self::W => write!(f, "w"),
211        }
212    }
213}
214
215impl fmt::Display for AtomicFunction {
216    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
217        match self {
218            Self::Add => write!(f, "atomicAdd"),
219            Self::Subtract => write!(f, "atomicSub"),
220            Self::And => write!(f, "atomicAnd"),
221            Self::ExclusiveOr => write!(f, "atomicXor"),
222            Self::InclusiveOr => write!(f, "atomicOr"),
223            Self::Min => write!(f, "atomicMin"),
224            Self::Max => write!(f, "atomicMax"),
225            Self::Exchange { compare: None } => write!(f, "atomicExchange"),
226            Self::Exchange { compare: Some(c) } => write!(f, "atomicCompareExchange({c:?})"),
227        }
228    }
229}
230
231/// Formats a type using the type arena for resolving inner references.
232pub fn format_type(ty: &Type, types: &UniqueArena<Type>) -> String {
233    if let Some(ref name) = ty.name {
234        return name.clone();
235    }
236    format_type_inner(&ty.inner, types)
237}
238
239/// Formats a [`TypeInner`] using the type arena for resolving references.
240pub fn format_type_inner(inner: &TypeInner, types: &UniqueArena<Type>) -> String {
241    match inner {
242        TypeInner::Scalar(s) => format!("{s}"),
243        TypeInner::Vector { size, scalar } => format!("vec{size}<{scalar}>"),
244        TypeInner::Matrix {
245            columns,
246            rows,
247            scalar,
248        } => format!("mat{columns}x{rows}<{scalar}>"),
249        TypeInner::Atomic(s) => format!("atomic<{s}>"),
250        TypeInner::Pointer { base, space } => {
251            let base_str = format_type(&types[*base], types);
252            format!("ptr<{space}, {base_str}>")
253        }
254        TypeInner::Array { base, size, stride } => {
255            let base_str = format_type(&types[*base], types);
256            match size {
257                ArraySize::Constant(n) => format!("array<{base_str}, {n}> /*stride {stride}*/"),
258                ArraySize::Dynamic => format!("array<{base_str}> /*stride {stride}*/"),
259            }
260        }
261        TypeInner::Struct { members, span } => {
262            format!("struct({} members, span {span})", members.len())
263        }
264        TypeInner::Tensor { scalar, shape } => {
265            let dims: Vec<String> = shape
266                .dims
267                .iter()
268                .map(|d| match d {
269                    crate::Dimension::Fixed(n) => n.to_string(),
270                    crate::Dimension::Symbolic(name) => name.clone(),
271                    crate::Dimension::Dynamic(Some(name)) => name.clone(),
272                    crate::Dimension::Dynamic(None) => "?".into(),
273                })
274                .collect();
275            format!("tensor<{scalar}>[{}]", dims.join(", "))
276        }
277    }
278}
279
280fn format_expr(handle: Handle<Expression>, exprs: &crate::Arena<Expression>) -> String {
281    match &exprs[handle] {
282        Expression::Literal(lit) => format!("{lit}"),
283        Expression::Compose { ty, components } => {
284            let args: Vec<_> = components.iter().map(|h| format!("{h:?}")).collect();
285            format!("Compose({ty:?}, [{}])", args.join(", "))
286        }
287        Expression::FunctionArgument(i) => format!("FunctionArgument({i})"),
288        Expression::GlobalVariable(h) => format!("GlobalVariable({h:?})"),
289        Expression::LocalVariable(h) => format!("LocalVariable({h:?})"),
290        Expression::Load { pointer } => format!("Load({pointer:?})"),
291        Expression::Access { base, index } => format!("Access({base:?}, {index:?})"),
292        Expression::AccessIndex { base, index } => format!("AccessIndex({base:?}, {index})"),
293        Expression::Swizzle {
294            size,
295            vector,
296            pattern,
297        } => {
298            let n = *size as usize;
299            let comps: Vec<_> = pattern[..n].iter().map(|c| format!("{c}")).collect();
300            format!("Swizzle({vector:?}).{}", comps.join(""))
301        }
302        Expression::Splat { size, value } => format!("Splat({value:?}, vec{size})"),
303        Expression::Unary { op, expr } => format!("{op}{expr:?}"),
304        Expression::Binary { op, left, right } => format!("{left:?} {op} {right:?}"),
305        Expression::Select {
306            condition,
307            accept,
308            reject,
309        } => format!("Select({condition:?}, {accept:?}, {reject:?})"),
310        Expression::Math {
311            fun,
312            arg,
313            arg1,
314            arg2,
315            arg3,
316        } => {
317            let mut args = format!("{arg:?}");
318            if let Some(a1) = arg1 {
319                args += &format!(", {a1:?}");
320            }
321            if let Some(a2) = arg2 {
322                args += &format!(", {a2:?}");
323            }
324            if let Some(a3) = arg3 {
325                args += &format!(", {a3:?}");
326            }
327            format!("{fun}({args})")
328        }
329        Expression::As {
330            expr,
331            kind,
332            convert,
333        } => match convert {
334            Some(w) => format!("As({expr:?} -> {kind}/{w})"),
335            None => format!("Bitcast({expr:?} -> {kind})"),
336        },
337        Expression::ArrayLength(expr) => format!("ArrayLength({expr:?})"),
338        Expression::CallResult(f) => format!("CallResult({f:?})"),
339        Expression::AtomicResult { ty, comparison } => {
340            format!("AtomicResult({ty:?}, cmp={comparison})")
341        }
342        Expression::ZeroValue(ty) => format!("ZeroValue({ty:?})"),
343    }
344}
345
346fn write_stmt(out: &mut String, stmt: &Statement, indent: usize) {
347    let pad = " ".repeat(indent);
348    match stmt {
349        Statement::Emit(range) => {
350            out.push_str(&format!("{pad}Emit({range:?})\n"));
351        }
352        Statement::Store { pointer, value } => {
353            out.push_str(&format!("{pad}Store {pointer:?} = {value:?}\n"));
354        }
355        Statement::If {
356            condition,
357            accept,
358            reject,
359        } => {
360            out.push_str(&format!("{pad}If ({condition:?}) {{\n"));
361            for s in accept {
362                write_stmt(out, s, indent + 4);
363            }
364            if !reject.is_empty() {
365                out.push_str(&format!("{pad}}} else {{\n"));
366                for s in reject {
367                    write_stmt(out, s, indent + 4);
368                }
369            }
370            out.push_str(&format!("{pad}}}\n"));
371        }
372        Statement::Loop {
373            body,
374            continuing,
375            break_if,
376        } => {
377            out.push_str(&format!("{pad}Loop {{\n"));
378            for s in body {
379                write_stmt(out, s, indent + 4);
380            }
381            if !continuing.is_empty() {
382                out.push_str(&format!("{pad}  Continuing {{\n"));
383                for s in continuing {
384                    write_stmt(out, s, indent + 8);
385                }
386                if let Some(brk) = break_if {
387                    out.push_str(&format!("{pad}    BreakIf({brk:?})\n"));
388                }
389                out.push_str(&format!("{pad}  }}\n"));
390            }
391            out.push_str(&format!("{pad}}}\n"));
392        }
393        Statement::Call {
394            function,
395            arguments,
396            result,
397        } => {
398            let args: Vec<_> = arguments.iter().map(|h| format!("{h:?}")).collect();
399            let res = match result {
400                Some(r) => format!(" -> {r:?}"),
401                None => String::new(),
402            };
403            out.push_str(&format!(
404                "{pad}Call {function:?}({}){res}\n",
405                args.join(", ")
406            ));
407        }
408        Statement::Atomic {
409            pointer,
410            fun,
411            value,
412            result,
413        } => {
414            let res = match result {
415                Some(r) => format!(" -> {r:?}"),
416                None => String::new(),
417            };
418            out.push_str(&format!("{pad}{fun}({pointer:?}, {value:?}){res}\n"));
419        }
420        Statement::Break => {
421            out.push_str(&format!("{pad}Break\n"));
422        }
423        Statement::Continue => {
424            out.push_str(&format!("{pad}Continue\n"));
425        }
426        Statement::Return { value } => match value {
427            Some(v) => out.push_str(&format!("{pad}Return {v:?}\n")),
428            None => out.push_str(&format!("{pad}Return\n")),
429        },
430        Statement::Barrier(b) => {
431            out.push_str(&format!("{pad}Barrier({b})\n"));
432        }
433    }
434}
435
436/// Produces a human-readable text dump of a [`Module`] for debugging.
437pub fn dump_module(module: &Module) -> String {
438    let mut out = String::new();
439
440    // Types
441    out.push_str("Types:\n");
442    for (handle, ty) in module.types.iter() {
443        let formatted = format_type(ty, &module.types);
444        out.push_str(&format!("  {handle:?} {formatted}\n"));
445    }
446
447    // Global variables
448    if !module.global_variables.is_empty() {
449        out.push_str("\nGlobal Variables:\n");
450        for (handle, var) in module.global_variables.iter() {
451            let name = var.name.as_deref().unwrap_or("_");
452            let ty_str = format_type(&module.types[var.ty], &module.types);
453            let binding_str = match &var.binding {
454                Some(b) => format!("{b} "),
455                None => String::new(),
456            };
457            out.push_str(&format!(
458                "  {handle:?} {binding_str}var<{}>  {name}: {ty_str}\n",
459                var.space
460            ));
461        }
462    }
463
464    // Global expressions
465    if !module.global_expressions.is_empty() {
466        out.push_str("\nGlobal Expressions:\n");
467        for (handle, _) in module.global_expressions.iter() {
468            let formatted = format_expr(handle, &module.global_expressions);
469            out.push_str(&format!("  {handle:?} {formatted}\n"));
470        }
471    }
472
473    // Helper functions
474    if !module.functions.is_empty() {
475        out.push_str("\nFunctions:\n");
476        for (handle, func) in module.functions.iter() {
477            dump_function(&mut out, &format!("{handle:?}"), func, &module.types);
478        }
479    }
480
481    // Entry points
482    if !module.entry_points.is_empty() {
483        out.push_str("\nEntry Points:\n");
484        for ep in &module.entry_points {
485            let [x, y, z] = ep.workgroup_size;
486            out.push_str(&format!("  @compute @workgroup_size({x}, {y}, {z})\n"));
487            dump_function(&mut out, &ep.name, &ep.function, &module.types);
488        }
489    }
490
491    out
492}
493
494fn dump_function(out: &mut String, label: &str, func: &crate::Function, types: &UniqueArena<Type>) {
495    let name = func.name.as_deref().unwrap_or("_");
496
497    // Signature
498    let args: Vec<_> = func
499        .arguments
500        .iter()
501        .map(|arg| {
502            let arg_name = arg.name.as_deref().unwrap_or("_");
503            let ty_str = format_type(&types[arg.ty], types);
504            let binding = match &arg.binding {
505                Some(b) => format!("{b} "),
506                None => String::new(),
507            };
508            format!("{binding}{arg_name}: {ty_str}")
509        })
510        .collect();
511    let ret = match &func.result {
512        Some(r) => format!(" -> {}", format_type(&types[r.ty], types)),
513        None => String::new(),
514    };
515    out.push_str(&format!(
516        "  fn {name}({})  [{label}]{ret} {{\n",
517        args.join(", ")
518    ));
519
520    // Local variables
521    for (handle, var) in func.local_variables.iter() {
522        let var_name = var.name.as_deref().unwrap_or("_");
523        let ty_str = format_type(&types[var.ty], types);
524        let init = match var.init {
525            Some(h) => format!(" = {}", format_expr(h, &func.expressions)),
526            None => String::new(),
527        };
528        out.push_str(&format!("    var {handle:?} {var_name}: {ty_str}{init}\n"));
529    }
530
531    // Expressions
532    if !func.expressions.is_empty() {
533        out.push_str("    Expressions:\n");
534        for (handle, _) in func.expressions.iter() {
535            let formatted = format_expr(handle, &func.expressions);
536            out.push_str(&format!("      {handle:?} {formatted}\n"));
537        }
538    }
539
540    // Body
541    if !func.body.is_empty() {
542        out.push_str("    Body:\n");
543        for stmt in &func.body {
544            write_stmt(out, stmt, 6);
545        }
546    }
547
548    out.push_str("  }\n");
549}
550
551#[cfg(test)]
552mod tests {
553    use super::*;
554
555    #[test]
556    fn display_scalar() {
557        assert_eq!(format!("{}", Scalar::F32), "f32");
558        assert_eq!(format!("{}", Scalar::I32), "i32");
559        assert_eq!(format!("{}", Scalar::U32), "u32");
560        assert_eq!(format!("{}", Scalar::F16), "f16");
561        assert_eq!(format!("{}", Scalar::BOOL), "bool");
562    }
563
564    #[test]
565    fn display_address_space() {
566        assert_eq!(format!("{}", AddressSpace::Uniform), "uniform");
567        assert_eq!(format!("{}", AddressSpace::Workgroup), "workgroup");
568        assert_eq!(
569            format!(
570                "{}",
571                AddressSpace::Storage {
572                    access: StorageAccess::LOAD | StorageAccess::STORE
573                }
574            ),
575            "storage, read_write"
576        );
577    }
578
579    #[test]
580    fn display_literal() {
581        assert_eq!(format!("{}", Literal::F32(3.125)), "3.125f");
582        assert_eq!(format!("{}", Literal::U32(42)), "42u");
583        assert_eq!(format!("{}", Literal::Bool(true)), "true");
584    }
585
586    #[test]
587    fn display_binary_op() {
588        assert_eq!(format!("{}", BinaryOp::Add), "+");
589        assert_eq!(format!("{}", BinaryOp::Equal), "==");
590        assert_eq!(format!("{}", BinaryOp::ShiftLeft), "<<");
591    }
592
593    #[test]
594    fn display_math_function() {
595        assert_eq!(format!("{}", MathFunction::Dot), "dot");
596        assert_eq!(format!("{}", MathFunction::Normalize), "normalize");
597    }
598
599    #[test]
600    fn display_binding() {
601        let b = Binding::BuiltIn(BuiltIn::GlobalInvocationId);
602        assert_eq!(format!("{b}"), "@builtin(global_invocation_id)");
603    }
604
605    #[test]
606    fn display_resource_binding() {
607        let rb = ResourceBinding {
608            group: 0,
609            binding: 2,
610        };
611        assert_eq!(format!("{rb}"), "@group(0) @binding(2)");
612    }
613
614    #[test]
615    fn dump_empty_module() {
616        let module = Module::default();
617        let dump = dump_module(&module);
618        assert!(dump.contains("Types:"));
619    }
620
621    #[test]
622    fn display_scalar_kind_all_variants() {
623        assert_eq!(format!("{}", ScalarKind::Bool), "bool");
624        assert_eq!(format!("{}", ScalarKind::Sint), "sint");
625        assert_eq!(format!("{}", ScalarKind::Uint), "uint");
626        assert_eq!(format!("{}", ScalarKind::Float), "float");
627        assert_eq!(format!("{}", ScalarKind::BFloat), "bfloat");
628    }
629
630    #[test]
631    fn display_scalar_bfloat() {
632        assert_eq!(format!("{}", Scalar::BF16), "bf16");
633    }
634
635    #[test]
636    fn display_storage_access_all_variants() {
637        assert_eq!(format!("{}", StorageAccess::LOAD), "read");
638        assert_eq!(format!("{}", StorageAccess::STORE), "write");
639        assert_eq!(
640            format!("{}", StorageAccess::LOAD | StorageAccess::STORE),
641            "read_write"
642        );
643        assert_eq!(format!("{}", StorageAccess::EMPTY), "none");
644    }
645
646    #[test]
647    fn display_address_space_all_variants() {
648        assert_eq!(format!("{}", AddressSpace::Function), "function");
649        assert_eq!(format!("{}", AddressSpace::Private), "private");
650        assert_eq!(format!("{}", AddressSpace::Workgroup), "workgroup");
651        assert_eq!(format!("{}", AddressSpace::Uniform), "uniform");
652        assert_eq!(
653            format!(
654                "{}",
655                AddressSpace::Storage {
656                    access: StorageAccess::LOAD
657                }
658            ),
659            "storage, read"
660        );
661    }
662
663    #[test]
664    fn display_builtin_all_variants() {
665        assert_eq!(
666            format!("{}", BuiltIn::GlobalInvocationId),
667            "global_invocation_id"
668        );
669        assert_eq!(
670            format!("{}", BuiltIn::LocalInvocationId),
671            "local_invocation_id"
672        );
673        assert_eq!(
674            format!("{}", BuiltIn::LocalInvocationIndex),
675            "local_invocation_index"
676        );
677        assert_eq!(format!("{}", BuiltIn::WorkgroupId), "workgroup_id");
678        assert_eq!(format!("{}", BuiltIn::NumWorkgroups), "num_workgroups");
679    }
680
681    #[test]
682    fn display_binding_location() {
683        let b = Binding::Location { location: 3 };
684        assert_eq!(format!("{b}"), "@location(3)");
685    }
686
687    #[test]
688    fn display_literal_all_variants() {
689        assert_eq!(format!("{}", Literal::Bool(false)), "false");
690        assert_eq!(format!("{}", Literal::I32(-7)), "-7i");
691        assert_eq!(format!("{}", Literal::U32(42)), "42u");
692        assert_eq!(format!("{}", Literal::F32(1.5)), "1.5f");
693        assert_eq!(format!("{}", Literal::F64(2.5)), "2.5lf");
694        assert_eq!(format!("{}", Literal::AbstractInt(99)), "99");
695        assert_eq!(format!("{}", Literal::AbstractFloat(1.23)), "1.23");
696    }
697
698    #[test]
699    fn display_unary_op_all_variants() {
700        assert_eq!(format!("{}", UnaryOp::Negate), "-");
701        assert_eq!(format!("{}", UnaryOp::LogicalNot), "!");
702        assert_eq!(format!("{}", UnaryOp::BitwiseNot), "~");
703    }
704
705    #[test]
706    fn display_binary_op_all_variants() {
707        assert_eq!(format!("{}", BinaryOp::Add), "+");
708        assert_eq!(format!("{}", BinaryOp::Subtract), "-");
709        assert_eq!(format!("{}", BinaryOp::Multiply), "*");
710        assert_eq!(format!("{}", BinaryOp::Divide), "/");
711        assert_eq!(format!("{}", BinaryOp::Modulo), "%");
712        assert_eq!(format!("{}", BinaryOp::Equal), "==");
713        assert_eq!(format!("{}", BinaryOp::NotEqual), "!=");
714        assert_eq!(format!("{}", BinaryOp::Less), "<");
715        assert_eq!(format!("{}", BinaryOp::LessEqual), "<=");
716        assert_eq!(format!("{}", BinaryOp::Greater), ">");
717        assert_eq!(format!("{}", BinaryOp::GreaterEqual), ">=");
718        assert_eq!(format!("{}", BinaryOp::LogicalAnd), "&&");
719        assert_eq!(format!("{}", BinaryOp::LogicalOr), "||");
720        assert_eq!(format!("{}", BinaryOp::BitwiseAnd), "&");
721        assert_eq!(format!("{}", BinaryOp::BitwiseOr), "|");
722        assert_eq!(format!("{}", BinaryOp::BitwiseXor), "^");
723        assert_eq!(format!("{}", BinaryOp::ShiftLeft), "<<");
724        assert_eq!(format!("{}", BinaryOp::ShiftRight), ">>");
725    }
726
727    #[test]
728    fn display_swizzle_component_all_variants() {
729        assert_eq!(format!("{}", SwizzleComponent::X), "x");
730        assert_eq!(format!("{}", SwizzleComponent::Y), "y");
731        assert_eq!(format!("{}", SwizzleComponent::Z), "z");
732        assert_eq!(format!("{}", SwizzleComponent::W), "w");
733    }
734
735    #[test]
736    fn display_atomic_function_all_variants() {
737        assert_eq!(format!("{}", AtomicFunction::Add), "atomicAdd");
738        assert_eq!(format!("{}", AtomicFunction::Subtract), "atomicSub");
739        assert_eq!(format!("{}", AtomicFunction::And), "atomicAnd");
740        assert_eq!(format!("{}", AtomicFunction::ExclusiveOr), "atomicXor");
741        assert_eq!(format!("{}", AtomicFunction::InclusiveOr), "atomicOr");
742        assert_eq!(format!("{}", AtomicFunction::Min), "atomicMin");
743        assert_eq!(format!("{}", AtomicFunction::Max), "atomicMax");
744        assert_eq!(
745            format!("{}", AtomicFunction::Exchange { compare: None }),
746            "atomicExchange"
747        );
748        let cmp_handle = {
749            let mut arena = crate::Arena::new();
750            arena.append(Expression::Literal(Literal::U32(0)))
751        };
752        assert!(
753            format!(
754                "{}",
755                AtomicFunction::Exchange {
756                    compare: Some(cmp_handle)
757                }
758            )
759            .starts_with("atomicCompareExchange(")
760        );
761    }
762
763    #[test]
764    fn display_barrier_all_variants() {
765        assert_eq!(format!("{}", Barrier::STORAGE), "storageBarrier");
766        assert_eq!(format!("{}", Barrier::WORKGROUP), "workgroupBarrier");
767        assert_eq!(
768            format!("{}", Barrier::STORAGE | Barrier::WORKGROUP),
769            "storageBarrier | workgroupBarrier"
770        );
771        assert_eq!(format!("{}", Barrier::EMPTY), "<no barrier>");
772    }
773
774    #[test]
775    fn format_type_named() {
776        let mut types = UniqueArena::new();
777        let h = types.insert(Type {
778            name: Some("MyStruct".into()),
779            inner: TypeInner::Scalar(Scalar::F32),
780        });
781        assert_eq!(format_type(&types[h], &types), "MyStruct");
782    }
783
784    #[test]
785    fn format_type_inner_all_variants() {
786        let mut types = UniqueArena::new();
787        let f32_ty = types.insert(Type {
788            name: None,
789            inner: TypeInner::Scalar(Scalar::F32),
790        });
791
792        // Scalar
793        assert_eq!(
794            format_type_inner(&TypeInner::Scalar(Scalar::F32), &types),
795            "f32"
796        );
797
798        // Vector
799        assert_eq!(
800            format_type_inner(
801                &TypeInner::Vector {
802                    size: VectorSize::Tri,
803                    scalar: Scalar::F32
804                },
805                &types
806            ),
807            "vec3<f32>"
808        );
809
810        // Matrix
811        assert_eq!(
812            format_type_inner(
813                &TypeInner::Matrix {
814                    columns: VectorSize::Quad,
815                    rows: VectorSize::Quad,
816                    scalar: Scalar::F32
817                },
818                &types
819            ),
820            "mat4x4<f32>"
821        );
822
823        // Atomic
824        assert_eq!(
825            format_type_inner(&TypeInner::Atomic(Scalar::U32), &types),
826            "atomic<u32>"
827        );
828
829        // Pointer
830        assert_eq!(
831            format_type_inner(
832                &TypeInner::Pointer {
833                    base: f32_ty,
834                    space: AddressSpace::Function
835                },
836                &types
837            ),
838            "ptr<function, f32>"
839        );
840
841        // Array (constant)
842        assert_eq!(
843            format_type_inner(
844                &TypeInner::Array {
845                    base: f32_ty,
846                    size: ArraySize::Constant(16),
847                    stride: 4
848                },
849                &types
850            ),
851            "array<f32, 16> /*stride 4*/"
852        );
853
854        // Array (dynamic)
855        assert_eq!(
856            format_type_inner(
857                &TypeInner::Array {
858                    base: f32_ty,
859                    size: ArraySize::Dynamic,
860                    stride: 4
861                },
862                &types
863            ),
864            "array<f32> /*stride 4*/"
865        );
866
867        // Struct
868        assert_eq!(
869            format_type_inner(
870                &TypeInner::Struct {
871                    members: vec![],
872                    span: 16
873                },
874                &types
875            ),
876            "struct(0 members, span 16)"
877        );
878
879        // Tensor (mixed dims)
880        assert_eq!(
881            format_type_inner(
882                &TypeInner::Tensor {
883                    scalar: Scalar::F32,
884                    shape: crate::TensorShape {
885                        dims: vec![
886                            crate::Dimension::Fixed(224),
887                            crate::Dimension::Dynamic(Some("batch".into())),
888                            crate::Dimension::Dynamic(None),
889                        ],
890                    }
891                },
892                &types
893            ),
894            "tensor<f32>[224, batch, ?]"
895        );
896
897        // Tensor with Symbolic dimension
898        assert_eq!(
899            format_type_inner(
900                &TypeInner::Tensor {
901                    scalar: Scalar::F32,
902                    shape: crate::TensorShape {
903                        dims: vec![
904                            crate::Dimension::Symbolic("batch".into()),
905                            crate::Dimension::Fixed(224),
906                            crate::Dimension::Fixed(224),
907                            crate::Dimension::Fixed(3),
908                        ],
909                    }
910                },
911                &types
912            ),
913            "tensor<f32>[batch, 224, 224, 3]"
914        );
915    }
916
917    #[test]
918    fn dump_module_with_globals_and_entry_point() {
919        use crate::{EntryPoint, Expression, Function, GlobalVariable, ResourceBinding, Statement};
920
921        let mut module = Module::default();
922
923        let f32_ty = module.types.insert(Type {
924            name: None,
925            inner: TypeInner::Scalar(Scalar::F32),
926        });
927        let arr_ty = module.types.insert(Type {
928            name: None,
929            inner: TypeInner::Array {
930                base: f32_ty,
931                size: ArraySize::Dynamic,
932                stride: 4,
933            },
934        });
935
936        module.global_variables.append(GlobalVariable {
937            name: Some("a".into()),
938            space: AddressSpace::Storage {
939                access: StorageAccess::LOAD,
940            },
941            binding: Some(ResourceBinding {
942                group: 0,
943                binding: 0,
944            }),
945            ty: arr_ty,
946            init: None,
947            layout: None,
948        });
949
950        module.global_variables.append(GlobalVariable {
951            name: None,
952            space: AddressSpace::Private,
953            binding: None,
954            ty: f32_ty,
955            init: None,
956            layout: None,
957        });
958
959        let mut func = Function::new("main");
960        let lit = func
961            .expressions
962            .append(Expression::Literal(Literal::F32(1.0)));
963        let gv = func.expressions.append(Expression::GlobalVariable(
964            module.global_variables.next_handle(),
965        ));
966        func.body.push(Statement::Store {
967            pointer: gv,
968            value: lit,
969        });
970
971        module.entry_points.push(EntryPoint {
972            name: "main".into(),
973            workgroup_size: [256, 1, 1],
974            function: func,
975        });
976
977        let dump = dump_module(&module);
978        assert!(dump.contains("Global Variables:"));
979        assert!(dump.contains("@group(0) @binding(0)"));
980        assert!(dump.contains("a:"));
981        assert!(dump.contains("storage, read"));
982        assert!(dump.contains("Entry Points:"));
983        assert!(dump.contains("@compute @workgroup_size(256, 1, 1)"));
984        assert!(dump.contains("Expressions:"));
985        assert!(dump.contains("Body:"));
986        assert!(dump.contains("Store"));
987    }
988
989    #[test]
990    fn dump_module_with_if_and_loop() {
991        use crate::{EntryPoint, Expression, Function, Statement};
992
993        let mut module = Module::default();
994        let mut func = Function::new("main");
995
996        let cond = func
997            .expressions
998            .append(Expression::Literal(Literal::Bool(true)));
999        let val = func
1000            .expressions
1001            .append(Expression::Literal(Literal::F32(1.0)));
1002
1003        func.body.push(Statement::If {
1004            condition: cond,
1005            accept: vec![Statement::Return { value: Some(val) }],
1006            reject: vec![Statement::Return { value: None }],
1007        });
1008        func.body.push(Statement::Loop {
1009            body: vec![Statement::Break],
1010            continuing: vec![Statement::Continue],
1011            break_if: Some(cond),
1012        });
1013        func.body.push(Statement::Barrier(Barrier::STORAGE));
1014
1015        module.entry_points.push(EntryPoint {
1016            name: "main".into(),
1017            workgroup_size: [1, 1, 1],
1018            function: func,
1019        });
1020
1021        let dump = dump_module(&module);
1022        assert!(dump.contains("If ("));
1023        assert!(dump.contains("} else {"));
1024        assert!(dump.contains("Return"));
1025        assert!(dump.contains("Loop {"));
1026        assert!(dump.contains("Continuing {"));
1027        assert!(dump.contains("BreakIf("));
1028        assert!(dump.contains("Break"));
1029        assert!(dump.contains("Continue"));
1030        assert!(dump.contains("Barrier(storageBarrier)"));
1031    }
1032
1033    #[test]
1034    fn dump_module_with_helper_function() {
1035        use crate::{Function, FunctionArgument, FunctionResult};
1036
1037        let mut module = Module::default();
1038        let f32_ty = module.types.insert(Type {
1039            name: None,
1040            inner: TypeInner::Scalar(Scalar::F32),
1041        });
1042
1043        let mut func = Function::new("helper");
1044        func.arguments.push(FunctionArgument {
1045            name: Some("x".into()),
1046            ty: f32_ty,
1047            binding: None,
1048        });
1049        func.result = Some(FunctionResult {
1050            ty: f32_ty,
1051            binding: None,
1052        });
1053        module.functions.append(func);
1054
1055        let dump = dump_module(&module);
1056        assert!(dump.contains("Functions:"));
1057        assert!(dump.contains("fn helper(x: f32)"));
1058        assert!(dump.contains("-> f32"));
1059    }
1060
1061    #[test]
1062    fn display_vector_size() {
1063        assert_eq!(format!("{}", VectorSize::Bi), "2");
1064        assert_eq!(format!("{}", VectorSize::Tri), "3");
1065        assert_eq!(format!("{}", VectorSize::Quad), "4");
1066    }
1067
1068    #[test]
1069    fn display_math_function_all_variants() {
1070        // Spot-check a representative sample beyond what existing tests cover
1071        assert_eq!(format!("{}", MathFunction::Abs), "abs");
1072        assert_eq!(format!("{}", MathFunction::Clamp), "clamp");
1073        assert_eq!(format!("{}", MathFunction::Fma), "fma");
1074        assert_eq!(format!("{}", MathFunction::Sin), "sin");
1075        assert_eq!(format!("{}", MathFunction::Tanh), "tanh");
1076        assert_eq!(format!("{}", MathFunction::Sqrt), "sqrt");
1077        assert_eq!(format!("{}", MathFunction::Pow), "pow");
1078        assert_eq!(format!("{}", MathFunction::Mix), "mix");
1079        assert_eq!(format!("{}", MathFunction::SmoothStep), "smoothStep");
1080        assert_eq!(format!("{}", MathFunction::InverseSqrt), "inverseSqrt");
1081    }
1082}
1083
1084#[cfg(test)]
1085mod math_function_display_tests {
1086    use super::*;
1087
1088    #[test]
1089    fn bit_manipulation_prints_as_wgsl_spells_it() {
1090        // These names go into IR dumps that people read next to the shader
1091        // they came from, so they match the WGSL builtin rather than the
1092        // Rust variant.
1093        assert_eq!(MathFunction::ExtractBits.to_string(), "extractBits");
1094        assert_eq!(MathFunction::InsertBits.to_string(), "insertBits");
1095    }
1096}
nxpu_ir/display.rs

nxpu_ir/
display.rs
