Add debug_printf! and debug_printfln! macros that uses the DebugPrintf extension

crates/spirv-std/macros/Cargo.toml

@@ -16,3 +16,4 @@ heck = "0.3.2"
 proc-macro2 = "1.0.24"
 quote = "1.0.8"
 syn = { version = "1.0.58", features=["full"] }
+regex = "1.5.4"

crates/spirv-std/macros/src/lib.rs

@@ -350,3 +350,167 @@ fn path_from_ident(ident: Ident) -> syn::Type {
         path: syn::Path::from(ident),
     })
 }
+
+/// Print a formatted string with a newline using the debug printf extension.
+///
+/// Examples:
+///
+/// ```rust,ignore
+/// debug_printfln!("uv: %v2f", uv);
+/// debug_printfln!("pos.x: %f, pos.z: %f, int: %i", pos.x, pos.z, int);
+/// ```
+///
+/// See <https://github.com/KhronosGroup/Vulkan-ValidationLayers/blob/master/docs/debug_printf.md#debug-printf-format-string> for formatting rules.
+#[proc_macro]
+pub fn debug_printf(input: TokenStream) -> TokenStream {
+    debug_printf_inner(syn::parse_macro_input!(input as DebugPrintfInput))
+}
+
+/// Similar to `debug_printf` but appends a newline to the format string.
+#[proc_macro]
+pub fn debug_printfln(input: TokenStream) -> TokenStream {
+    let mut input = syn::parse_macro_input!(input as DebugPrintfInput);
+    input.format_string.push('\n');
+    debug_printf_inner(input)
+}
+
+struct DebugPrintfInput {
+    span: proc_macro2::Span,
+    format_string: String,
+    variables: Vec<syn::Expr>,
+}
+
+impl syn::parse::Parse for DebugPrintfInput {
+    fn parse(input: syn::parse::ParseStream<'_>) -> syn::parse::Result<Self> {
+        let span = input.span();
+
+        if input.is_empty() {
+            return Ok(Self {
+                span,
+                format_string: Default::default(),
+                variables: Default::default(),
+            });
+        }
+
+        let format_string = input.parse::<syn::LitStr>()?;
+        if !input.is_empty() {
+            input.parse::<syn::token::Comma>()?;
+        }
+        let variables =
+            syn::punctuated::Punctuated::<syn::Expr, syn::token::Comma>::parse_terminated(input)?;
+
+        Ok(Self {
+            span,
+            format_string: format_string.value(),
+            variables: variables.into_iter().collect(),
+        })
+    }
+}
+
+fn debug_printf_inner(input: DebugPrintfInput) -> TokenStream {
+    let DebugPrintfInput {
+        format_string,
+        variables,
+        span,
+    } = input;
+
+    let number_of_arguments =
+        format_string.matches('%').count() - format_string.matches("%%").count() * 2;
+
+    if number_of_arguments != variables.len() {
+        return syn::Error::new(
+            span,
+            &format!(
+                "{} % arguments were found, but {} variables were given",
+                number_of_arguments,
+                variables.len()
+            ),
+        )
+        .to_compile_error()
+        .into();
+    }
+
+    let specifiers = "d|i|o|u|x|X|a|A|e|E|f|F|g|G|ul|lu|lx";
+
+    let regex = regex::Regex::new(&format!(
+        r"%\d*\.?\d*(v(2|3|4)({specifiers})|{specifiers})",
+        specifiers = specifiers
+    ))
+    .unwrap();
+
+    fn map_specifier_to_type(specifier: &str) -> proc_macro2::TokenStream {
+        match specifier {
+            "d" => quote::quote! { i32 },
+            "i" => quote::quote! { i32 },
+            "o" => quote::quote! { u32 },
+            "u" => quote::quote! { u32 },
+            "x" => quote::quote! { u32 },
+            "X" => quote::quote! { u32 },
+            "a" => quote::quote! { f32 },
+            "A" => quote::quote! { f32 },
+            "e" => quote::quote! { f32 },
+            "E" => quote::quote! { f32 },
+            "f" => quote::quote! { f32 },
+            "F" => quote::quote! { f32 },
+            "g" => quote::quote! { f32 },
+            "G" => quote::quote! { f32 },
+            "ul" => quote::quote! { u64 },
+            "lu" => quote::quote! { u64 },
+            "lx" => quote::quote! { u64 },
+            _ => unreachable!(),
+        }
+    }
+
+    let assert_fns = regex.captures_iter(&format_string).map(|captures| {
+        let specifier = &captures[1][0..1];
+
+        if specifier == "v" {
+            let count = &captures[2].parse::<usize>().unwrap();
+            let ty = map_specifier_to_type(&captures[3][0..1]);
+            quote::quote! { spirv_std::debug_printf_assert_is_vector::<#ty, _, #count> }
+        } else {
+            let ty = map_specifier_to_type(specifier);
+            quote::quote! { spirv_std::debug_printf_assert_is_type::<#ty> }
+        }
+    });
+
+    let mut variable_idents = String::new();
+    let mut input_registers = Vec::new();
+    let mut op_loads = Vec::new();
+
+    for (i, (variable, assert_fn)) in variables.into_iter().zip(assert_fns).enumerate() {
+        let ident = quote::format_ident!("_{}", i);
+
+        variable_idents.push_str(&format!("%{} ", ident));
+
+        input_registers.push(quote::quote! {
+            #ident = in(reg) &#assert_fn(#variable),
+        });
+
+        let op_load = format!("%{ident} = OpLoad _ {{{ident}}}", ident = ident);
+
+        op_loads.push(quote::quote! {
+            #op_load,
+        });
+    }
+
+    let input_registers = input_registers
+        .into_iter()
+        .collect::<proc_macro2::TokenStream>();
+    let op_loads = op_loads.into_iter().collect::<proc_macro2::TokenStream>();
+
+    let op_string = format!("%string = OpString {:?}", format_string);
+
+    let output = quote::quote! {
+        asm!(
+            "%void = OpTypeVoid",
+            #op_string,
+            "%debug_printf = OpExtInstImport \"NonSemantic.DebugPrintf\"",
+            #op_loads
+            concat!("%result = OpExtInst %void %debug_printf 1 %string ", #variable_idents),
+            #input_registers
+        )
+    };
+
+    output.into()
+}

crates/spirv-std/src/lib.rs

@@ -130,3 +130,19 @@ extern "C" fn rust_eh_personality() {}
 #[doc(hidden)]
 /// [spirv_types]
 pub fn workaround_rustdoc_ice_84738() {}
+
+#[doc(hidden)]
+pub fn debug_printf_assert_is_type<T>(ty: T) -> T {
+    ty
+}
+
+#[doc(hidden)]
+pub fn debug_printf_assert_is_vector<
+    TY: crate::scalar::Scalar,
+    V: crate::vector::Vector<TY, SIZE>,
+    const SIZE: usize,
+>(
+    vec: V,
+) -> V {
+    vec
+}

tests/ui/arch/debug_printf.rs

@@ -0,0 +1,51 @@
+// build-pass
+// compile-flags: -Ctarget-feature=+ext:SPV_KHR_non_semantic_info
+
+use spirv_std::{
+    glam::{IVec2, UVec2, Vec2, Vec3, Vec4},
+    macros::{debug_printf, debug_printfln},
+};
+
+fn func(a: f32, b: f32) -> f32 {
+    a * b + 1.0
+}
+
+struct Struct {
+    a: f32,
+}
+
+impl Struct {
+    fn method(&self, b: f32, c: f32) -> f32 {
+        self.a * b + c
+    }
+}
+
+#[spirv(fragment)]
+pub fn main() {
+    unsafe {
+        debug_printf!();
+        debug_printfln!();
+        debug_printfln!("Hello World");
+        debug_printfln!("Hello World",);
+        debug_printfln!(r#"Hello "World""#);
+        debug_printfln!(
+            r#"Hello "World"
+"#
+        );
+        debug_printfln!("Hello \"World\"\n\n");
+    }
+
+    let vec = Vec2::new(1.52, 25.1);
+
+    unsafe {
+        debug_printfln!("%v2f", vec);
+        debug_printfln!("%1v2f", { vec * 2.0 });
+        debug_printfln!("%1.2v2f", vec * 3.0);
+        debug_printfln!("%% %v2f %%", vec * 4.0);
+        debug_printfln!("%u %i %f 🐉", 11_u32, -11_i32, 11.0_f32);
+        debug_printfln!("%f", func(33.0, 44.0));
+        debug_printfln!("%f", Struct { a: 33.0 }.method(44.0, 55.0));
+        debug_printfln!("%v3f %v4f", Vec3::new(1.0, 1.0, 1.0), Vec4::splat(5.0));
+        debug_printfln!("%v2u %v2i", UVec2::new(1, 1), IVec2::splat(-5));
+    }
+}