makepad/libs/render/examples/ao_render.rs

//! Software render of real Kenney props, with baked AO on and off.
//!
//!   cargo run -p makepad-render --release --example ao_render
//!
//! The GPU path lives in the app, which this crate cannot drive — so this
//! reproduces the shader's shading maths on the CPU and writes a JPEG. That is
//! enough to judge the ONE thing in question: what the baked AO term does to a
//! prop. Anything muddy, striped or banded shows up here exactly as it would
//! on screen, because the arithmetic is the same.
//!
//! Left half of each tile: AO off. Right half: AO on.

use jpeg_encoder::{ColorType, Encoder};
use makepad_draw::makepad_math::Vec3f;
use makepad_render::model::{StaticModel, MODEL_VERTEX_FLOATS};

const TILE: usize = 300;
const COLS: usize = 3;

/// Same terms the shader uses, so what this shows is what the GPU would draw.
const SUN_DIR: Vec3f = Vec3f { x: 0.45, y: 0.78, z: 0.44 };
const SUN_COLOR: f32 = 0.72;
const SKY: f32 = 0.40;
const GROUND: f32 = 0.20;

fn main() {
    let root = "apps/sandbox/resources/models/kenney";
    // A prop with eaves, an arch, a slatted bench, a barrel, a tree, a wall.
    let wanted = [
        ("house", "fantasy-town-kit"),
        ("arch", "castle-kit"),
        ("bench", "graveyard-kit"),
        ("barrel", "survival-kit"),
        ("tree", "nature-kit"),
        ("wall", "castle-kit"),
    ];

    let mut picks: Vec<(String, std::path::PathBuf)> = Vec::new();
    for (needle, pack) in wanted {
        let dir = std::path::Path::new(root).join(pack);
        let Ok(rd) = std::fs::read_dir(&dir) else { continue };
        let mut best: Option<std::path::PathBuf> = None;
        for e in rd.flatten() {
            let p = e.path();
            if p.extension().map(|x| x != "glb").unwrap_or(true) {
                continue;
            }
            let name = p.file_stem().unwrap().to_string_lossy().to_lowercase();
            if name.contains(needle) {
                // Shortest matching name = the plainest variant.
                if best.as_ref().map(|b| name.len() < b.to_string_lossy().len()).unwrap_or(true) {
                    best = Some(p);
                }
            }
        }
        if let Some(b) = best {
            picks.push((format!("{pack}/{}", b.file_stem().unwrap().to_string_lossy()), b));
        }
    }
    if picks.is_empty() {
        eprintln!("no models found under {root} — run apps/sandbox/download_assets.sh");
        return;
    }

    let rows = picks.len().div_ceil(COLS);
    let (w, h) = (TILE * COLS, TILE * rows);
    let mut fb = vec![0u8; w * h * 3];
    // Mid grey, so both over- and under-darkening are visible against it.
    for p in fb.iter_mut() {
        *p = 96;
    }

    for (i, (name, path)) in picks.iter().enumerate() {
        let Ok(bytes) = std::fs::read(path) else { continue };
        let Ok(m) = StaticModel::parse_glb(&bytes) else { continue };
        let ox = (i % COLS) * TILE;
        let oy = (i / COLS) * TILE;
        // Left half without AO, right half with — one model, one camera, so
        // the only difference in the image is the term being judged.
        draw(&mut fb, w, ox, oy, TILE / 2, TILE, &m, false);
        draw(&mut fb, w, ox + TILE / 2, oy, TILE / 2, TILE, &m, true);
        println!("{i}: {name}  ({} verts)", m.vertex_count());
    }

    let out = std::env::args().nth(1).unwrap_or_else(|| {
        "/private/tmp/claude-501/-Users-admin-makepad-makepad/99a6fb4f-a075-40cd-8cda-fdb93c19da1d/scratchpad/ao_compare.jpg".into()
    });
    let enc = Encoder::new_file(&out, 92).expect("encoder");
    enc.encode(&fb, w as u16, h as u16, ColorType::Rgb).expect("encode");
    println!("wrote {out} ({w}x{h})");
}

