vello/piet-gpu-derive/src/glsl.rs

//! 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)
    }
}