[derive] Add layout

Add a layout pass to the struct derive logic.

piet-gpu-derive/src/glsl.rs

@@ -0,0 +1,196 @@
+//! Generation of GLSL struct definitions and accessor functions.
+
+use std::fmt::Write;
+use std::ops::Deref;
+
+use crate::layout::{LayoutModule, LayoutType, LayoutTypeDef};
+use crate::parse::{GpuScalar, GpuType};
+
+pub fn gen_glsl(module: &LayoutModule) -> String {
+    let mut r = String::new();
+    writeln!(&mut r, "// Code auto-generated by piet-gpu-derive\n").unwrap();
+    // Note: GLSL needs definitions before uses. We could do a topological sort here,
+    // but easiest for now to just require that in spec.
+    for name in &module.def_names {
+        gen_refdef(&mut r, &name);
+    }
+    for name in &module.def_names {
+        let def = module.defs.get(name).unwrap();
+        if let (size, LayoutTypeDef::Struct(fields)) = def {
+            gen_struct_def(&mut r, &name, size.size, fields);
+        }
+    }
+    for name in &module.def_names {
+        let def = module.defs.get(name).unwrap();
+        match def {
+            (size, LayoutTypeDef::Struct(fields)) => {
+                gen_struct_read(&mut r, &module.name, &name, size.size, fields)
+            }
+            (_size, LayoutTypeDef::Enum(_)) => gen_enum_read(&mut r, &module.name, &name),
+        }
+    }
+    r
+}
+
+fn gen_refdef(r: &mut String, name: &str) {
+    writeln!(r, "struct {}Ref {{", name).unwrap();
+    writeln!(r, "    uint offset;").unwrap();
+    writeln!(r, "}};\n").unwrap();
+}
+
+fn gen_struct_def(r: &mut String, name: &str, size: usize, 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();
+
+    writeln!(r, "#define {}_size {}\n", name, size).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) {
+        // TODO: don't generate raw reads for inline structs
+        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();
+    }
+    writeln!(r, "    return s;").unwrap();
+    writeln!(r, "}}\n").unwrap();
+}
+
+fn gen_enum_read(r: &mut String, bufname: &str, name: &str) {
+    writeln!(r, "uint {}_tag({}Ref ref) {{", name, name).unwrap();
+    writeln!(r, "    return {}[ref.offset >> 2];", bufname).unwrap();
+    writeln!(r, "}}\n").unwrap();
+}
+
+fn gen_extract(offset: usize, ty: &GpuType) -> String {
+    match ty {
+        GpuType::Scalar(scalar) => gen_extract_scalar(offset, scalar),
+        GpuType::Vector(scalar, size) => {
+            let mut r = glsl_type(ty);
+            r.push_str("(");
+            for i in 0..*size {
+                if i != 0 {
+                    r.push_str(", ");
+                }
+                let el_offset = offset + i * scalar.size();
+                r.push_str(&gen_extract_scalar(el_offset, scalar));
+            }
+            r.push_str(")");
+            r
+        }
+        GpuType::InlineStruct(name) => format!(
+            "{}_read({}Ref({}))",
+            name,
+            name,
+            simplified_add("ref.offset", offset)
+        ),
+        GpuType::Ref(inner) => {
+            if let GpuType::InlineStruct(name) = inner.deref() {
+                format!(
+                    "{}Ref({})",
+                    name,
+                    gen_extract_scalar(offset, &GpuScalar::U32)
+                )
+            } else {
+                panic!("only know how to deal with Ref of struct")
+            }
+        }
+    }
+}
+
+fn gen_extract_scalar(offset: usize, ty: &GpuScalar) -> String {
+    match ty {
+        GpuScalar::F32 => format!("uintBitsToFloat(raw{})", offset / 4),
+        GpuScalar::U8 | GpuScalar::U16 | GpuScalar::U32 => extract_ubits(offset, ty.size()),
+        GpuScalar::I8 | GpuScalar::I16 | GpuScalar::I32 => extract_ibits(offset, ty.size()),
+    }
+}
+
+fn extract_ubits(offset: usize, nbytes: usize) -> String {
+    if nbytes == 4 {
+        return format!("raw{}", offset / 4);
+    }
+    let mask = (1 << (nbytes * 8)) - 1;
+    if offset % 4 == 0 {
+        format!("raw{} & 0x{:x}", offset / 4, mask)
+    } else if offset % 4 + nbytes == 4 {
+        format!("raw{} >> {}", offset / 4, (offset % 4) * 8)
+    } else {
+        format!("(raw{} >> {}) & 0x{:x}", offset / 4, (offset % 4) * 8, mask)
+    }
+}
+
+fn extract_ibits(offset: usize, nbytes: usize) -> String {
+    if nbytes == 4 {
+        return format!("int(raw{})", offset / 4);
+    }
+    if offset % 4 + nbytes == 4 {
+        format!("int(raw{}) >> {}", offset / 4, (offset % 4) * 8)
+    } else {
+        format!(
+            "(int(raw{}) << {}) >> {}",
+            offset / 4,
+            (3 - offset % 4) * 8,
+            (4 - nbytes) * 8
+        )
+    }
+}
+
+fn glsl_type(ty: &GpuType) -> String {
+    match ty {
+        GpuType::Scalar(scalar) => glsl_scalar(scalar).into(),
+        GpuType::Vector(scalar, size) => {
+            if *size == 1 {
+                glsl_scalar(scalar).into()
+            } else {
+                format!("{}{}", glsl_vecname(scalar), size)
+            }
+        }
+        GpuType::InlineStruct(name) => name.clone(),
+        GpuType::Ref(inner) => {
+            if let GpuType::InlineStruct(name) = inner.deref() {
+                format!("{}Ref", name)
+            } else {
+                panic!("only know how to deal with Ref of struct")
+            }
+        }
+    }
+}
+
+// GLSL type that can contain the scalar value.
+fn glsl_scalar(s: &GpuScalar) -> &'static str {
+    match s {
+        GpuScalar::F32 => "float",
+        GpuScalar::I8 | GpuScalar::I16 | GpuScalar::I32 => "int",
+        GpuScalar::U8 | GpuScalar::U16 | GpuScalar::U32 => "uint",
+    }
+}
+
+fn glsl_vecname(s: &GpuScalar) -> &'static str {
+    match s {
+        GpuScalar::F32 => "vec",
+        GpuScalar::I8 | GpuScalar::I16 | GpuScalar::I32 => "ivec",
+        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 {
+        String::from(var_name)
+    } else {
+        format!("{} + {}", var_name, c)
+    }
+}

