[derive] Add Rust encoding of structs

This also requires some traits and structs in an `encode` module (the
one in piet-metal/piet-gpu-types works, and was used to test).

piet-gpu-derive/src/derive.rs

@@ -0,0 +1,162 @@
+//! Generation of Rust derive functions for encoding.
+
+use quote::{format_ident, quote};
+
+use crate::layout::{LayoutModule, LayoutTypeDef};
+use crate::parse::{GpuScalar, GpuType};
+
+pub fn gen_derive(module: &LayoutModule) -> proc_macro2::TokenStream {
+    let mut ts = proc_macro2::TokenStream::new();
+    let module_name = format_ident!("{}", module.name);
+    for name in &module.def_names {
+        let def = module.defs.get(name).unwrap();
+        ts.extend(gen_derive_def(name, def.0.size, &def.1));
+    }
+    quote! {
+        mod #module_name {
+            #ts
+        }
+    }
+}
+
+fn gen_derive_def(name: &str, size: usize, def: &LayoutTypeDef) -> proc_macro2::TokenStream {
+    let name_id = format_ident!("{}", name);
+    match def {
+        LayoutTypeDef::Struct(fields) => {
+            let mut gen_fields = proc_macro2::TokenStream::new();
+            let mut encode_fields = proc_macro2::TokenStream::new();
+            for (field_name, offset, ty) in fields {
+                let field_name_id = format_ident!("{}", field_name);
+                let gen_ty = gen_derive_ty(&ty.ty);
+                let gen_field = quote! {
+                    pub #field_name_id: #gen_ty,
+                };
+                gen_fields.extend(gen_field);
+
+                encode_fields.extend(gen_encode_field(field_name, *offset, &ty.ty));
+            }
+            quote! {
+                pub struct #name_id {
+                    #gen_fields
+                }
+
+                impl crate::encoder::Encode for #name_id {
+                    fn fixed_size() -> usize {
+                        #size
+                    }
+                    fn encode_to(&self, buf: &mut [u8]) {
+                        #encode_fields
+                    }
+                }
+            }
+        }
+        LayoutTypeDef::Enum(variants) => {
+            let mut gen_variants = proc_macro2::TokenStream::new();
+            let mut cases = proc_macro2::TokenStream::new();
+            for (variant_ix, (variant_name, payload)) in variants.iter().enumerate() {
+                let variant_id = format_ident!("{}", variant_name);
+                let field_tys = payload.iter().map(|(_offset, ty)| gen_derive_ty(&ty.ty));
+                let variant = quote! {
+                    #variant_id(#(#field_tys),*),
+                };
+                gen_variants.extend(variant);
+
+                let mut args = Vec::new();
+                let mut field_encoders = proc_macro2::TokenStream::new();
+                for (i, (offset, _ty)) in payload.iter().enumerate() {
+                    let field_id = format_ident!("f{}", i);
+                    let field_encoder = quote! {
+                        #field_id.encode_to(&mut buf[#offset..]);
+                    };
+                    field_encoders.extend(field_encoder);
+                    args.push(field_id);
+                }
+                let case = quote! {
+                    #name_id::#variant_id(#(#args),*) => {
+                        buf[0..4].copy_from_slice(&#variant_ix.to_le_bytes());
+                        #field_encoders
+                    }
+                };
+                cases.extend(case);
+            }
+            quote! {
+                pub enum #name_id {
+                    #gen_variants
+                }
+
+                impl crate::encoder::Encode for #name_id {
+                    fn fixed_size() -> usize {
+                        #size
+                    }
+                    fn encode_to(&self, buf: &mut [u8]) {
+                        match self {
+                            #cases
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+/// Generate a Rust type.
+fn gen_derive_ty(ty: &GpuType) -> proc_macro2::TokenStream {
+    match ty {
+        GpuType::Scalar(s) => gen_derive_scalar_ty(s),
+        GpuType::Vector(s, len) => {
+            let scalar = gen_derive_scalar_ty(s);
+            quote! { [#scalar; #len] }
+        }
+        GpuType::InlineStruct(name) => {
+            let name_id = format_ident!("{}", name);
+            quote! { #name_id }
+        }
+        GpuType::Ref(ty) => {
+            let gen_ty = gen_derive_ty(ty);
+            quote! { crate::encoder::Ref<#gen_ty> }
+        }
+    }
+}
+
+fn gen_derive_scalar_ty(ty: &GpuScalar) -> proc_macro2::TokenStream {
+    match ty {
+        GpuScalar::F32 => quote!(f32),
+        GpuScalar::I8 => quote!(i8),
+        GpuScalar::I16 => quote!(i16),
+        GpuScalar::I32 => quote!(i32),
+        GpuScalar::U8 => quote!(u8),
+        GpuScalar::U16 => quote!(u16),
+        GpuScalar::U32 => quote!(u32),
+    }
+}
+
+fn gen_encode_field(name: &str, offset: usize, ty: &GpuType) -> proc_macro2::TokenStream {
+    let name_id = format_ident!("{}", name);
+    match ty {
+        GpuType::Scalar(s) => {
+            let end = offset + s.size();
+            quote! {
+                buf[#offset..#end].copy_from_slice(&self.#name_id.to_le_bytes());
+            }
+        }
+        GpuType::Vector(s, len) => {
+            let size = s.size();
+            quote! {
+                for i in 0..#len {
+                    let offset = #offset + i * #size;
+                    buf[offset..offset + #size].copy_from_slice(&self.#name_id[i].to_le_bytes());
+                }
+            }
+        }
+        GpuType::Ref(_) => {
+            quote! {
+                buf[#offset..#offset + 4].copy_from_slice(&self.#name_id.offset().to_le_bytes());
+            }
+        }
+        _ => {
+            quote! {
+                &self.#name_id.encode_to(&mut buf[#offset..]);
+            }
+        }
+    }
+}

piet-gpu-derive/src/lib.rs

@@ -1,3 +1,4 @@
+mod derive;
 mod glsl;
 mod layout;
 mod parse;
@@ -17,10 +18,13 @@ pub fn piet_gpu(input: TokenStream) -> TokenStream {
     let layout = LayoutModule::from_gpu(&module);
     let glsl = glsl::gen_glsl(&layout);
     let gen_gpu_fn = format_ident!("gen_gpu_{}", layout.name);
-    let expanded = quote! {
+    let mut expanded = quote! {
         fn #gen_gpu_fn() -> String {
             #glsl.into()
         }
     };
+    if layout.rust_encode {
+        expanded.extend(derive::gen_derive(&layout));
+    }
     expanded.into()
 }