Skip to main content

oxframe_2d/
shader.rs

1//! A shader node: WGSL `Params` plus `shade`, checked with naga.
2//!
3//! The waves preset is rasterized on the CPU, so a check without a GPU
4//! still paints it. Both backends then draw that image, which keeps the
5//! preset the same on either one. The lens_drops preset reads the frame
6//! drawn before it (see [`crate::post`]). A custom `shade` is checked on
7//! every draw (a changed file is picked up without a rebuild) and run on
8//! the CPU by [`crate::wgsl_cpu`], natively and in the browser.
9
10use std::cell::Cell;
11
12use kurbo::{Affine, Point};
13use oxframe_core::{Color, Signal};
14
15use crate::draw::Primitive;
16use crate::graph::{NodeData, NodeKind};
17use crate::image::RasterImage;
18
19/// Where the WGSL text comes from. Stored on a shader node.
20#[doc(hidden)]
21#[derive(Clone)]
22pub enum ShaderCode {
23    /// Source kept on the node.
24    Wgsl(String),
25    /// Read again on every draw, so an edit shows up without a rebuild.
26    File(std::path::PathBuf),
27}
28
29/// Which built-in a shader node runs. Stored on a shader node.
30#[doc(hidden)]
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum Preset {
33    /// User WGSL: checked, and run on the CPU per pixel (not in wasm).
34    Custom,
35    /// [`waves`](crate::waves): painted inside the node's box.
36    Waves,
37    /// [`lens_drops`](crate::lens_drops): changes the frame drawn before it.
38    LensDrops,
39}
40
41impl Preset {
42    /// `"custom"`, `"waves"` or `"lens_drops"`.
43    pub fn name(self) -> &'static str {
44        match self {
45            Self::Custom => "custom",
46            Self::Waves => "waves",
47            Self::LensDrops => "lens_drops",
48        }
49    }
50}
51
52/// Parameters of the lens_drops preset, all numbers.
53const LENS_PARAMS: [&str; 3] = ["wetness", "seed", "drop_size"];
54
55/// Inputs the lens_drops preset reads.
56const LENS_INPUTS: [&str; 1] = ["scene"];
57
58/// One value passed to `Params`. Stored on a shader node.
59#[doc(hidden)]
60pub struct Uniform {
61    pub name: String,
62    pub value: UniformValue,
63}
64
65/// A number, a point, a color, or a bool stored as 0 or 1.
66#[doc(hidden)]
67#[allow(dead_code)]
68pub enum UniformValue {
69    /// An `f32` field.
70    Num(Signal<f64>),
71    /// A `vec2<f32>` field.
72    Vec2(Signal<kurbo::Vec2>),
73    /// A `vec4<f32>` color field.
74    Color(Signal<Color>),
75    /// A `bool` field, stored as 0 or 1 so it can be animated.
76    Bool(Signal<f64>),
77}
78
79/// A WGSL problem found while painting.
80#[derive(Debug, Clone, PartialEq)]
81pub struct ShaderDiag {
82    /// `.name(..)`, or the node kind and id.
83    pub node: String,
84    /// `file:line` of the `add` call, when it was recorded.
85    pub source: Option<String>,
86    /// `shader_error`, `param_unknown`, or `shader_slow`.
87    pub rule: &'static str,
88    /// What to change.
89    pub hint: String,
90}
91
92const WAVES: &str = r#"struct Params { amp: f32, speed: f32, deep: vec4<f32>, foam: vec4<f32> }
93@group(1) @binding(0) var<uniform> p: Params;
94
95fn shade(i: ShadeIn) -> vec4<f32> {
96    let h = sin(i.world.x * 0.008 + ox.time * p.speed) * p.amp;
97    let t = smoothstep(h - 0.02, h + 0.02, i.uv.y - 0.5);
98    return mix(p.foam, p.deep, t);
99}
100"#;
101
102/// WGSL of the waves preset.
103pub(crate) fn waves_wgsl() -> &'static str {
104    WAVES
105}
106
107const PRELUDE: &str = r#"struct Ox {
108    time: f32,
109    frame: f32,
110    resolution: vec2<f32>,
111    size: vec2<f32>,
112    seed: f32,
113    view: mat4x4<f32>,
114    progress: f32,
115}
116@group(0) @binding(0) var<uniform> ox: Ox;
117
118struct ShadeIn {
119    uv: vec2<f32>,
120    local: vec2<f32>,
121    world: vec2<f32>,
122    screen: vec2<f32>,
123}
124
125"#;
126
127const EPILOGUE: &str = r#"
128struct _OxVertexOut {
129    @builtin(position) pos: vec4<f32>,
130    @location(0) uv: vec2<f32>,
131    @location(1) local: vec2<f32>,
132    @location(2) world: vec2<f32>,
133    @location(3) screen: vec2<f32>,
134}
135
136@vertex
137fn _ox_vs(@builtin(vertex_index) index: u32) -> _OxVertexOut {
138    var positions = array<vec2<f32>, 3>(
139        vec2<f32>(-1.0, -1.0),
140        vec2<f32>(3.0, -1.0),
141        vec2<f32>(-1.0, 3.0),
142    );
143    let p = positions[index];
144    var out: _OxVertexOut;
145    out.pos = vec4<f32>(p, 0.0, 1.0);
146    out.uv = vec2<f32>(p.x * 0.5 + 0.5, 0.5 - p.y * 0.5);
147    out.local = out.uv;
148    out.world = out.uv;
149    out.screen = out.uv;
150    return out;
151}
152
153@fragment
154fn _ox_fs(in: _OxVertexOut) -> @location(0) vec4<f32> {
155    var s: ShadeIn;
156    s.uv = in.uv;
157    s.local = in.local;
158    s.world = in.world;
159    s.screen = in.screen;
160    return shade(s);
161}
162"#;
163
164thread_local! {
165    /// Scene size and frame rate of the graph being drawn, for `ox`.
166    static SCENE: Cell<((f64, f64), f64)> = const { Cell::new(((1920.0, 1080.0), 60.0)) };
167    static DRAFT: Cell<bool> = const { Cell::new(false) };
168    static BUDGET_MS: Cell<f64> = const { Cell::new(250.0) };
169    static DIAGS: std::cell::RefCell<Vec<ShaderDiag>> = const { std::cell::RefCell::new(Vec::new()) };
170}
171
172/// Sets the scene size for `ox.resolution` and `ox.screen`; the graph
173/// calls this before it builds a draw list.
174pub(crate) fn set_scene_size(size: (f64, f64)) {
175    SCENE.with(|s| s.set((size, s.get().1)));
176}
177
178/// Frames per second, for `ox.frame` in custom shaders. The engine sets it
179/// before it draws.
180#[doc(hidden)]
181pub fn set_frame_rate(fps: f64) {
182    SCENE.with(|s| s.set((s.get().0, fps.max(1.0))));
183}
184
185pub(crate) fn clear_diags() {
186    DIAGS.with(|d| d.borrow_mut().clear());
187}
188
189pub(crate) fn push_diag(diag: ShaderDiag) {
190    DIAGS.with(|d| d.borrow_mut().push(diag));
191}
192
193pub(crate) fn take_diags() -> Vec<ShaderDiag> {
194    DIAGS.with(|d| std::mem::take(&mut *d.borrow_mut()))
195}
196
197/// Restores the previous draft flag when dropped.
198pub struct ShaderDraft {
199    prev: bool,
200}
201
202impl Drop for ShaderDraft {
203    fn drop(&mut self) {
204        DRAFT.with(|d| d.set(self.prev));
205    }
206}
207
208/// Paints shader nodes at half resolution until dropped.
209///
210/// Timings and positions stay the same. `oxframe check` uses this so a
211/// shader layer does not dominate the budget.
212pub fn draft_shaders() -> ShaderDraft {
213    let prev = DRAFT.with(Cell::get);
214    DRAFT.with(|d| d.set(true));
215    ShaderDraft { prev }
216}
217
218/// Milliseconds a shader raster may take before `shader_slow`.
219pub fn set_shader_budget_ms(ms: f64) {
220    BUDGET_MS.with(|b| b.set(ms));
221}
222
223fn draft() -> bool {
224    DRAFT.with(Cell::get)
225}
226
227fn budget_ms() -> f64 {
228    BUDGET_MS.with(Cell::get)
229}
230
231/// Compile and parameter problems for this node, and a video that could
232/// not be opened. Empty for other nodes.
233pub(crate) fn problems(n: &NodeData) -> Vec<ShaderDiag> {
234    if let NodeKind::Video { data, .. } = &n.kind {
235        return data
236            .problem
237            .iter()
238            .map(|p| ShaderDiag {
239                node: label(n),
240                source: n.source.clone(),
241                rule: "video_unavailable",
242                hint: format!("{}: {p}", data.path.display()),
243            })
244            .collect();
245    }
246    let NodeKind::Shader {
247        code,
248        uniforms,
249        preset,
250        inputs,
251        ..
252    } = &n.kind
253    else {
254        return Vec::new();
255    };
256    let label = label(n);
257    let source = n.source.clone();
258    let diag = |hint: String| {
259        vec![ShaderDiag {
260            node: label.clone(),
261            source: source.clone(),
262            rule: "param_unknown",
263            hint,
264        }]
265    };
266    let known: &[&str] = match preset {
267        Preset::LensDrops => &LENS_INPUTS,
268        Preset::Waves => &[],
269        Preset::Custom => &[],
270    };
271    if *preset == Preset::Custom {
272        for (name, _) in inputs {
273            if !is_ident(name) {
274                return diag(format!(
275                    "input `{name}` is read as `sample_{name}(uv)`, so its name must be letters, digits and `_`."
276                ));
277            }
278        }
279    } else {
280        for (name, _) in inputs {
281            if !known.contains(&name.as_str()) {
282                return diag(unknown("input", name, known));
283            }
284        }
285    }
286    if *preset == Preset::LensDrops {
287        for u in uniforms {
288            if !LENS_PARAMS.contains(&u.name.as_str()) {
289                return diag(unknown("parameter", &u.name, &LENS_PARAMS));
290            }
291            if !matches!(u.value, UniformValue::Num(_)) {
292                return diag(format!("`{}` of lens_drops is a number.", u.name));
293            }
294        }
295        return Vec::new();
296    }
297    let text = match load(code) {
298        Ok(text) => text,
299        Err(hint) => {
300            return vec![ShaderDiag {
301                node: label,
302                source,
303                rule: "shader_error",
304                hint,
305            }];
306        }
307    };
308    let names: Vec<&str> = inputs.iter().map(|(n, _)| n.as_str()).collect();
309    match check_wgsl(&text, uniforms, &names) {
310        Ok(()) => Vec::new(),
311        Err(diag) => vec![ShaderDiag {
312            node: label,
313            source,
314            rule: diag.0,
315            hint: diag.1,
316        }],
317    }
318}
319
320/// Raster of a waves node, or nothing when this node is not the preset.
321pub(crate) fn raster(
322    n: &NodeData,
323    size: (f64, f64),
324    tf: Affine,
325    time: f64,
326    opacity: f64,
327) -> Option<Primitive> {
328    let NodeKind::Shader {
329        preset,
330        uniforms,
331        code,
332        ..
333    } = &n.kind
334    else {
335        return None;
336    };
337    let (w, h) = (size.0.max(1.0), size.1.max(1.0));
338    let scale = if draft() { 0.5 } else { 1.0 };
339    let iw = ((w * scale).round() as u32).clamp(1, 1920);
340    let ih = ((h * scale).round() as u32).clamp(1, 1080);
341    let started = Stopwatch::start();
342    let rgba = match preset {
343        Preset::Waves => {
344            let (amp, speed, deep, foam) = wave_uniforms(uniforms);
345            paint(
346                iw,
347                ih,
348                |uv, local| {
349                    let world = tf * Point::new(local.0, local.1);
350                    let crest = (world.x * 0.008 + time * speed).sin() * amp;
351                    let t = smoothstep(crest - 0.02, crest + 0.02, uv.1 - 0.5);
352                    Ok(mix(foam, deep, t, opacity))
353                },
354                (w, h),
355            )
356        }
357        Preset::Custom if !reads_frame_input(n) => {
358            custom(n, code, uniforms, (w, h), tf, time, opacity, (iw, ih))
359        }
360        _ => {
361            let _ = code;
362            return None;
363        }
364    };
365    let rgba = match rgba {
366        Ok(rgba) => rgba,
367        Err(hint) => {
368            push_diag(ShaderDiag {
369                node: label(n),
370                source: n.source.clone(),
371                rule: "shader_error",
372                hint,
373            });
374            return None;
375        }
376    };
377    let ms = started.millis();
378    if ms > budget_ms() {
379        push_diag(ShaderDiag {
380            node: label(n),
381            source: n.source.clone(),
382            rule: "shader_slow",
383            hint: format!(
384                "This shader took {ms:.0} ms in draft. The budget is {:.0} ms.",
385                budget_ms()
386            ),
387        });
388    }
389    let rect = kurbo::Rect::new(-w / 2.0, -h / 2.0, w / 2.0, h / 2.0);
390    Some(Primitive::Image {
391        image: std::rc::Rc::new(RasterImage {
392            width: iw,
393            height: ih,
394            id: hash_bytes(&rgba),
395            rgba,
396        }),
397        rect,
398        opacity: 1.0,
399    })
400}
401
402/// The frame marker of a lens_drops node, or nothing for other shaders.
403///
404/// `area` is the node's box in scene coordinates. The node reads the frame
405/// drawn before it whatever its input is named: a backdrop outside a layer
406/// is the frame.
407pub(crate) fn read_frame(
408    nodes: &[NodeData],
409    n: &NodeData,
410    size: (f64, f64),
411    tf: Affine,
412    time: f64,
413    opacity: f64,
414) -> Option<crate::post::Shade> {
415    if reads_frame_input(n) {
416        return custom_read_frame(nodes, n, size, tf, time, opacity);
417    }
418    let NodeKind::Shader {
419        preset: Preset::LensDrops,
420        uniforms,
421        ..
422    } = &n.kind
423    else {
424        return None;
425    };
426    let num = |name: &str, default: f64| {
427        uniforms
428            .iter()
429            .find_map(|u| match (&u.name[..], &u.value) {
430                (m, UniformValue::Num(v)) if m == name => Some(v.get()),
431                _ => None,
432            })
433            .unwrap_or(default)
434    };
435    let (w, h) = (size.0.max(0.0), size.1.max(0.0));
436    let area = tf.transform_rect_bbox(kurbo::Rect::new(-w / 2.0, -h / 2.0, w / 2.0, h / 2.0));
437    Some(crate::post::Shade {
438        op: crate::post::ShadeOp::LensDrops {
439            wetness: num("wetness", 0.6),
440            seed: num("seed", 0.0).max(0.0) as u64,
441            scale: num("drop_size", 1.0),
442            time,
443        },
444        opacity,
445        area: Some(area),
446    })
447}
448
449/// The frame-reading pass of a custom shader, or nothing (with a
450/// `shader_error`) when it does not compile.
451fn custom_read_frame(
452    nodes: &[NodeData],
453    n: &NodeData,
454    size: (f64, f64),
455    tf: Affine,
456    time: f64,
457    opacity: f64,
458) -> Option<crate::post::Shade> {
459    let NodeKind::Shader {
460        code,
461        uniforms,
462        inputs,
463        ..
464    } = &n.kind
465    else {
466        return None;
467    };
468    let (program, globals) = match prepare(n, code, uniforms, size, time) {
469        Ok(ok) => ok,
470        Err(hint) => {
471            push_diag(ShaderDiag {
472                node: label(n),
473                source: n.source.clone(),
474                rule: "shader_error",
475                hint,
476            });
477            return None;
478        }
479    };
480    let frame = inputs
481        .iter()
482        .filter(|(_, i)| matches!(i, crate::post::Input::Frame | crate::post::Input::Backdrop))
483        .map(|(name, _)| name.clone())
484        .collect();
485    let offscreen = inputs
486        .iter()
487        .filter_map(|(name, i)| match i {
488            crate::post::Input::Node(id) if (id.0 as usize) < nodes.len() => {
489                let (items, bounds) = crate::graph::offscreen_items(nodes, *id, time);
490                Some(NodeInput {
491                    name: name.clone(),
492                    items,
493                    bounds,
494                })
495            }
496            _ => None,
497        })
498        .collect();
499    let ((sw, sh), _) = SCENE.with(Cell::get);
500    let (w, h) = (size.0.max(0.0), size.1.max(0.0));
501    let area = tf.transform_rect_bbox(kurbo::Rect::new(-w / 2.0, -h / 2.0, w / 2.0, h / 2.0));
502    Some(crate::post::Shade {
503        op: crate::post::ShadeOp::Custom(Box::new(CustomShade {
504            program,
505            globals,
506            frame,
507            offscreen,
508            tf,
509            size: (w, h),
510            scene: (sw, sh),
511        })),
512        opacity,
513        area: Some(area),
514    })
515}
516
517/// `no input `x`. Did you mean `y`? Available: …` for a preset.
518fn unknown(what: &str, name: &str, known: &[&str]) -> String {
519    let list = if known.is_empty() {
520        "none".to_owned()
521    } else {
522        known.join(", ")
523    };
524    match closest(name, known) {
525        Some(n) => format!("no {what} `{name}`. Did you mean `{n}`? Available: {list}."),
526        None => format!("no {what} `{name}`. Available: {list}."),
527    }
528}
529
530/// Wall clock for `shader_slow`.
531///
532/// `Instant::now()` panics on `wasm32-unknown-unknown`, which ended the
533/// playground's waves example in "unreachable executed". A check never runs
534/// in the browser, so there the watch reads 0 and `shader_slow` stays quiet.
535#[cfg(not(target_arch = "wasm32"))]
536struct Stopwatch(std::time::Instant);
537
538#[cfg(target_arch = "wasm32")]
539struct Stopwatch;
540
541impl Stopwatch {
542    fn start() -> Self {
543        #[cfg(not(target_arch = "wasm32"))]
544        {
545            Self(std::time::Instant::now())
546        }
547        #[cfg(target_arch = "wasm32")]
548        {
549            Self
550        }
551    }
552
553    fn millis(&self) -> f64 {
554        #[cfg(not(target_arch = "wasm32"))]
555        {
556            self.0.elapsed().as_secs_f64() * 1000.0
557        }
558        #[cfg(target_arch = "wasm32")]
559        {
560            0.0
561        }
562    }
563}
564
565/// RGBA bytes of an `iw` × `ih` raster of a `size` box: `pixel` gets the
566/// uv (0..1, from the top left) and the local point (from the center).
567fn paint(
568    iw: u32,
569    ih: u32,
570    mut pixel: impl FnMut((f64, f64), (f64, f64)) -> Result<Color, String>,
571    size: (f64, f64),
572) -> Result<Vec<u8>, String> {
573    let (w, h) = size;
574    let mut rgba = vec![0u8; iw as usize * ih as usize * 4];
575    for py in 0..ih {
576        for px in 0..iw {
577            let u = (f64::from(px) + 0.5) / f64::from(iw);
578            let v = (f64::from(py) + 0.5) / f64::from(ih);
579            let c = pixel((u, v), (u * w - w / 2.0, v * h - h / 2.0))?;
580            let i = (py as usize * iw as usize + px as usize) * 4;
581            rgba[i] = byte(c.r);
582            rgba[i + 1] = byte(c.g);
583            rgba[i + 2] = byte(c.b);
584            rgba[i + 3] = byte(c.a);
585        }
586    }
587    Ok(rgba)
588}
589
590/// A custom `shade`, run per pixel of its own box.
591#[allow(clippy::too_many_arguments)]
592fn custom(
593    n: &NodeData,
594    code: &ShaderCode,
595    uniforms: &[Uniform],
596    size: (f64, f64),
597    tf: Affine,
598    time: f64,
599    opacity: f64,
600    raster: (u32, u32),
601) -> Result<Vec<u8>, String> {
602    let (program, globals) = prepare(n, code, uniforms, size, time)?;
603    let ((sw, sh), _) = SCENE.with(Cell::get);
604    let (iw, ih) = raster;
605    let a = opacity.clamp(0.0, 1.0) as f32;
606    let (w, h) = size;
607    let total = iw as usize * ih as usize;
608    let mut rgba = vec![0u8; total * 4];
609    run_tiles(
610        &program,
611        &globals,
612        total,
613        |i| {
614            let (px, py) = ((i % iw as usize) as f64, (i / iw as usize) as f64);
615            let u = (px + 0.5) / f64::from(iw);
616            let v = (py + 0.5) / f64::from(ih);
617            let local = (u * w - w / 2.0, v * h - h / 2.0);
618            let world = tf * Point::new(local.0, local.1);
619            [
620                u,
621                v,
622                local.0,
623                local.1,
624                world.x,
625                world.y,
626                world.x + sw / 2.0,
627                world.y + sh / 2.0,
628            ]
629        },
630        |i, [r, g, b, alpha]| {
631            let o = i * 4;
632            rgba[o] = byte(r);
633            rgba[o + 1] = byte(g);
634            rgba[o + 2] = byte(b);
635            rgba[o + 3] = byte(alpha * a);
636        },
637    )?;
638    Ok(rgba)
639}
640
641/// The compiled shader of `n` and the values it reads: `ox`, the `Params`
642/// fields, and its image inputs.
643fn prepare(
644    n: &NodeData,
645    code: &ShaderCode,
646    uniforms: &[Uniform],
647    size: (f64, f64),
648    time: f64,
649) -> Result<
650    (
651        std::rc::Rc<crate::wgsl_cpu::Program>,
652        crate::wgsl_cpu::Globals,
653    ),
654    String,
655> {
656    use crate::wgsl_cpu::{Globals, Value};
657    let inputs = match &n.kind {
658        NodeKind::Shader { inputs, .. } => inputs.as_slice(),
659        _ => &[],
660    };
661    let names: Vec<&str> = inputs.iter().map(|(n, _)| n.as_str()).collect();
662    let text = load(code)?;
663    let program = compiled(&text, &names).map_err(|(_, hint)| hint)?;
664    let ((sw, sh), fps) = SCENE.with(Cell::get);
665    let identity = Value::C(
666        (0..4)
667            .map(|c| {
668                let mut col = [0.0; 4];
669                col[c] = 1.0;
670                Value::fv(&col)
671            })
672            .collect(),
673    );
674    let ox = Value::C(vec![
675        Value::f(time as f32),
676        Value::f((time * fps + 1e-6).floor() as f32),
677        Value::fv(&[sw as f32, sh as f32]),
678        Value::fv(&[size.0 as f32, size.1 as f32]),
679        Value::f(0.0),
680        identity,
681        Value::f(0.0),
682    ]);
683    let params = program
684        .params()
685        .into_iter()
686        .map(|name| {
687            uniforms
688                .iter()
689                .find(|u| u.name == name)
690                .map_or(Value::f(0.0), |u| match &u.value {
691                    UniformValue::Num(v) => Value::f(v.get() as f32),
692                    UniformValue::Vec2(v) => {
693                        let v = v.get();
694                        Value::fv(&[v.x as f32, v.y as f32])
695                    }
696                    UniformValue::Color(v) => {
697                        let c = v.get();
698                        Value::fv(&[c.r, c.g, c.b, c.a])
699                    }
700                    UniformValue::Bool(v) => Value::B(1, [v.get() > 0.5, false, false, false]),
701                })
702        })
703        .collect();
704    let globals = Globals {
705        ox,
706        params: Value::C(params),
707        inputs: inputs
708            .iter()
709            .filter_map(|(name, input)| match input {
710                crate::post::Input::Image(img) => Some((name.clone(), img.clone())),
711                _ => None,
712            })
713            .collect(),
714    };
715    Ok((program, globals))
716}
717
718/// Runs `shade` for pixels `0..total`, a tile at a time. `at` gives a
719/// pixel's `ShadeIn` (uv, local, world, screen); `put` takes its color.
720fn run_tiles(
721    program: &crate::wgsl_cpu::Program,
722    globals: &crate::wgsl_cpu::Globals,
723    total: usize,
724    at: impl Fn(usize) -> [f64; 8],
725    mut put: impl FnMut(usize, [f32; 4]),
726) -> Result<(), String> {
727    use crate::wgsl_cpu::{LANES, V};
728    let mut runner = program.runner(globals);
729    let mut start = 0;
730    while start < total {
731        let count = (total - start).min(LANES);
732        let mut comps = vec![[0.0f32; LANES]; 8];
733        for k in 0..count {
734            for (c, x) in comps.iter_mut().zip(at(start + k)) {
735                c[k] = x as f32;
736            }
737        }
738        let input = V::C(comps.chunks(2).map(|pair| V::F(pair.to_vec())).collect());
739        let out = runner
740            .run(&input, count)
741            .map_err(|e| format!("{e} (when oxframe runs the shader on the CPU)"))?;
742        for k in 0..count {
743            let rgba = out
744                .lane(k)
745                .floats()
746                .ok_or("`shade` must return vec4<f32>")?;
747            put(start + k, rgba);
748        }
749        start += count;
750    }
751    Ok(())
752}
753
754/// A custom `shade` with a frame or backdrop input. It runs between the
755/// parts of the draw list, over the pixels drawn before it (see
756/// [`crate::post`]), and is painted over them inside its box.
757#[derive(Clone)]
758pub struct CustomShade {
759    program: std::rc::Rc<crate::wgsl_cpu::Program>,
760    globals: crate::wgsl_cpu::Globals,
761    /// Inputs that read the frame.
762    frame: Vec<String>,
763    /// Inputs that read a node drawn alone.
764    offscreen: Vec<NodeInput>,
765    tf: Affine,
766    size: (f64, f64),
767    scene: (f64, f64),
768}
769
770/// A node a custom shader reads: its draw list and its box in scene
771/// coordinates.
772#[derive(Clone)]
773struct NodeInput {
774    name: String,
775    items: Vec<crate::DrawItem>,
776    bounds: kurbo::Rect,
777}
778
779impl PartialEq for NodeInput {
780    fn eq(&self, other: &Self) -> bool {
781        self.name == other.name
782            && self.bounds == other.bounds
783            && self.items.len() == other.items.len()
784            && self.items.iter().zip(&other.items).all(|(a, b)| {
785                a.node == b.node
786                    && a.transform == b.transform
787                    && crate::graph::prim_same(&a.prim, &b.prim)
788            })
789    }
790}
791
792impl std::fmt::Debug for CustomShade {
793    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
794        f.debug_struct("CustomShade")
795            .field("frame", &self.frame)
796            .field("size", &self.size)
797            .finish_non_exhaustive()
798    }
799}
800
801impl PartialEq for CustomShade {
802    fn eq(&self, other: &Self) -> bool {
803        std::rc::Rc::ptr_eq(&self.program, &other.program)
804            && self.globals.ox == other.globals.ox
805            && self.globals.params == other.globals.params
806            && self.globals.inputs.len() == other.globals.inputs.len()
807            && self
808                .globals
809                .inputs
810                .iter()
811                .zip(&other.globals.inputs)
812                .all(|(a, b)| a.0 == b.0 && a.1.id == b.1.id)
813            && self.frame == other.frame
814            && self.offscreen == other.offscreen
815            && self.tf == other.tf
816            && self.size == other.size
817            && self.scene == other.scene
818    }
819}
820
821impl CustomShade {
822    /// Draw lists of the nodes this shader reads, in order. The engine
823    /// draws each alone on a clear target of the frame's size and passes
824    /// the images to [`crate::post::apply_with`].
825    pub fn offscreen_lists(&self) -> Vec<&[crate::DrawItem]> {
826        self.offscreen.iter().map(|n| n.items.as_slice()).collect()
827    }
828
829    /// Runs the shader over `img` inside the node's box. `to_px` maps
830    /// scene coordinates to the pixels of `img`; `drawn` holds the images
831    /// of [`offscreen_lists`](Self::offscreen_lists). A shader that fails
832    /// while running leaves the pixels as they were; `check` reports it.
833    pub(crate) fn apply(
834        &self,
835        img: &mut oxframe_core::Image,
836        to_px: Affine,
837        opacity: f32,
838        drawn: &[oxframe_core::Image],
839    ) {
840        let (w, h) = self.size;
841        if img.width == 0 || img.height == 0 || w <= 0.0 || h <= 0.0 {
842            return;
843        }
844        let snapshot = std::rc::Rc::new(RasterImage {
845            width: img.width,
846            height: img.height,
847            rgba: img.data.clone(),
848            id: 0,
849        });
850        let mut globals = self.globals.clone();
851        for name in &self.frame {
852            globals.inputs.push((name.clone(), snapshot.clone()));
853        }
854        for (input, pixels) in self.offscreen.iter().zip(drawn) {
855            globals.inputs.push((
856                input.name.clone(),
857                std::rc::Rc::new(crop(pixels, to_px, input.bounds)),
858            ));
859        }
860        let node_px = to_px * self.tf;
861        let Some(inv) = (node_px.determinant().abs() > 1e-12).then(|| node_px.inverse()) else {
862            return;
863        };
864        let r = node_px.transform_rect_bbox(kurbo::Rect::new(-w / 2.0, -h / 2.0, w / 2.0, h / 2.0));
865        let x0 = r.x0.floor().clamp(0.0, f64::from(img.width)) as usize;
866        let y0 = r.y0.floor().clamp(0.0, f64::from(img.height)) as usize;
867        let x1 = r.x1.ceil().clamp(0.0, f64::from(img.width)) as usize;
868        let y1 = r.y1.ceil().clamp(0.0, f64::from(img.height)) as usize;
869        if x1 <= x0 || y1 <= y0 {
870            return;
871        }
872        let cols = x1 - x0;
873        let total = cols * (y1 - y0);
874        let (sw, sh) = self.scene;
875        let tf = self.tf;
876        let width = img.width as usize;
877        let pixel = |i: usize| (x0 + i % cols, y0 + i / cols);
878        let local = |i: usize| {
879            let (x, y) = pixel(i);
880            inv * Point::new(x as f64 + 0.5, y as f64 + 0.5)
881        };
882        let src = &snapshot.rgba;
883        let data = &mut img.data;
884        let _ = run_tiles(
885            &self.program,
886            &globals,
887            total,
888            |i| {
889                let l = local(i);
890                let world = tf * l;
891                [
892                    (l.x + w / 2.0) / w,
893                    (l.y + h / 2.0) / h,
894                    l.x,
895                    l.y,
896                    world.x,
897                    world.y,
898                    world.x + sw / 2.0,
899                    world.y + sh / 2.0,
900                ]
901            },
902            |i, [r, g, b, a]| {
903                let l = local(i);
904                if l.x.abs() > w / 2.0 || l.y.abs() > h / 2.0 {
905                    return;
906                }
907                let (x, y) = pixel(i);
908                let o = (y * width + x) * 4;
909                let a = (a.clamp(0.0, 1.0) * opacity).clamp(0.0, 1.0);
910                let under = f32::from(src[o + 3]) / 255.0;
911                let out_a = a + under * (1.0 - a);
912                for (c, v) in [r, g, b].into_iter().enumerate() {
913                    let below = f32::from(src[o + c]) / 255.0;
914                    let mixed = if out_a > 0.0 {
915                        (v.clamp(0.0, 1.0) * a + below * under * (1.0 - a)) / out_a
916                    } else {
917                        0.0
918                    };
919                    data[o + c] = byte(mixed);
920                }
921                data[o + 3] = byte(out_a);
922            },
923        );
924    }
925}
926
927/// The part of `img` under `bounds` (scene coordinates), at least a pixel.
928fn crop(img: &oxframe_core::Image, to_px: Affine, bounds: kurbo::Rect) -> RasterImage {
929    let r = to_px.transform_rect_bbox(bounds);
930    let clampx = |v: f64| v.clamp(0.0, f64::from(img.width)) as u32;
931    let clampy = |v: f64| v.clamp(0.0, f64::from(img.height)) as u32;
932    let (x0, y0) = (clampx(r.x0.floor()), clampy(r.y0.floor()));
933    let (x1, y1) = (
934        clampx(r.x1.ceil()).max(x0 + 1),
935        clampy(r.y1.ceil()).max(y0 + 1),
936    );
937    let (w, h) = (x1 - x0, y1 - y0);
938    let mut rgba = vec![0u8; (w * h * 4) as usize];
939    for y in 0..h {
940        for x in 0..w {
941            let (sx, sy) = (x0 + x, y0 + y);
942            if sx < img.width && sy < img.height {
943                let s = ((sy * img.width + sx) * 4) as usize;
944                let d = ((y * w + x) * 4) as usize;
945                rgba[d..d + 4].copy_from_slice(&img.data[s..s + 4]);
946            }
947        }
948    }
949    RasterImage {
950        width: w,
951        height: h,
952        rgba,
953        id: 0,
954    }
955}
956
957/// True when a custom shader reads the frame, the backdrop or a node.
958pub(crate) fn reads_frame_input(n: &NodeData) -> bool {
959    matches!(&n.kind, NodeKind::Shader { preset: Preset::Custom, inputs, .. }
960        if inputs.iter().any(|(_, i)| matches!(i, crate::post::Input::Frame | crate::post::Input::Backdrop | crate::post::Input::Node(_))))
961}
962
963fn wave_uniforms(uniforms: &[Uniform]) -> (f64, f64, Color, Color) {
964    let mut amp = 0.05;
965    let mut speed = 1.0;
966    let mut deep = Color::BLACK;
967    let mut foam = Color::WHITE;
968    for u in uniforms {
969        match (&u.name[..], &u.value) {
970            ("amp", UniformValue::Num(v)) => amp = v.get(),
971            ("speed", UniformValue::Num(v)) => speed = v.get(),
972            ("deep", UniformValue::Color(v)) => deep = v.get(),
973            ("foam", UniformValue::Color(v)) => foam = v.get(),
974            _ => {}
975        }
976    }
977    (amp, speed, deep, foam)
978}
979
980fn load(code: &ShaderCode) -> Result<String, String> {
981    match code {
982        ShaderCode::Wgsl(text) => Ok(text.clone()),
983        ShaderCode::File(path) => read_file(path),
984    }
985}
986
987#[cfg(not(target_arch = "wasm32"))]
988fn read_file(path: &std::path::Path) -> Result<String, String> {
989    std::fs::read_to_string(path).map_err(|e| format!("{}: {e}", path.display()))
990}
991
992#[cfg(target_arch = "wasm32")]
993fn read_file(path: &std::path::Path) -> Result<String, String> {
994    let _ = path;
995    Err("a shader file is read by the native renderer".into())
996}
997
998fn label(n: &NodeData) -> String {
999    n.name
1000        .clone()
1001        .unwrap_or_else(|| format!("{}#{}", n.kind.name(), n.id.0))
1002}
1003
1004type Compiled = Result<std::rc::Rc<crate::wgsl_cpu::Program>, (&'static str, String)>;
1005
1006thread_local! {
1007    /// Parsed and validated shaders by the hash of their text: a file is
1008    /// read on every draw, but parsed once per change.
1009    static PROGRAMS: std::cell::RefCell<std::collections::HashMap<u64, Compiled>> =
1010        std::cell::RefCell::new(std::collections::HashMap::new());
1011}
1012
1013/// The shader for `text`, parsed and validated once.
1014fn compiled(text: &str, inputs: &[&str]) -> Compiled {
1015    let key = hash_bytes(format!("{text}\0{}", inputs.join(",")).as_bytes());
1016    if let Some(hit) = PROGRAMS.with(|p| p.borrow().get(&key).cloned()) {
1017        return hit;
1018    }
1019    let result = compile(text, inputs);
1020    PROGRAMS.with(|p| {
1021        let mut p = p.borrow_mut();
1022        if p.len() > 64 {
1023            p.clear();
1024        }
1025        p.insert(key, result.clone());
1026    });
1027    result
1028}
1029
1030fn compile(text: &str, inputs: &[&str]) -> Compiled {
1031    // After the user's text, so naga's line numbers still point into it.
1032    let samplers: String = inputs
1033        .iter()
1034        .map(|name| {
1035            format!("\nfn sample_{name}(uv: vec2<f32>) -> vec4<f32> {{ return vec4<f32>(0.0); }}\n")
1036        })
1037        .collect();
1038    let full = format!("{PRELUDE}{text}{samplers}{EPILOGUE}");
1039    let module = naga::front::wgsl::parse_str(&full).map_err(|e| {
1040        let mut hint = shift_lines(&e.emit_to_string(&full), prelude_lines());
1041        if let Some(name) = missing_sampler(&hint) {
1042            hint.push_str(&format!(
1043                "\nRead an image with `.input(\"{name}\", Input::file(\"…png\"))` on the shader."
1044            ));
1045        }
1046        ("shader_error", hint)
1047    })?;
1048    naga::valid::Validator::new(
1049        naga::valid::ValidationFlags::all(),
1050        naga::valid::Capabilities::all(),
1051    )
1052    .validate(&module)
1053    .map_err(|e| ("shader_error", shift_lines(&e.to_string(), prelude_lines())))?;
1054    crate::wgsl_cpu::Program::new(module)
1055        .map(std::rc::Rc::new)
1056        .ok_or(("shader_error", "the shader has no `fn shade`".to_string()))
1057}
1058
1059fn check_wgsl(
1060    text: &str,
1061    uniforms: &[Uniform],
1062    inputs: &[&str],
1063) -> Result<(), (&'static str, String)> {
1064    let program = compiled(text, inputs)?;
1065    let fields = params_of(program.module()).ok_or((
1066        "shader_error",
1067        "the shader has no `struct Params`".to_string(),
1068    ))?;
1069    match params(uniforms, &fields) {
1070        Ok(()) => Ok(()),
1071        Err(hint) => Err(("param_unknown", hint)),
1072    }
1073}
1074
1075/// The input name in naga's error for a `sample_<name>` that no input
1076/// declares.
1077fn missing_sampler(message: &str) -> Option<String> {
1078    let at = message.find("sample_")?;
1079    let name: String = message[at + 7..]
1080        .chars()
1081        .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
1082        .collect();
1083    (message.contains("no definition in scope") && !name.is_empty()).then_some(name)
1084}
1085
1086/// A WGSL identifier: a letter or `_`, then letters, digits and `_`.
1087fn is_ident(name: &str) -> bool {
1088    let mut chars = name.chars();
1089    chars
1090        .next()
1091        .is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
1092        && chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
1093}
1094
1095fn prelude_lines() -> usize {
1096    PRELUDE.lines().count()
1097}
1098
1099/// naga reports lines in the wrapped source. Subtract the prelude so the
1100/// number points at the user's text.
1101fn shift_lines(message: &str, by: usize) -> String {
1102    let chars: Vec<char> = message.chars().collect();
1103    let mut out = String::with_capacity(message.len());
1104    let mut i = 0;
1105    while i < chars.len() {
1106        if chars[i] == ':' && i + 1 < chars.len() && chars[i + 1].is_ascii_digit() {
1107            let start = i + 1;
1108            let mut end = start;
1109            while end < chars.len() && chars[end].is_ascii_digit() {
1110                end += 1;
1111            }
1112            let digits: String = chars[start..end].iter().collect();
1113            let line: usize = digits.parse().unwrap_or(0);
1114            out.push(':');
1115            out.push_str(&line.saturating_sub(by).to_string());
1116            i = end;
1117        } else {
1118            out.push(chars[i]);
1119            i += 1;
1120        }
1121    }
1122    out
1123}
1124
1125enum WgslTy {
1126    F32,
1127    Vec2,
1128    Vec4,
1129    Bool,
1130    Other,
1131}
1132
1133fn params_of(module: &naga::Module) -> Option<Vec<(String, WgslTy)>> {
1134    for (_, ty) in module.types.iter() {
1135        if ty.name.as_deref() != Some("Params") {
1136            continue;
1137        }
1138        let naga::TypeInner::Struct { members, .. } = &ty.inner else {
1139            continue;
1140        };
1141        return Some(
1142            members
1143                .iter()
1144                .map(|m| {
1145                    let name = m.name.clone().unwrap_or_default();
1146                    (name, wgsl_ty(&module.types[m.ty].inner))
1147                })
1148                .collect(),
1149        );
1150    }
1151    None
1152}
1153
1154fn wgsl_ty(inner: &naga::TypeInner) -> WgslTy {
1155    match inner {
1156        naga::TypeInner::Scalar(s) if s.kind == naga::ScalarKind::Float && s.width == 4 => {
1157            WgslTy::F32
1158        }
1159        naga::TypeInner::Scalar(s) if s.kind == naga::ScalarKind::Bool => WgslTy::Bool,
1160        naga::TypeInner::Vector { size, scalar }
1161            if scalar.kind == naga::ScalarKind::Float
1162                && scalar.width == 4
1163                && *size == naga::VectorSize::Bi =>
1164        {
1165            WgslTy::Vec2
1166        }
1167        naga::TypeInner::Vector { size, scalar }
1168            if scalar.kind == naga::ScalarKind::Float
1169                && scalar.width == 4
1170                && *size == naga::VectorSize::Quad =>
1171        {
1172            WgslTy::Vec4
1173        }
1174        _ => WgslTy::Other,
1175    }
1176}
1177
1178fn params(uniforms: &[Uniform], fields: &[(String, WgslTy)]) -> Result<(), String> {
1179    let names: Vec<&str> = fields.iter().map(|(n, _)| n.as_str()).collect();
1180    let list = names.join(", ");
1181    for u in uniforms {
1182        let Some((_, ty)) = fields.iter().find(|(n, _)| n == &u.name) else {
1183            let near = closest(&u.name, &names);
1184            let hint = match near {
1185                Some(n) => format!(
1186                    "no parameter `{}`. Did you mean `{n}`? Available: {list}.",
1187                    u.name
1188                ),
1189                None => format!("no parameter `{}`. Available: {list}.", u.name),
1190            };
1191            return Err(hint);
1192        };
1193        if !type_matches(&u.value, ty) {
1194            return Err(format!(
1195                "`{}` is {} in Params, but the scene sets {}.",
1196                u.name,
1197                ty_name(ty),
1198                value_name(&u.value)
1199            ));
1200        }
1201    }
1202    for (name, _) in fields {
1203        if !uniforms.iter().any(|u| u.name == *name) {
1204            return Err(format!(
1205                "Params.{name} has no value. Set it with `.uniform(\"{name}\", ...)`."
1206            ));
1207        }
1208    }
1209    Ok(())
1210}
1211
1212fn type_matches(value: &UniformValue, ty: &WgslTy) -> bool {
1213    matches!(
1214        (value, ty),
1215        (UniformValue::Num(_), WgslTy::F32)
1216            | (UniformValue::Vec2(_), WgslTy::Vec2)
1217            | (UniformValue::Color(_), WgslTy::Vec4)
1218            | (UniformValue::Bool(_), WgslTy::Bool)
1219    )
1220}
1221
1222fn ty_name(ty: &WgslTy) -> &'static str {
1223    match ty {
1224        WgslTy::F32 => "f32",
1225        WgslTy::Vec2 => "vec2<f32>",
1226        WgslTy::Vec4 => "vec4<f32>",
1227        WgslTy::Bool => "bool",
1228        WgslTy::Other => "another type",
1229    }
1230}
1231
1232fn value_name(value: &UniformValue) -> &'static str {
1233    match value {
1234        UniformValue::Num(_) => "a number",
1235        UniformValue::Vec2(_) => "a point",
1236        UniformValue::Color(_) => "a color",
1237        UniformValue::Bool(_) => "a bool",
1238    }
1239}
1240
1241fn closest<'a>(name: &str, options: &[&'a str]) -> Option<&'a str> {
1242    options
1243        .iter()
1244        .copied()
1245        .filter(|o| edit(name, o) <= 3 && *o != name)
1246        .min_by_key(|o| edit(name, o))
1247}
1248
1249fn edit(a: &str, b: &str) -> usize {
1250    let a: Vec<char> = a.chars().collect();
1251    let b: Vec<char> = b.chars().collect();
1252    let mut prev: Vec<usize> = (0..=b.len()).collect();
1253    let mut cur = vec![0; b.len() + 1];
1254    for (i, ca) in a.iter().enumerate() {
1255        cur[0] = i + 1;
1256        for (j, cb) in b.iter().enumerate() {
1257            let cost = usize::from(ca != cb);
1258            cur[j + 1] = (prev[j + 1] + 1).min(cur[j] + 1).min(prev[j] + cost);
1259        }
1260        std::mem::swap(&mut prev, &mut cur);
1261    }
1262    prev[b.len()]
1263}
1264
1265fn smoothstep(e0: f64, e1: f64, x: f64) -> f64 {
1266    let t = if (e1 - e0).abs() < 1e-9 {
1267        if x < e0 { 0.0 } else { 1.0 }
1268    } else {
1269        ((x - e0) / (e1 - e0)).clamp(0.0, 1.0)
1270    };
1271    t * t * (3.0 - 2.0 * t)
1272}
1273
1274fn mix(foam: Color, deep: Color, t: f64, opacity: f64) -> Color {
1275    let t = t as f32;
1276    let a = opacity.clamp(0.0, 1.0) as f32;
1277    Color {
1278        r: foam.r + (deep.r - foam.r) * t,
1279        g: foam.g + (deep.g - foam.g) * t,
1280        b: foam.b + (deep.b - foam.b) * t,
1281        a: (foam.a + (deep.a - foam.a) * t) * a,
1282    }
1283}
1284
1285fn byte(c: f32) -> u8 {
1286    (c.clamp(0.0, 1.0) * 255.0).round() as u8
1287}
1288
1289fn hash_bytes(bytes: &[u8]) -> u64 {
1290    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
1291    for &b in bytes {
1292        h = (h ^ u64::from(b)).wrapping_mul(0x0100_0000_01b3);
1293    }
1294    h
1295}
1296
1297#[cfg(test)]
1298mod tests {
1299    use super::*;
1300    use crate::builder::NodeBuilder;
1301    use crate::{SceneGraph, View};
1302    use std::rc::Rc;
1303
1304    fn paint(f: impl FnOnce(View)) -> Rc<SceneGraph> {
1305        let graph = SceneGraph::new(320, 180);
1306        graph.enter(|| f(View::current()));
1307        graph
1308    }
1309
1310    fn image_of(graph: &SceneGraph) -> Rc<RasterImage> {
1311        let items = graph.draw_list();
1312        for item in &items {
1313            if let Primitive::Image { image, .. } = &item.prim {
1314                return image.clone();
1315            }
1316        }
1317        panic!("waves did not paint, diags: {:?}", graph.shader_diags());
1318    }
1319
1320    #[test]
1321    fn waves_put_foam_above_deep_water() {
1322        let graph = paint(|view| {
1323            view.add(crate::waves().width(200.0).height(80.0));
1324        });
1325        let image = image_of(&graph);
1326        let px = |x: u32, y: u32| {
1327            let i = (y * image.width + x) as usize * 4;
1328            image.rgba[i + 2]
1329        };
1330        let mid = image.width / 2;
1331        assert!(
1332            px(mid, 2) > px(mid, image.height - 3),
1333            "the crest is the lighter foam"
1334        );
1335        assert!(
1336            graph.shader_diags().is_empty(),
1337            "{:?}",
1338            graph.shader_diags()
1339        );
1340    }
1341
1342    #[test]
1343    fn custom_wgsl_runs_like_the_hand_written_preset() {
1344        // The waves preset's own WGSL, run by the interpreter, against the
1345        // preset's Rust code: the same picture up to f32 rounding.
1346        let custom = paint(|view| {
1347            view.add(
1348                crate::Shader::source(waves_wgsl())
1349                    .uniform("amp", 0.2)
1350                    .uniform("speed", 1.0)
1351                    .uniform_color("deep", Color::rgba(0.1, 0.3, 0.5, 1.0))
1352                    .uniform_color("foam", Color::rgba(0.9, 0.95, 1.0, 1.0))
1353                    .width(200.0)
1354                    .height(80.0),
1355            );
1356        });
1357        let preset = paint(|view| {
1358            view.add(
1359                crate::waves()
1360                    .amp(0.2)
1361                    .deep(Color::rgba(0.1, 0.3, 0.5, 1.0))
1362                    .foam(Color::rgba(0.9, 0.95, 1.0, 1.0))
1363                    .width(200.0)
1364                    .height(80.0),
1365            );
1366        });
1367        for t in [0.0, 1.3] {
1368            custom.set_time(t);
1369            preset.set_time(t);
1370            let (a, b) = (image_of(&custom), image_of(&preset));
1371            assert_eq!((a.width, a.height), (b.width, b.height));
1372            let worst = a
1373                .rgba
1374                .iter()
1375                .zip(&b.rgba)
1376                .map(|(x, y)| x.abs_diff(*y))
1377                .max()
1378                .unwrap_or(0);
1379            assert!(worst <= 2, "differs by {worst}/255 at t={t}");
1380        }
1381        assert!(
1382            custom.shader_diags().is_empty(),
1383            "{:?}",
1384            custom.shader_diags()
1385        );
1386    }
1387
1388    fn two_pixels() -> crate::Input {
1389        crate::Input::Image(Rc::new(RasterImage {
1390            width: 2,
1391            height: 1,
1392            rgba: vec![255, 0, 0, 255, 0, 0, 255, 255],
1393            id: 9_001,
1394        }))
1395    }
1396
1397    const SAMPLE_PHOTO: &str = "struct Params { k: f32 }
1398@group(1) @binding(0) var<uniform> p: Params;
1399fn shade(i: ShadeIn) -> vec4<f32> {
1400    return sample_photo(i.uv);
1401}";
1402
1403    #[test]
1404    fn a_custom_shade_reads_an_image_input() {
1405        let graph = paint(|view| {
1406            view.add(
1407                crate::Shader::source(SAMPLE_PHOTO)
1408                    .uniform("k", 0.0)
1409                    .input("photo", two_pixels())
1410                    .width(40.0)
1411                    .height(4.0),
1412            );
1413        });
1414        let img = image_of(&graph);
1415        let px = |x: u32| {
1416            let o = ((img.height / 2 * img.width + x) * 4) as usize;
1417            [img.rgba[o], img.rgba[o + 1], img.rgba[o + 2]]
1418        };
1419        assert_eq!(px(0), [255, 0, 0]);
1420        assert_eq!(px(img.width - 1), [0, 0, 255]);
1421        let mid = px(img.width / 2);
1422        assert!(mid[0] > 60 && mid[2] > 60, "bilinear between: {mid:?}");
1423        assert!(rules(&graph).is_empty(), "{:?}", graph.shader_diags());
1424    }
1425
1426    #[test]
1427    fn a_custom_shade_with_a_frame_input_reads_what_was_drawn() {
1428        // A mirror over the right half of a 320x180 frame whose left half
1429        // is red and right half blue: the right half turns red.
1430        let graph = paint(|view| {
1431            view.add(
1432                crate::Shader::source(
1433                    "struct Params { k: f32 }
1434@group(1) @binding(0) var<uniform> p: Params;
1435fn shade(i: ShadeIn) -> vec4<f32> {
1436    let s = i.screen / ox.resolution;
1437    return sample_scene(vec2<f32>(1.0 - s.x, s.y));
1438}",
1439                )
1440                .uniform("k", 0.0)
1441                .input("scene", crate::Input::frame())
1442                .width(160.0)
1443                .height(180.0)
1444                .x(80.0),
1445            );
1446        });
1447        assert!(reads_frame(&graph));
1448        assert!(rules(&graph).is_empty(), "{:?}", graph.shader_diags());
1449        let shade = graph
1450            .draw_list()
1451            .iter()
1452            .find_map(|i| match &i.prim {
1453                Primitive::ReadFrame(s) => Some(s.clone()),
1454                _ => None,
1455            })
1456            .expect("a frame pass");
1457        assert_eq!(shade.op.name(), "custom");
1458        let mut img = oxframe_core::Image {
1459            width: 320,
1460            height: 180,
1461            data: (0..320 * 180)
1462                .flat_map(|i| {
1463                    if i % 320 < 160 {
1464                        [255, 0, 0, 255]
1465                    } else {
1466                        [0, 0, 255, 255]
1467                    }
1468                })
1469                .collect(),
1470        };
1471        let to_px = Affine::translate((160.0, 90.0));
1472        crate::post::apply(&shade, &mut img, to_px);
1473        let px = |x: usize| &img.data[(90 * 320 + x) * 4..(90 * 320 + x) * 4 + 3];
1474        assert_eq!(px(10), [255, 0, 0], "left half is not touched");
1475        assert_eq!(px(300), [255, 0, 0], "right half mirrors the left");
1476    }
1477
1478    #[test]
1479    fn an_offscreen_node_is_not_drawn_and_check_skips_it() {
1480        let graph = paint(|view| {
1481            let card = view.add(crate::Rect::new().size(20.0).name("card").offscreen());
1482            view.add(
1483                crate::Shader::source(SAMPLE_PHOTO.replace("photo", "card"))
1484                    .uniform("k", 0.0)
1485                    .input("card", crate::Input::node(card))
1486                    .width(40.0)
1487                    .height(40.0),
1488            );
1489        });
1490        let list = graph.draw_list();
1491        assert!(
1492            list.iter().all(|i| i.prim.is_marker()),
1493            "only the shader's pass is on screen"
1494        );
1495        let shade = list
1496            .iter()
1497            .find_map(|i| match &i.prim {
1498                Primitive::ReadFrame(s) => Some(s),
1499                _ => None,
1500            })
1501            .expect("a node input makes a frame pass");
1502        assert_eq!(shade.offscreen_lists().len(), 1);
1503        assert!(!shade.offscreen_lists()[0].is_empty());
1504        let card = graph
1505            .inspect_nodes()
1506            .into_iter()
1507            .find(|n| n.name == "card")
1508            .expect("card");
1509        assert_eq!(card.opacity, 0.0);
1510        assert!(graph.has_post());
1511    }
1512
1513    #[test]
1514    fn sampling_an_input_that_is_not_declared_says_how_to_add_it() {
1515        let graph = paint(|view| {
1516            view.add(
1517                crate::Shader::source(SAMPLE_PHOTO)
1518                    .uniform("k", 0.0)
1519                    .width(4.0)
1520                    .height(4.0),
1521            );
1522        });
1523        let _ = graph.draw_list();
1524        let diags = graph.shader_diags();
1525        let d = diags
1526            .iter()
1527            .find(|d| d.rule == "shader_error")
1528            .expect("shader_error");
1529        assert!(d.hint.contains(".input(\"photo\""), "{}", d.hint);
1530    }
1531
1532    #[test]
1533    fn an_input_name_must_be_a_wgsl_name() {
1534        let graph = paint(|view| {
1535            view.add(
1536                crate::Shader::source(SAMPLE_PHOTO)
1537                    .uniform("k", 0.0)
1538                    .input("my photo", two_pixels())
1539                    .width(4.0)
1540                    .height(4.0),
1541            );
1542        });
1543        assert!(rules(&graph).contains(&"param_unknown"));
1544    }
1545
1546    #[test]
1547    fn a_shader_that_fails_while_running_is_a_shader_error() {
1548        let graph = paint(|view| {
1549            view.add(
1550                crate::Shader::source(
1551                    "struct Params { k: f32 }
1552@group(1) @binding(0) var<uniform> p: Params;
1553fn shade(i: ShadeIn) -> vec4<f32> {
1554    var x = 0.0;
1555    loop { x += p.k; }
1556    return vec4<f32>(x);
1557}",
1558                )
1559                .uniform("k", 1.0)
1560                .width(4.0)
1561                .height(4.0),
1562            );
1563        });
1564        let items = graph.draw_list();
1565        assert!(
1566            !items
1567                .iter()
1568                .any(|i| matches!(i.prim, Primitive::Image { .. }))
1569        );
1570        let diags = graph.shader_diags();
1571        assert!(
1572            diags
1573                .iter()
1574                .any(|d| d.rule == "shader_error" && d.hint.contains("loops")),
1575            "{diags:?}"
1576        );
1577    }
1578
1579    #[test]
1580    fn the_crest_moves_along_the_wave() {
1581        let graph = paint(|view| {
1582            view.add(crate::waves().amp(0.2).width(200.0).height(80.0));
1583        });
1584        let a = image_of(&graph).rgba.clone();
1585        graph.set_time(1.0);
1586        let b = image_of(&graph).rgba.clone();
1587        assert_ne!(a, b);
1588    }
1589
1590    #[test]
1591    fn a_misspelled_uniform_names_the_real_one() {
1592        let graph = paint(|view| {
1593            view.add(
1594                crate::Shader::source(waves_wgsl())
1595                    .uniform("ampp", 0.1)
1596                    .uniform("speed", 1.0)
1597                    .uniform_color("deep", Color::BLACK)
1598                    .uniform_color("foam", Color::WHITE)
1599                    .width(40.0)
1600                    .height(20.0)
1601                    .name("sea"),
1602            );
1603        });
1604        let _ = graph.draw_list();
1605        let diag = graph
1606            .shader_diags()
1607            .into_iter()
1608            .find(|d| d.rule == "param_unknown")
1609            .expect("param_unknown");
1610        assert_eq!(diag.node, "sea");
1611        assert!(diag.hint.contains("amp"), "{diag:?}");
1612    }
1613
1614    #[test]
1615    fn a_broken_shader_is_an_error() {
1616        let graph = paint(|view| {
1617            view.add(
1618                crate::Shader::source("this is not wgsl")
1619                    .width(40.0)
1620                    .height(20.0),
1621            );
1622        });
1623        let _ = graph.draw_list();
1624        assert!(
1625            graph
1626                .shader_diags()
1627                .iter()
1628                .any(|d| d.rule == "shader_error"),
1629            "{:?}",
1630            graph.shader_diags()
1631        );
1632    }
1633
1634    #[cfg(not(target_arch = "wasm32"))]
1635    #[test]
1636    fn editing_the_file_is_seen_on_the_next_draw() {
1637        let dir = std::env::temp_dir().join(format!("oxframe-shader-{}", std::process::id()));
1638        std::fs::create_dir_all(&dir).unwrap();
1639        let path = dir.join("sea.wgsl");
1640        std::fs::write(&path, "this is not wgsl").unwrap();
1641        let graph = paint(|view| {
1642            view.add(crate::Shader::file(&path).width(40.0).height(20.0));
1643        });
1644        let _ = graph.draw_list();
1645        assert!(
1646            graph
1647                .shader_diags()
1648                .iter()
1649                .any(|d| d.rule == "shader_error"),
1650            "{:?}",
1651            graph.shader_diags()
1652        );
1653        std::fs::write(&path, waves_wgsl()).unwrap();
1654        let _ = graph.draw_list();
1655        let diags = graph.shader_diags();
1656        assert!(diags.iter().all(|d| d.rule != "shader_error"), "{diags:?}");
1657        let _ = std::fs::remove_dir_all(&dir);
1658    }
1659
1660    fn rules(graph: &SceneGraph) -> Vec<&'static str> {
1661        let _ = graph.draw_list();
1662        graph.shader_diags().iter().map(|d| d.rule).collect()
1663    }
1664
1665    fn reads_frame(graph: &SceneGraph) -> bool {
1666        graph
1667            .draw_list()
1668            .iter()
1669            .any(|i| matches!(i.prim, Primitive::ReadFrame(_)))
1670    }
1671
1672    #[test]
1673    fn lens_drops_read_the_frame() {
1674        let graph = paint(|view| {
1675            view.add(
1676                crate::lens_drops()
1677                    .fullscreen()
1678                    .input("scene", crate::Input::frame()),
1679            );
1680        });
1681        assert!(graph.has_post());
1682        assert!(reads_frame(&graph));
1683        assert!(rules(&graph).is_empty(), "{:?}", graph.shader_diags());
1684        let passes = crate::pass_graph(&graph.draw_list());
1685        assert!(
1686            passes
1687                .iter()
1688                .any(|p| p.kind == crate::PassKind::Effect && p.shader == Some("lens_drops"))
1689        );
1690    }
1691
1692    #[test]
1693    fn a_misspelled_input_or_parameter_names_the_real_one() {
1694        let graph = paint(|view| {
1695            view.add(
1696                crate::lens_drops()
1697                    .input("scen", crate::Input::frame())
1698                    .name("lens"),
1699            );
1700        });
1701        let _ = graph.draw_list();
1702        let diag = &graph.shader_diags()[0];
1703        assert_eq!((diag.rule, diag.node.as_str()), ("param_unknown", "lens"));
1704        assert!(diag.hint.contains("`scene`"), "{diag:?}");
1705
1706        let graph = paint(|view| {
1707            view.add(crate::lens_drops().uniform("wetnes", 1.0));
1708        });
1709        let _ = graph.draw_list();
1710        let diag = &graph.shader_diags()[0];
1711        assert!(diag.hint.contains("`wetness`"), "{diag:?}");
1712
1713        let graph = paint(|view| {
1714            view.add(crate::waves().input("scene", crate::Input::frame()));
1715        });
1716        assert_eq!(rules(&graph), ["param_unknown"]);
1717    }
1718
1719    #[test]
1720    fn a_lens_inside_a_fading_group_reads_the_layer() {
1721        let graph = paint(|view| {
1722            let group = view.add(crate::Group::new().opacity(0.5));
1723            group.add(crate::Rect::new().size(20.0));
1724            group.add(crate::lens_drops().fullscreen());
1725        });
1726        assert!(rules(&graph).is_empty(), "{:?}", graph.shader_diags());
1727        assert!(reads_frame(&graph));
1728    }
1729
1730    #[test]
1731    fn an_effect_inside_an_effect_runs_both() {
1732        let graph = paint(|view| {
1733            let outer = view.add(crate::Group::new());
1734            let inner = outer.add(crate::Rect::new().size(20.0));
1735            outer.effect(crate::ripple());
1736            inner.effect(crate::ripple());
1737        });
1738        assert!(rules(&graph).is_empty(), "{:?}", graph.shader_diags());
1739        let items = graph.draw_list();
1740        let begins = items
1741            .iter()
1742            .filter(|i| matches!(i.prim, Primitive::BeginEffect(_)))
1743            .count();
1744        assert_eq!(begins, 2);
1745        let passes = crate::pass_graph(&items);
1746        assert!(
1747            passes
1748                .iter()
1749                .any(|p| p.kind == crate::PassKind::Effect && p.shader == Some("ripple"))
1750        );
1751    }
1752
1753    #[test]
1754    fn effect_signals_animate_the_marker() {
1755        let graph = SceneGraph::new(320, 180);
1756        let effect = graph.enter(|| {
1757            View::current()
1758                .add(crate::Rect::new().size(40.0))
1759                .effect(crate::ripple().strength(0.0))
1760        });
1761        let strength = |graph: &SceneGraph| {
1762            graph.draw_list().iter().find_map(|i| match &i.prim {
1763                Primitive::BeginEffect(s) => match s.op {
1764                    crate::post::ShadeOp::Ripple { strength, .. } => Some(strength),
1765                    _ => None,
1766                },
1767                _ => None,
1768            })
1769        };
1770        assert_eq!(strength(&graph), Some(0.0));
1771        graph.enter(|| effect.strength().set(0.7));
1772        assert_eq!(strength(&graph), Some(0.7));
1773    }
1774
1775    #[test]
1776    fn a_slow_raster_is_reported() {
1777        let graph = paint(|view| {
1778            view.add(crate::waves().width(180.0).height(80.0).name("sea"));
1779        });
1780        set_shader_budget_ms(0.0);
1781        let _ = graph.draw_list();
1782        set_shader_budget_ms(250.0);
1783        assert!(
1784            graph.shader_diags().iter().any(|d| d.rule == "shader_slow"),
1785            "{:?}",
1786            graph.shader_diags()
1787        );
1788    }
1789}
1790
1791#[cfg(test)]
1792mod speed {
1793    use super::*;
1794
1795    #[test]
1796    #[ignore = "timing, run by hand in release"]
1797    fn custom_waves_full_frame() {
1798        let graph = crate::SceneGraph::new(1920, 1080);
1799        graph.enter(|| {
1800            crate::View::current().add(
1801                crate::Shader::source(waves_wgsl())
1802                    .uniform("amp", 0.2)
1803                    .uniform("speed", 1.0)
1804                    .uniform_color("deep", Color::BLACK)
1805                    .uniform_color("foam", Color::WHITE)
1806                    .width(1920.0)
1807                    .height(1080.0),
1808            );
1809        });
1810        let t = std::time::Instant::now();
1811        let _ = graph.draw_list();
1812        eprintln!("custom waves 1920x1080: {:?}", t.elapsed());
1813        let graph = crate::SceneGraph::new(1920, 1080);
1814        graph.enter(|| {
1815            crate::View::current().add(crate::waves().width(1920.0).height(1080.0));
1816        });
1817        let t = std::time::Instant::now();
1818        let _ = graph.draw_list();
1819        eprintln!("preset waves 1920x1080: {:?}", t.elapsed());
1820    }
1821}