piet-gpu-derive/src/layout.rs

@@ -0,0 +1,177 @@
+//! Logic for layout of structures in memory.
+
+use std::collections::{HashMap, HashSet};
+
+use crate::parse::{GpuModule, GpuType, GpuTypeDef};
+
+#[derive(Clone)]
+pub struct LayoutType {
+    size: Size,
+    pub ty: GpuType,
+}
+
+#[derive(Clone)]
+pub enum LayoutTypeDef {
+    /// Name, offset, field type. Make a separate struct?
+    Struct(Vec<(String, usize, LayoutType)>),
+    Enum(Vec<(String, Vec<(usize, LayoutType)>)>),
+}
+
+pub struct LayoutModule {
+    pub name: String,
+    pub def_names: Vec<String>,
+    pub defs: HashMap<String, (Size, LayoutTypeDef)>,
+}
+
+struct LayoutSession<'a> {
+    enum_variants: HashSet<String>,
+    orig_defs: HashMap<String, &'a GpuTypeDef>,
+    defs: HashMap<String, (Size, LayoutTypeDef)>,
+}
+
+#[derive(Clone, Copy)]
+pub struct Size {
+    pub size: usize,
+    alignment: usize,
+}
+
+impl LayoutType {
+    fn from_gpu(ty: &GpuType, session: &mut LayoutSession) -> LayoutType {
+        let size = session.get_size(ty);
+        LayoutType {
+            size,
+            ty: ty.clone(),
+        }
+    }
+}
+
+impl LayoutTypeDef {
+    // Maybe have a type representing the tuple?
+    fn from_gpu(def: &GpuTypeDef, session: &mut LayoutSession) -> (Size, LayoutTypeDef) {
+        match def {
+            GpuTypeDef::Struct(_name, fields) => {
+                let mut offset = 0;
+                let mut result = Vec::new();
+                for field in fields {
+                    let layout_ty = LayoutType::from_gpu(&field.1, session);
+                    offset += align_padding(offset, layout_ty.size.alignment);
+                    let size = layout_ty.size.size;
+                    result.push((field.0.clone(), offset, layout_ty));
+                    offset += size;
+                }
+                offset += align_padding(offset, 4);
+                let size = Size::new_struct(offset);
+                (size, LayoutTypeDef::Struct(result))
+            }
+            GpuTypeDef::Enum(en) => {
+                let mut result = Vec::new();
+                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) {
+                        4
+                    } else {
+                        0
+                    };
+                    for field in &variant.1 {
+                        let layout_ty = LayoutType::from_gpu(field, session);
+                        offset += align_padding(offset, layout_ty.size.alignment);
+                        let size = layout_ty.size.size;
+                        r2.push((offset, layout_ty));
+                        offset += size;
+                    }
+                    max_offset = max_offset.max(offset);
+                    result.push((variant.0.clone(), r2));
+                }
+                max_offset += align_padding(max_offset, 4);
+                let size = Size::new_struct(max_offset);
+                (size, LayoutTypeDef::Enum(result))
+            }
+        }
+    }
+}
+
+impl LayoutModule {
+    pub fn from_gpu(module: &GpuModule) -> LayoutModule {
+        let def_names = module
+            .defs
+            .iter()
+            .map(|def| def.name().to_owned())
+            .collect::<Vec<_>>();
+        let mut session = LayoutSession::new(module);
+        for def in &module.defs {
+            let _ = session.layout_def(def.name());
+        }
+        LayoutModule {
+            name: module.name.clone(),
+            def_names,
+            defs: session.defs,
+        }
+    }
+}
+
+impl<'a> LayoutSession<'a> {
+    fn new(module: &GpuModule) -> LayoutSession {
+        let mut orig_defs = HashMap::new();
+        let mut enum_variants = HashSet::new();
+        for def in &module.defs {
+            orig_defs.insert(def.name().to_owned(), def.clone());
+            if let GpuTypeDef::Enum(en) = def {
+                for variant in &en.variants {
+                    if let Some(GpuType::InlineStruct(name)) = variant.1.first() {
+                        enum_variants.insert(name.clone());
+                    }
+                }
+            }
+        }
+        LayoutSession {
+            enum_variants,
+            orig_defs,
+            defs: HashMap::new(),
+        }
+    }
+
+    /// Do layout of one def.
+    ///
+    /// This might be called recursively.
+    /// Note: expect stack overflow for circular dependencies.
+    fn layout_def(&mut self, name: &str) -> Size {
+        if let Some(def) = self.defs.get(name) {
+            return def.0;
+        }
+        let def = self.orig_defs.get(name).unwrap();
+        let layout = LayoutTypeDef::from_gpu(def, self);
+        let size = layout.0;
+        self.defs.insert(name.to_owned(), layout);
+        size
+    }
+
+    fn get_size(&mut self, ty: &GpuType) -> Size {
+        match ty {
+            GpuType::Scalar(scalar) => Size::new(scalar.size()),
+            GpuType::Vector(scalar, len) => Size::new(scalar.size() * len),
+            GpuType::Ref(_) => Size::new(4),
+            GpuType::InlineStruct(name) => self.layout_def(name),
+        }
+    }
+
+    fn is_enum_variant(&self, name: &str) -> bool {
+        self.enum_variants.contains(name)
+    }
+}
+
+impl Size {
+    fn new(size: usize) -> Size {
+        let alignment = if size < 4 { 1 } else { 4 };
+        Size { size, alignment }
+    }
+
+    fn new_struct(size: usize) -> Size {
+        let alignment = 4;
+        Size { size, alignment }
+    }
+}
+
+fn align_padding(offset: usize, alignment: usize) -> usize {
+    offset.wrapping_neg() & (alignment - 1)
+}

