[derive] Add writers, enums

This adds shader generation of writers (and cleans up some other
things) and much better support for enums.

piet-gpu-derive/src/glsl.rs

@@ -15,18 +15,32 @@ pub fn gen_glsl(module: &LayoutModule) -> String {
         gen_refdef(&mut r, &name);
     }
     for name in &module.def_names {
-        let def = module.defs.get(name).unwrap();
+        match module.defs.get(name).unwrap() {
-        if let (size, LayoutTypeDef::Struct(fields)) = def {
+            (size, LayoutTypeDef::Struct(fields)) => {
-            gen_struct_def(&mut r, &name, size.size, fields);
+                gen_struct_def(&mut r, name, fields);
+                gen_item_def(&mut r, name, size.size);
+            }
+            (size, LayoutTypeDef::Enum(en)) => {
+                gen_enum_def(&mut r, name, en);
+                gen_item_def(&mut r, name, size.size);
+            }
         }
     }
     for name in &module.def_names {
         let def = module.defs.get(name).unwrap();
         match def {
-            (size, LayoutTypeDef::Struct(fields)) => {
+            (_size, LayoutTypeDef::Struct(fields)) => {
-                gen_struct_read(&mut r, &module.name, &name, size.size, fields)
+                gen_struct_read(&mut r, &module.name, &name, fields);
+                if module.gpu_write {
+                    gen_struct_write(&mut r, &module.name, &name, fields);
+                }
+            }
+            (_size, LayoutTypeDef::Enum(en)) => {
+                gen_enum_read(&mut r, &module.name, &name, en);
+                if module.gpu_write {
+                    gen_enum_write(&mut r, &module.name, &name, en);
+                }
             }
-            (_size, LayoutTypeDef::Enum(_)) => gen_enum_read(&mut r, &module.name, &name),
         }
     }
     r
@@ -38,29 +52,51 @@ fn gen_refdef(r: &mut String, name: &str) {
     writeln!(r, "}};\n").unwrap();
 }
 
-fn gen_struct_def(r: &mut String, name: &str, size: usize, fields: &[(String, usize, LayoutType)]) {
+fn gen_struct_def(r: &mut String, name: &str, fields: &[(String, usize, LayoutType)]) {
     writeln!(r, "struct {} {{", name).unwrap();
     for (name, _offset, ty) in fields {
         writeln!(r, "    {} {};", glsl_type(&ty.ty), name).unwrap();
     }
     writeln!(r, "}};\n").unwrap();
+}
 
+fn gen_enum_def(r: &mut String, name: &str, variants: &[(String, Vec<(usize, LayoutType)>)]) {
+    for (i, (var_name, _payload)) in variants.iter().enumerate() {
+        writeln!(r, "#define {}_{} {}", name, var_name, i).unwrap();
+    }
+}
+
+fn gen_item_def(r: &mut String, name: &str, size: usize) {
     writeln!(r, "#define {}_size {}\n", name, size).unwrap();
+    writeln!(
+        r,
+        "{}Ref {}_index({}Ref ref, uint index) {{",
+        name, name, name
+    )
+    .unwrap();
+    writeln!(
+        r,
+        "    return {}Ref(ref.offset + index * {}_size);",
+        name, name
+    )
+    .unwrap();
+    writeln!(r, "}}\n").unwrap();
 }
 
 fn gen_struct_read(
     r: &mut String,
     bufname: &str,
     name: &str,
-    size: usize,
     fields: &[(String, usize, LayoutType)],
 ) {
     writeln!(r, "{} {}_read({}Ref ref) {{", name, name, name).unwrap();
     writeln!(r, "    uint ix = ref.offset >> 2;").unwrap();
-    for i in 0..(size / 4) {
+    let coverage = crate::layout::struct_coverage(fields, false);
-        // TODO: don't generate raw reads for inline structs
+    for (i, fields) in coverage.iter().enumerate() {
+        if !fields.is_empty() {
             writeln!(r, "    uint raw{} = {}[ix + {}];", i, bufname, i).unwrap();
         }
+    }
     writeln!(r, "    {} s;", name).unwrap();
     for (name, offset, ty) in fields {
         writeln!(r, "    s.{} = {};", name, gen_extract(*offset, &ty.ty)).unwrap();
@@ -69,10 +105,35 @@ fn gen_struct_read(
     writeln!(r, "}}\n").unwrap();
 }
 
-fn gen_enum_read(r: &mut String, bufname: &str, name: &str) {
+fn gen_enum_read(
+    r: &mut String,
+    bufname: &str,
+    name: &str,
+    variants: &[(String, Vec<(usize, LayoutType)>)],
+) {
     writeln!(r, "uint {}_tag({}Ref ref) {{", name, name).unwrap();
     writeln!(r, "    return {}[ref.offset >> 2];", bufname).unwrap();
     writeln!(r, "}}\n").unwrap();
+    for (var_name, payload) in variants {
+        if payload.len() == 1 {
+            if let GpuType::InlineStruct(structname) = &payload[0].1.ty {
+                writeln!(
+                    r,
+                    "{} {}_{}_read({}Ref ref) {{",
+                    structname, name, var_name, name
+                )
+                .unwrap();
+                writeln!(
+                    r,
+                    "    return {}_read({}Ref(ref.offset + {}));",
+                    structname, structname, payload[0].0
+                )
+                .unwrap();
+                writeln!(r, "}}\n").unwrap();
+            }
+        }
+        // TODO: support for variants that aren't one struct.
+    }
 }
 
 fn gen_extract(offset: usize, ty: &GpuType) -> String {
@@ -141,14 +202,146 @@ fn extract_ibits(offset: usize, nbytes: usize) -> String {
         format!("int(raw{}) >> {}", offset / 4, (offset % 4) * 8)
     } else {
         format!(
-            "(int(raw{}) << {}) >> {}",
+            "int(raw{} << {}) >> {}",
             offset / 4,
-            (3 - offset % 4) * 8,
+            ((4 - nbytes) - offset % 4) * 8,
             (4 - nbytes) * 8
         )
     }
 }
 
+// Writing
+
+fn gen_struct_write(
+    r: &mut String,
+    bufname: &str,
+    name: &str,
+    fields: &[(String, usize, LayoutType)],
+) {
+    writeln!(r, "void {}_write({}Ref ref, {} s) {{", name, name, name).unwrap();
+    let coverage = crate::layout::struct_coverage(fields, true);
+    for (i, field_ixs) in coverage.iter().enumerate() {
+        let mut pieces = Vec::new();
+        for field_ix in field_ixs {
+            let (name, offset, ty) = &fields[*field_ix];
+            match &ty.ty {
+                GpuType::Scalar(scalar) => {
+                    let inner = format!("s.{}", name);
+                    pieces.push(gen_pack_bits_scalar(scalar, *offset, &inner));
+                }
+                GpuType::Vector(scalar, len) => {
+                    let size = scalar.size();
+                    let ix_lo = (i * 4 - offset) / size;
+                    let ix_hi = ((4 + i * 4 - offset) / size).min(*len);
+                    for ix in ix_lo..ix_hi {
+                        let scalar_offset = offset + ix * size;
+                        let inner = format!("s.{}.{}", name, &"xyzw"[ix..ix + 1]);
+                        pieces.push(gen_pack_bits_scalar(scalar, scalar_offset, &inner));
+                    }
+                }
+                GpuType::InlineStruct(structname) => {
+                    writeln!(
+                        r,
+                        "    {}_write({}Ref({}), s.{});",
+                        structname,
+                        structname,
+                        simplified_add("ref.offset", *offset),
+                        name
+                    )
+                    .unwrap();
+                }
+                GpuType::Ref(_) => pieces.push(format!("s.{}.offset", name)),
+            }
+        }
+        if !pieces.is_empty() {
+            write!(r, "    {}[{}] = ", bufname, i).unwrap();
+            for (j, piece) in pieces.iter().enumerate() {
+                if j != 0 {
+                    write!(r, " | ").unwrap();
+                }
+                write!(r, "{}", piece).unwrap();
+            }
+            writeln!(r, ";").unwrap();
+        }
+    }
+    writeln!(r, "}}\n").unwrap();
+}
+
+fn gen_pack_bits_scalar(ty: &GpuScalar, offset: usize, inner: &str) -> String {
+    let shift = (offset % 4) * 8;
+    let bits = match ty {
+        GpuScalar::F32 => format!("floatBitsToUint({})", inner),
+        // Note: this doesn't mask small unsigned int types; the caller is
+        // responsible for making sure they don't overflow.
+        GpuScalar::U8 | GpuScalar::U16 | GpuScalar::U32 => inner.into(),
+        GpuScalar::I8 => {
+            if shift == 24 {
+                format!("uint({})", inner)
+            } else {
+                format!("(uint({}) & 0xff)", inner)
+            }
+        }
+        GpuScalar::I16 => {
+            if shift == 16 {
+                format!("uint({})", inner)
+            } else {
+                format!("(uint({}) & 0xffff)", inner)
+            }
+        }
+        GpuScalar::I32 => format!("uint({})", inner),
+    };
+    if shift == 0 {
+        bits
+    } else {
+        format!("({} << {})", bits, shift)
+    }
+}
+
+fn gen_enum_write(
+    r: &mut String,
+    bufname: &str,
+    name: &str,
+    variants: &[(String, Vec<(usize, LayoutType)>)],
+) {
+    for (var_name, payload) in variants {
+        if payload.is_empty() {
+            writeln!(r, "void {}_{}_write({}Ref ref) {{", name, var_name, name).unwrap();
+            writeln!(
+                r,
+                "    {}[ref.offset >> 2] = {}_{};",
+                bufname, name, var_name
+            )
+            .unwrap();
+            writeln!(r, "}}\n").unwrap();
+        } else if payload.len() == 1 {
+            if let GpuType::InlineStruct(structname) = &payload[0].1.ty {
+                writeln!(
+                    r,
+                    "void {}_{}_write({}Ref ref, {} s) {{",
+                    name, var_name, name, structname
+                )
+                .unwrap();
+                writeln!(
+                    r,
+                    "    {}[ref.offset >> 2] = {}_{};",
+                    bufname, name, var_name
+                )
+                .unwrap();
+                writeln!(
+                    r,
+                    "    {}_write({}Ref(ref.offset + {}), s);",
+                    structname, structname, payload[0].0
+                )
+                .unwrap();
+                writeln!(r, "}}\n").unwrap();
+            }
+        }
+        // TODO: support for variants that aren't one struct.
+    }
+}
+
+// Utility functions
+
 fn glsl_type(ty: &GpuType) -> String {
     match ty {
         GpuType::Scalar(scalar) => glsl_scalar(scalar).into(),
@@ -186,6 +379,7 @@ fn glsl_vecname(s: &GpuScalar) -> &'static str {
         GpuScalar::U8 | GpuScalar::U16 | GpuScalar::U32 => "uvec",
     }
 }
+
 /// If `c = 0`, return `"var_name"`, else `"var_name + c"`
 fn simplified_add(var_name: &str, c: usize) -> String {
     if c == 0 {

piet-gpu-derive/src/layout.rs

@@ -1,5 +1,11 @@
 //! Logic for layout of structures in memory.
 
+// This is fairly simple now, but there are some extensions that are likely:
+// * Addition of f16 types
+//   + These will probably have 2-byte alignments to support `packHalf2x16`
+// * 1 byte tag values (so small struct fields can be packed along with tag)
+// * (Possibly) reordering for better packing
+
 use std::collections::{HashMap, HashSet};
 
 use crate::parse::{GpuModule, GpuType, GpuTypeDef};
@@ -21,6 +27,16 @@ pub struct LayoutModule {
     pub name: String,
     pub def_names: Vec<String>,
     pub defs: HashMap<String, (Size, LayoutTypeDef)>,
+    enum_variants: HashSet<String>,
+
+    /// Generate shader code to write the module.
+    ///
+    /// This is derived from the presence of the `gpu_write` attribute in the source module.
+    pub gpu_write: bool,
+    /// Generate Rust code to encode the module.
+    ///
+    /// This is derived from the presence of the `rust_encode` attribute in the source module.
+    pub rust_encode: bool,
 }
 
 struct LayoutSession<'a> {
@@ -50,6 +66,13 @@ impl LayoutTypeDef {
     fn from_gpu(def: &GpuTypeDef, session: &mut LayoutSession) -> (Size, LayoutTypeDef) {
         match def {
             GpuTypeDef::Struct(_name, fields) => {
+                // TODO: We want to be able to pack enums more tightly, in particular
+                // other struct fields along with the enum tag. Structs in that category
+                // (first field has an alignment < 4, serve as enum variant) will have a
+                // different layout. This is why we're tracking `is_enum_variant`.
+                //
+                // But it's a bit of YAGNI for now; we're currently reserving 4 bytes for
+                // the tag, so structure layout doesn't care.
                 let mut offset = 0;
                 let mut result = Vec::new();
                 for field in fields {
@@ -68,11 +91,7 @@ impl LayoutTypeDef {
                 let mut max_offset = 0;
                 for variant in &en.variants {
                     let mut r2 = Vec::new();
-                    let mut offset = if session.is_enum_variant(&en.name) {
+                    let mut offset = 4;
-                        4
-                    } else {
-                        0
-                    };
                     for field in &variant.1 {
                         let layout_ty = LayoutType::from_gpu(field, session);
                         offset += align_padding(offset, layout_ty.size.alignment);
@@ -102,12 +121,22 @@ impl LayoutModule {
         for def in &module.defs {
             let _ = session.layout_def(def.name());
         }
+        let gpu_write = module.attrs.contains("gpu_write");
+        let rust_encode = module.attrs.contains("rust_encode");
         LayoutModule {
             name: module.name.clone(),
+            gpu_write,
+            rust_encode,
             def_names,
+            enum_variants: session.enum_variants,
             defs: session.defs,
         }
     }
+
+    #[allow(unused)]
+    pub fn is_enum_variant(&self, name: &str) -> bool {
+        self.enum_variants.contains(name)
+    }
 }
 
 impl<'a> LayoutSession<'a> {
@@ -155,14 +184,52 @@ impl<'a> LayoutSession<'a> {
         }
     }
 
+    #[allow(unused)]
     fn is_enum_variant(&self, name: &str) -> bool {
         self.enum_variants.contains(name)
     }
 }
 
+/// Compute coverage of fields.
+///
+/// Each element of the result represents a list of fields for one 4-byte chunk of
+/// the struct layout. Inline structs are only included if requested.
+pub fn struct_coverage(
+    fields: &[(String, usize, LayoutType)],
+    include_inline: bool,
+) -> Vec<Vec<usize>> {
+    let mut result: Vec<Vec<usize>> = Vec::new();
+    for (i, (_name, offset, ty)) in fields.iter().enumerate() {
+        let size = match ty.ty {
+            GpuType::Scalar(scalar) => scalar.size(),
+            GpuType::Vector(scalar, len) => scalar.size() * len,
+            GpuType::Ref(_) => 4,
+            GpuType::InlineStruct(_) => {
+                if include_inline {
+                    4
+                } else {
+                    0
+                }
+            }
+        };
+        if size > 0 {
+            for ix in (offset / 4)..(offset + size + 3) / 4 {
+                if ix >= result.len() {
+                    result.resize_with(ix + 1, Default::default);
+                }
+                result[ix].push(i);
+            }
+        }
+    }
+    result
+}
+
 impl Size {
     fn new(size: usize) -> Size {
-        let alignment = if size < 4 { 1 } else { 4 };
+        // Note: there is special case we could do better:
+        // `(u8, u16, u8)`, where the alignment could be 1. However,
+        // this case can also be solved by reordering.
+        let alignment = size.min(4);
         Size { size, alignment }
     }
 

piet-gpu-derive/src/parse.rs

@@ -2,6 +2,8 @@
 
 extern crate proc_macro;
 
+use std::collections::HashSet;
+
 use syn::{
     Expr, ExprLit, Fields, FieldsNamed, FieldsUnnamed, GenericArgument, ItemEnum, ItemStruct, Lit,
     PathArguments, TypeArray, TypePath,
@@ -42,6 +44,7 @@ pub enum GpuTypeDef {
 
 pub struct GpuModule {
     pub name: String,
+    pub attrs: HashSet<String>,
     pub defs: Vec<GpuTypeDef>,
 }
 
@@ -72,22 +75,6 @@ impl GpuScalar {
     }
 }
 
-fn ty_as_single_ident(ty: &syn::Type) -> Option<String> {
-    if let syn::Type::Path(TypePath {
-        path: syn::Path { segments, .. },
-        ..
-    }) = ty
-    {
-        if segments.len() == 1 {
-            let seg = &segments[0];
-            if seg.arguments == PathArguments::None {
-                return Some(seg.ident.to_string());
-            }
-        }
-    }
-    None
-}
-
 impl GpuType {
     fn from_syn(ty: &syn::Type) -> Result<Self, String> {
         //println!("gputype {:#?}", ty);
@@ -189,6 +176,12 @@ impl GpuTypeDef {
 impl GpuModule {
     pub fn from_syn(module: &syn::ItemMod) -> Result<Self, String> {
         let name = module.ident.to_string();
+        let mut attrs = HashSet::new();
+        for attr in &module.attrs {
+            if let Some(id) = path_as_single_ident(&attr.path) {
+                attrs.insert(id.to_owned());
+            }
+        }
         let mut defs = Vec::new();
         if let Some((_brace, items)) = &module.content {
             for item in items {
@@ -196,7 +189,25 @@ impl GpuModule {
                 defs.push(def);
             }
         }
-        Ok(GpuModule { name, defs })
+        Ok(GpuModule { name, attrs, defs })
+    }
+}
+
+fn path_as_single_ident(path: &syn::Path) -> Option<String> {
+    if path.segments.len() == 1 {
+        let seg = &path.segments[0];
+        if seg.arguments == PathArguments::None {
+            return Some(seg.ident.to_string());
+        }
+    }
+    None
+}
+
+fn ty_as_single_ident(ty: &syn::Type) -> Option<String> {
+    if let syn::Type::Path(TypePath { path, .. }) = ty {
+        path_as_single_ident(path)
+    } else {
+        None
     }
 }