cgpu 0.1.0

#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct ShaderBuilder {
    bindings: Vec<BindingDesc>,
    structs: Vec<WgslStruct>,
    helpers: Vec<String>,
    main_body: Vec<String>,
    workgroup_size: [u32; 3],
}

impl ShaderBuilder {
    pub fn new() -> Self {
        Self {
            workgroup_size: [1, 1, 1],
            ..Self::default()
        }
    }

    pub fn workgroup_size(mut self, x: u32, y: u32, z: u32) -> Self {
        self.workgroup_size = [x, y, z];
        self
    }

    pub fn add_struct(mut self, s: WgslStruct) -> Self {
        self.structs.push(s);
        self
    }

    pub fn add_binding(mut self, desc: BindingDesc) -> Self {
        self.bindings.push(desc);
        self
    }

    pub fn add_helper(mut self, wgsl: &str) -> Self {
        self.helpers.push(wgsl.to_string());
        self
    }

    pub fn main_line(mut self, wgsl: &str) -> Self {
        self.main_body.push(format!("    {wgsl}"));
        self
    }

    pub fn main_block(mut self, wgsl: &str) -> Self {
        for line in wgsl.trim().lines() {
            self.main_body.push(format!("    {}", line.trim_end()));
        }
        self
    }

    pub fn build(self) -> String {
        let mut source = Vec::new();

        for s in self.structs {
            source.push(s.build());
        }

        for binding in self.bindings {
            source.push(binding.build());
        }

        source.extend(self.helpers);
        source.push(format!(
            "@compute @workgroup_size({}, {}, {})",
            self.workgroup_size[0], self.workgroup_size[1], self.workgroup_size[2]
        ));
        source.push("fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {".to_string());
        source.extend(self.main_body);
        source.push("}".to_string());
        source.join("\n\n")
    }

    #[cfg(feature = "gpu")]
    pub fn compile(self, device: &wgpu::Device) -> wgpu::ShaderModule {
        device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("cgpu-inline-shader"),
            source: wgpu::ShaderSource::Wgsl(self.build().into()),
        })
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct BindingDesc {
    pub group: u32,
    pub binding: u32,
    pub name: String,
    pub address_space: String,
    pub access: Option<String>,
    pub ty: String,
}

impl BindingDesc {
    pub fn storage_buffer_read_only(binding: u32, name: &str, ty: &str) -> Self {
        Self {
            group: 0,
            binding,
            name: name.to_string(),
            address_space: "storage".to_string(),
            access: Some("read".to_string()),
            ty: ty.to_string(),
        }
    }

    pub fn storage_buffer(binding: u32, name: &str, ty: &str) -> Self {
        Self {
            group: 0,
            binding,
            name: name.to_string(),
            address_space: "storage".to_string(),
            access: Some("read_write".to_string()),
            ty: ty.to_string(),
        }
    }

    pub fn uniform(binding: u32, name: &str, ty: &str) -> Self {
        Self {
            group: 0,
            binding,
            name: name.to_string(),
            address_space: "uniform".to_string(),
            access: None,
            ty: ty.to_string(),
        }
    }

    pub fn group(mut self, group: u32) -> Self {
        self.group = group;
        self
    }

    fn build(self) -> String {
        let access = self
            .access
            .map(|access| format!(", {access}"))
            .unwrap_or_default();
        format!(
            "@group({}) @binding({})\nvar<{}{}> {}: {};",
            self.group, self.binding, self.address_space, access, self.name, self.ty
        )
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct WgslStruct {
    pub name: String,
    pub fields: Vec<WgslField>,
}

impl WgslStruct {
    pub fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
            fields: Vec::new(),
        }
    }

    pub fn field(mut self, name: &str, ty: &str) -> Self {
        self.fields.push(WgslField {
            name: name.to_string(),
            ty: ty.to_string(),
        });
        self
    }

    fn build(self) -> String {
        let mut source = format!("struct {} {{\n", self.name);
        for field in self.fields {
            source.push_str(&format!("    {}: {},\n", field.name, field.ty));
        }
        source.push('}');
        source
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct WgslField {
    pub name: String,
    pub ty: String,
}