piet-gpu-derive/src/lib.rs

@@ -1,9 +1,12 @@
+mod glsl;
+mod layout;
 mod parse;
 
 use proc_macro::TokenStream;
-use quote::quote;
+use quote::{format_ident, quote};
 use syn::parse_macro_input;
 
+use layout::LayoutModule;
 use parse::GpuModule;
 
 #[proc_macro]
@@ -11,6 +14,13 @@ pub fn piet_gpu(input: TokenStream) -> TokenStream {
     let input = parse_macro_input!(input as syn::ItemMod);
     //println!("input: {:#?}", input);
     let module = GpuModule::from_syn(&input).unwrap();
-    let expanded = quote! {};
+    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! {
+        fn #gen_gpu_fn() -> String {
+            #glsl.into()
+        }
+    };
     expanded.into()
 }

piet-gpu-derive/src/parse.rs

@@ -2,11 +2,9 @@
 
 extern crate proc_macro;
 
-use std::collections::HashSet;
-
 use syn::{
-    Expr, ExprLit, Fields, FieldsNamed, FieldsUnnamed, GenericArgument, ItemEnum, ItemStruct,
-    Lit, PathArguments, TypeArray, TypePath,
+    Expr, ExprLit, Fields, FieldsNamed, FieldsUnnamed, GenericArgument, ItemEnum, ItemStruct, Lit,
+    PathArguments, TypeArray, TypePath,
 };
 
 /// A scalar that can be represented in a packed data structure.
