Proc macro for running picasso and include!ing the compiled shader

.gitignore

@@ -17,3 +17,6 @@ Cargo.lock
 rust-toolchain
 rust-toolchain.toml
 .cargo/
+
+# Pica200 output files
+*.shbin

Cargo.toml

@@ -1,6 +1,15 @@
 [workspace]
-members = ["citro3d-sys", "citro3d", "bindgen-citro3d"]
-default-members = ["citro3d", "citro3d-sys"]
+members = [
+    "bindgen-citro3d",
+    "citro3d-macros",
+    "citro3d-sys",
+    "citro3d",
+]
+default-members = [
+    "citro3d",
+    "citro3d-sys",
+    "citro3d-macros",
+]
 resolver = "2"
 
 [patch."https://github.com/rust3ds/citro3d-rs.git"]

citro3d-macros/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "citro3d-macros"
+version = "0.1.0"
+edition = "2021"
+authors = ["Rust3DS Org"]
+license = "MIT OR Apache-2.0"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+litrs = { version = "0.4.0", default-features = false }
+quote = "1.0.32"

citro3d-macros/build.rs

@@ -0,0 +1,8 @@
+//! This build script mainly exists just to ensure `OUT_DIR` is set for the macro,
+//! but we can also use it to force a re-evaluation if `DEVKITPRO` changes.
+
+fn main() {
+    for var in ["OUT_DIR", "DEVKITPRO"] {
+        println!("cargo:rerun-if-env-changed={var}");
+    }
+}

citro3d-macros/src/lib.rs