/// Rasterise the model into a sub-rect, with a z-buffer and the shader's
/// lighting. `ao` selects whether the baked term is applied.
#[allow(clippy::too_many_arguments)]
fn draw(
    fb: &mut [u8],
    fb_w: usize,
    ox: usize,
    oy: usize,
    vw: usize,
    vh: usize,
    m: &StaticModel,
    ao: bool,
) {
    let mut zbuf = vec![f32::MAX; vw * vh];
    let centre = Vec3f {
        x: (m.min.x + m.max.x) * 0.5,
        y: (m.min.y + m.max.y) * 0.5,
        z: (m.min.z + m.max.z) * 0.5,
    };
    let span = (m.max.x - m.min.x)
        .max(m.max.y - m.min.y)
        .max(m.max.z - m.min.z)
        .max(1.0e-4);

    // Fixed three-quarter view: high enough to see the ground-facing crevices
    // that AO is mostly about.
    let yaw = 0.7f32;
    let pitch = 0.45f32;
    let (sy, cy) = (yaw.sin(), yaw.cos());
    let (sp, cp) = (pitch.sin(), pitch.cos());
    let scale = vw.min(vh) as f32 / (span * 1.45);

    let project = |p: Vec3f| -> (f32, f32, f32) {
        let x = p.x - centre.x;
        let y = p.y - centre.y;
        let z = p.z - centre.z;
        let rx = x * cy + z * sy;
        let rz = -x * sy + z * cy;
        let ry = y * cp + rz * sp;
        let depth = -y * sp + rz * cp;
        (
            vw as f32 * 0.5 + rx * scale,
            vh as f32 * 0.5 - ry * scale,
            depth,
        )
    };

    let vert = |i: usize| -> (Vec3f, Vec3f, f32) {
        let b = i * MODEL_VERTEX_FLOATS;
        let p = Vec3f {
            x: m.vertices[b],
            y: m.vertices[b + 1],
            z: m.vertices[b + 2],
        };
        let n = oct_decode(m.vertices[b + 3]);
        let packed = m.vertices[b + 5].to_bits();
        let unorm = |shift: u32| ((packed >> shift) & 0xff) as f32 / 255.0;
        // rgb = material tint, w = baked AO (see model.rs).
        let tint = (unorm(0) + unorm(8) + unorm(16)) / 3.0;
        let a = unorm(24);
        (p, n, if ao { a } else { 1.0 } * (0.25 + 0.75 * tint))
    };

    for t in 0..m.indices.len() / 3 {
        let (p0, n0, a0) = vert(m.indices[t * 3] as usize);
        let (p1, n1, a1) = vert(m.indices[t * 3 + 1] as usize);
        let (p2, n2, a2) = vert(m.indices[t * 3 + 2] as usize);
        let s0 = project(p0);
        let s1 = project(p1);
        let s2 = project(p2);

        let area = (s1.0 - s0.0) * (s2.1 - s0.1) - (s2.0 - s0.0) * (s1.1 - s0.1);
        if area.abs() < 1.0e-6 {
            continue;
        }
        let minx = s0.0.min(s1.0).min(s2.0).floor().max(0.0) as usize;
        let maxx = (s0.0.max(s1.0).max(s2.0).ceil() as usize).min(vw.saturating_sub(1));
        let miny = s0.1.min(s1.1).min(s2.1).floor().max(0.0) as usize;
        let maxy = (s0.1.max(s1.1).max(s2.1).ceil() as usize).min(vh.saturating_sub(1));

        for py in miny..=maxy {
            for px in minx..=maxx {
                let fx = px as f32 + 0.5;
                let fy = py as f32 + 0.5;
                let w0 = ((s1.0 - fx) * (s2.1 - fy) - (s2.0 - fx) * (s1.1 - fy)) / area;
                let w1 = ((s2.0 - fx) * (s0.1 - fy) - (s0.0 - fx) * (s2.1 - fy)) / area;
                let w2 = 1.0 - w0 - w1;
                if w0 < 0.0 || w1 < 0.0 || w2 < 0.0 {
                    continue;
                }
                let depth = s0.2 * w0 + s1.2 * w1 + s2.2 * w2;
                let zi = py * vw + px;
                if depth >= zbuf[zi] {
                    continue;
                }
                zbuf[zi] = depth;

                let n = Vec3f {
                    x: n0.x * w0 + n1.x * w1 + n2.x * w2,
                    y: n0.y * w0 + n1.y * w1 + n2.y * w2,
                    z: n0.z * w0 + n1.z * w1 + n2.z * w2,
                };
                let l = (n.x * n.x + n.y * n.y + n.z * n.z).sqrt().max(1.0e-6);
                let ny = n.y / l;
                let dp = ((n.x * SUN_DIR.x + n.y * SUN_DIR.y + n.z * SUN_DIR.z) / l).max(0.0);
                let shade = a0 * w0 + a1 * w1 + a2 * w2;
                // The shader's arithmetic: AO scales AMBIENT only, direct is
                // untouched. Reproduced exactly so this image is predictive.
                let hemi = (ny * 0.5 + 0.5).clamp(0.0, 1.0);
                let ambient = GROUND + (SKY - GROUND) * hemi;
                let lit = (ambient * shade + SUN_COLOR * dp).clamp(0.0, 1.4);
                let v = (lit * 210.0).clamp(0.0, 255.0) as u8;
                let o = ((oy + py) * fb_w + ox + px) * 3;
                fb[o] = v;
                fb[o + 1] = v;
                fb[o + 2] = v;
            }
        }
    }
}

/// Inverse of skin.rs's oct_encode — the same fold the shader does.
fn oct_decode(packed: f32) -> Vec3f {
    let bits = packed.to_bits();
    let f16 = |h: u32| -> f32 {
        let s = ((h >> 15) & 1) as i32;
        let e = ((h >> 10) & 0x1f) as i32;
        let m = (h & 0x3ff) as i32;
        let v = if e == 0 {
            (m as f32) * 2.0f32.powi(-24)
        } else {
            (1.0 + m as f32 / 1024.0) * 2.0f32.powi(e - 15)
        };
        if s == 1 {
            -v
        } else {
            v
        }
    };
    let ex = f16(bits & 0xffff);
    let ey = f16((bits >> 16) & 0xffff);
    let nz = 1.0 - ex.abs() - ey.abs();
    let t = (-nz).max(0.0);
    let sx = if ex >= 0.0 { 1.0 } else { -1.0 };
    let sy = if ey >= 0.0 { 1.0 } else { -1.0 };
    let v = Vec3f {
        x: ex - t * sx,
        y: ey - t * sy,
        z: nz,
    };
    let l = (v.x * v.x + v.y * v.y + v.z * v.z).sqrt().max(1.0e-6);
    Vec3f {
        x: v.x / l,
        y: v.y / l,
        z: v.z / l,
    }
}