@@ -60,6 +58,18 @@ impl GpuScalar {
             _ => None,
         })
     }
+
+    /// Size of scalar type.
+    ///
+    /// This is arguably a concern at the layout level, not syntax, but it's here because
+    /// it's not likely to be variable, so reduces the total number of types.
+    pub fn size(self) -> usize {
+        match self {
+            GpuScalar::F32 | GpuScalar::I32 | GpuScalar::U32 => 4,
+            GpuScalar::I8 | GpuScalar::U8 => 1,
+            GpuScalar::I16 | GpuScalar::U16 => 2,
+        }
+    }
 }
 
 fn ty_as_single_ident(ty: &syn::Type) -> Option<String> {
@@ -166,7 +176,14 @@ impl GpuTypeDef {
                 Err("unknown item".into())
             }
         }
-    }   
+    }
+
+    pub fn name(&self) -> &str {
+        match self {
+            GpuTypeDef::Struct(name, _) => name,
+            GpuTypeDef::Enum(en) => &en.name,
+        }
+    }
 }
 
 impl GpuModule {
@@ -179,10 +196,7 @@ impl GpuModule {
                 defs.push(def);
             }
         }
-        Ok(GpuModule {
-            name,
-            defs,
-        })
+        Ok(GpuModule { name, defs })
     }
 }
 

piet-gpu-hal/src/lib.rs

@@ -2,7 +2,6 @@
 ///
 /// This abstraction is inspired by gfx-hal, but is specialized to the needs of piet-gpu.
 /// In time, it may go away and be replaced by either gfx-hal or wgpu.
-
 pub mod vulkan;
 
 /// This isn't great but is expedient.

piet-gpu-hal/src/vulkan.rs

@@ -74,17 +74,14 @@ impl VkInstance {
                 None,
             )?;
 
-            Ok(VkInstance {
-                entry,
-                instance,
-            })
+            Ok(VkInstance { entry, instance })
         }
     }
 
     /// Create a device from the instance, suitable for compute.
-    /// 
+    ///
     /// # Safety
-    /// 
+    ///
     /// The caller is responsible for making sure that the instance outlives the device.
     /// We could enforce that, for example having an `Arc` of the raw instance, but for
     /// now keep things simple.
@@ -127,11 +124,7 @@ impl crate::Device for VkDevice {
     type Pipeline = Pipeline;
     type MemFlags = MemFlags;
 
-    fn create_buffer(
-        &self,
-        size: u64,
-        mem_flags: MemFlags,
-    ) -> Result<Buffer, Error> {
+    fn create_buffer(&self, size: u64, mem_flags: MemFlags) -> Result<Buffer, Error> {
         unsafe {
             let device = &self.device.device;
             let buffer = device.create_buffer(
@@ -197,20 +190,22 @@ impl crate::Device for VkDevice {
             None,
         )?;
 
-        let pipeline = device.create_compute_pipelines(
-            vk::PipelineCache::null(),
-            &[vk::ComputePipelineCreateInfo::builder()
-                .stage(
-                    vk::PipelineShaderStageCreateInfo::builder()
-                        .stage(vk::ShaderStageFlags::COMPUTE)
-                        .module(compute_shader_module)
-                        .name(&entry_name)
-                        .build(),
-                )
-                .layout(pipeline_layout)
-                .build()],
-            None,
-        ).map_err(|(_pipeline, err)| err)?[0];
+        let pipeline = device
+            .create_compute_pipelines(
+                vk::PipelineCache::null(),
+                &[vk::ComputePipelineCreateInfo::builder()
+                    .stage(
+                        vk::PipelineShaderStageCreateInfo::builder()
+                            .stage(vk::ShaderStageFlags::COMPUTE)
+                            .module(compute_shader_module)
+                            .name(&entry_name)
+                            .build(),
+                    )
+                    .layout(pipeline_layout)
+                    .build()],
+                None,
+            )
+            .map_err(|(_pipeline, err)| err)?[0];
         Ok(Pipeline {
             pipeline,
             pipeline_layout,
@@ -333,11 +328,7 @@ impl crate::Device for VkDevice {
         Ok(())
     }
 
-    unsafe fn write_buffer<T: Sized>(
-        &self,
-        buffer: &Buffer,
-        contents: &[T],
-    ) -> Result<(), Error> {
+    unsafe fn write_buffer<T: Sized>(&self, buffer: &Buffer, contents: &[T]) -> Result<(), Error> {
         let device = &self.device.device;
         let buf = device.map_memory(
             buffer.buffer_memory,