@@ -0,0 +1,167 @@
+// we're already nightly-only so might as well use unstable proc macro APIs.
+#![feature(proc_macro_span)]
+
+use std::error::Error;
+use std::fs::DirBuilder;
+use std::path::PathBuf;
+use std::{env, process};
+
+use litrs::StringLit;
+use proc_macro::TokenStream;
+use quote::quote;
+
+/// Compiles the given PICA200 shader using [`picasso`](https://github.com/devkitPro/picasso)
+/// and returns the compiled bytes directly as a `&[u8]` slice.
+///
+/// This is similar to the standard library's [`include_bytes!`](std::include_bytes) macro, for which
+/// file paths are relative to the source file where the macro is invoked.
+///
+/// The compiled shader binary will be saved in the caller's `$OUT_DIR`.
+///
+/// # Errors
+///
+/// This macro will fail to compile if the input is not a single string literal.
+/// In other words, inputs like `concat!("foo", "/bar")` are not supported.
+///
+/// # Example
+///
+/// ```
+/// use citro3d_macros::include_shader;
+///
+/// static SHADER_BYTES: &[u8] = include_shader!("../tests/integration.pica");
+/// ```
+///
+/// # Errors
+///
+/// The macro will fail to compile if the `.pica` file cannot be found, or contains
+/// `picasso` syntax errors.
+///
+/// ```compile_fail
+/// # use citro3d_macros::include_shader;
+/// static _ERROR: &[u8] = include_shader!("../tests/nonexistent.pica");
+/// ```
+///
+/// ```compile_fail
+/// # use citro3d_macros::include_shader;
+/// static _ERROR: &[u8] = include_shader!("../tests/bad-shader.pica");
+/// ```
+#[proc_macro]
+pub fn include_shader(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    match include_shader_impl(input) {
+        Ok(tokens) => tokens,
+        Err(err) => {
+            let err_str = err.to_string();
+            quote! { compile_error!( #err_str ) }.into()
+        }
+    }
+}
+
+fn include_shader_impl(input: TokenStream) -> Result<TokenStream, Box<dyn Error>> {
+    let tokens: Vec<_> = input.into_iter().collect();
+
+    if tokens.len() != 1 {
+        return Err(format!("expected exactly one input token, got {}", tokens.len()).into());
+    }
+
+    let shader_source_filename = &tokens[0];
+
+    let string_lit = match StringLit::try_from(shader_source_filename) {
+        Ok(lit) => lit,
+        Err(err) => return Ok(err.to_compile_error()),
+    };
+
+    // The cwd can change depending on whether this is running in a doctest or not:
+    // https://users.rust-lang.org/t/which-directory-does-a-proc-macro-run-from/71917
+    //
+    // But the span's `source_file()` seems to always be relative to the cwd.
+    let cwd = env::current_dir()
+        .map_err(|err| format!("unable to determine current directory: {err}"))?;
+
+    let invoking_source_file = shader_source_filename.span().source_file().path();
+    let Some(invoking_source_dir) = invoking_source_file.parent() else {
+        return Ok(quote! {
+            compile_error!(
+                concat!(
+                    "unable to find parent directory of current source file \"",
+                    file!(),
+                    "\""
+                )
+            )
+        }
+        .into());
+    };
+
+    // By joining these three pieces, we arrive at approximately the same behavior as `include_bytes!`
+    let shader_source_file = cwd
+        .join(invoking_source_dir)
+        .join(string_lit.value())
+        // This might be overkill, but it ensures we get a unique path if different
+        // shaders with the same relative path are used within one program
+        .canonicalize()
+        .map_err(|err| format!("unable to resolve absolute path of shader source: {err}"))?;
+
+    let shader_out_file: PathBuf = shader_source_file.with_extension("shbin");
+
+    let out_dir = PathBuf::from(env!("OUT_DIR"));
+
+    let out_path = out_dir.join(shader_out_file.components().skip(1).collect::<PathBuf>());
+    // UNWRAP: we already canonicalized the source path, so it should have a parent.
+    let out_parent = out_path.parent().unwrap();
+
+    DirBuilder::new()
+        .recursive(true)
+        .create(out_parent)
+        .map_err(|err| format!("unable to create output directory {out_parent:?}: {err}"))?;
+
+    let devkitpro = PathBuf::from(env!("DEVKITPRO"));
+    let picasso = devkitpro.join("tools/bin/picasso");
+
+    let output = process::Command::new(&picasso)
+        .arg("--out")
+        .args([&out_path, &shader_source_file])
+        .output()
+        .map_err(|err| format!("unable to run {picasso:?}: {err}"))?;
+
+    let error_code = match output.status.code() {
+        Some(0) => None,
+        code => Some(code.map_or_else(|| String::from("<unknown>"), |c| c.to_string())),
+    };
+
+    if let Some(code) = error_code {
+        return Err(format!(
+            "failed to compile shader: `picasso` exited with status {code}: {}",
+            String::from_utf8_lossy(&output.stderr),
+        )
+        .into());
+    }
+
+    let bytes = std::fs::read(&out_path)
+        .map_err(|err| format!("unable to read output file {out_path:?}: {err}"))?;
+
+    let source_file_path = shader_source_file.to_string_lossy();
+
+    let result = quote! {
+        {
+            // ensure the source is re-evaluted if the input file changes
+            const _SOURCE: &[u8] = include_bytes! ( #source_file_path );
+
+            // https://users.rust-lang.org/t/can-i-conveniently-compile-bytes-into-a-rust-program-with-a-specific-alignment/24049/2
+            #[repr(C)]
+            struct AlignedAsU32<Bytes: ?Sized> {
+                _align: [u32; 0],
+                bytes: Bytes,
+            }
+
+            // this assignment is made possible by CoerceUnsized
+            const ALIGNED: &AlignedAsU32<[u8]> = &AlignedAsU32 {
+                _align: [],
+                // emits a token stream like `[10u8, 11u8, ... ]`
+                bytes: [ #(#bytes),* ]
+            };
+
+            &ALIGNED.bytes
+        }
+    };
+
+    Ok(result.into())
+}

citro3d-macros/tests/bad-shader.pica

@@ -0,0 +1,11 @@
+; Vertex shader that won't compile
+
+.out outpos position
+.out outclr color
+
+.proc main
+    mov outpos, 1
+    mov outclr, 0
+
+    end
+.end

citro3d-macros/tests/integration.pica

@@ -0,0 +1,14 @@
+; Trivial vertex shader
+
+.out outpos position
+.out outclr color
+
+.alias inpos v1
+.alias inclr v0
+
+.proc main
+    mov outpos, inpos
+    mov outclr, inclr
+
+    end
+.end

citro3d-macros/tests/integration.rs

@@ -0,0 +1,15 @@
+use citro3d_macros::include_shader;
+
+#[test]
+fn includes_shader_static() {
+    static SHADER_BYTES: &[u8] = include_shader!("integration.pica");
+
+    assert_eq!(SHADER_BYTES.len() % 4, 0);
+}
+
+#[test]
+fn includes_shader_const() {
+    const SHADER_BYTES: &[u8] = include_shader!("integration.pica");
+
+    assert_eq!(SHADER_BYTES.len() % 4, 0);
+}

citro3d/Cargo.toml

@@ -1,11 +1,12 @@
 [package]
 name = "citro3d"
-authors = [ "Rust3DS Org" ]
+authors = ["Rust3DS Org"]
 license = "MIT OR Apache-2.0"
 version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+citro3d-macros = { version = "0.1.0", path = "../citro3d-macros" }
 bitflags = "1.3.2"
 bytemuck = { version = "1.10.0", features = ["extern_crate_std"] }
 citro3d-sys = { git = "https://github.com/rust3ds/citro3d-rs.git" }

citro3d/examples/triangle.rs

@@ -3,17 +3,16 @@
 
 #![feature(allocator_api)]
 
-use citro3d::attrib;
-use citro3d::buffer;
+use std::ffi::CStr;
+use std::mem::MaybeUninit;
+
+use citro3d::macros::include_shader;
 use citro3d::render::{self, ClearFlags};
-use citro3d::{include_aligned_bytes, shader};
+use citro3d::{attrib, buffer, shader};
 use citro3d_sys::C3D_Mtx;
 use ctru::prelude::*;
 use ctru::services::gfx::{RawFrameBuffer, Screen, TopScreen3D};
 
-use std::ffi::CStr;
-use std::mem::MaybeUninit;
-
 #[repr(C)]
 #[derive(Copy, Clone)]
 struct Vec3 {
@@ -50,8 +49,7 @@ static VERTICES: &[Vertex] = &[
     },
 ];
 
-static SHADER_BYTES: &[u8] =
-    include_aligned_bytes!(concat!(env!("OUT_DIR"), "/examples/assets/vshader.shbin"));
+static SHADER_BYTES: &[u8] = include_shader!("assets/vshader.pica");
 
 fn main() {
     ctru::use_panic_handler();

citro3d/src/lib.rs

@@ -9,6 +9,11 @@ pub mod shader;
 use citro3d_sys::C3D_FrameDrawOn;
 pub use error::{Error, Result};
 
+pub mod macros {
+    //! Helper macros for working with shaders.
+    pub use citro3d_macros::*;
+}
+
 /// The single instance for using `citro3d`. This is the base type that an application
 /// should instantiate to use this library.
 #[non_exhaustive]

citro3d/src/shader.rs

@@ -7,11 +7,6 @@
 use std::error::Error;
 use std::mem::MaybeUninit;
 
-// Macros get exported at the crate root, so no reason to document this module.
-// It still needs to be `pub` for the helper struct it exports.
-#[doc(hidden)]
-pub mod macros;
-
 /// A PICA200 shader program. It may have one or both of:
 ///
 /// * A vertex shader [`Library`]