perf: editing theme doesn't require a complete reredenring

src/fractal/renderer.rs

@@ -17,6 +17,37 @@ use eframe::egui_wgpu::{self, wgpu};
 /// bounds the iteration count. 128k points * 8 bytes = 1 MiB.
 pub const MAX_REF_POINTS: usize = 1 << 17;
 
+/// Format of the intermediate iteration-data texture holding, per pixel,
+/// `(ci, DE factor, interior fraction)`. 32-bit float keeps the smooth iteration
+/// count precise at deep zoom. Color-renderable and read with nearest sampling
+/// (iteration data must never be linearly filtered across escape boundaries), so
+/// no `float32-filterable` feature is needed.
+const DATA_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba32Float;
+
+/// True when the two uniforms differ in any field the iteration pass depends on
+/// (i.e. anything except the palette / colour scale / offset).
+fn geom_differs(a: &Uniforms, b: &Uniforms) -> bool {
+    a.span != b.span
+        || a.max_iter != b.max_iter
+        || a.ref_len != b.ref_len
+        || a.bailout_sq != b.bailout_sq
+        || a.is_julia != b.is_julia
+        || a.aa_level != b.aa_level
+        || a.kind != b.kind
+        || a.power != b.power
+        || a.dc_offset != b.dc_offset
+        || a.phoenix_p != b.phoenix_p
+        || a.de_coloring != b.de_coloring
+}
+
+/// True when the two uniforms differ in a colour-only field (remappable by the
+/// cheap colourise pass without re-iterating).
+fn color_differs(a: &Uniforms, b: &Uniforms) -> bool {
+    a.color_offset != b.color_offset
+        || a.color_scale != b.color_scale
+        || a.palette_id != b.palette_id
+}
+
 /// GPU-side view + coloring parameters. Layout must match `Uniforms` in the
 /// WGSL shader; total size is a multiple of 16 bytes for uniform-buffer rules.
 #[repr(C)]
@@ -53,26 +84,44 @@ pub struct Uniforms {
     pub _pad: [u32; 3],
 }
 
-/// Offscreen texture the fractal is rendered into, plus the bind group used to
+/// Offscreen textures for the two-pass render, recreated whenever the widget's
-/// blit it. Recreated whenever the widget's pixel size changes.
+/// pixel size changes:
+/// * `data_view` — the iteration pass's output (see [`DATA_FORMAT`]).
+/// * `color_view` — the colourise pass's output; the blit source.
+/// plus the bind groups that read them.
 struct CacheTarget {
-    view: wgpu::TextureView,
+    data_view: wgpu::TextureView,
+    color_view: wgpu::TextureView,
+    /// Colourise pass input: uniforms + the data texture.
+    colorize_bind_group: wgpu::BindGroup,
+    /// Blit pass input: the colour texture + sampler.
     blit_bind_group: wgpu::BindGroup,
     width: u32,
     height: u32,
 }
 
-/// State the cache texture was last rendered with. If the next frame's inputs
+/// What the iteration-data texture was last computed with. If the next frame's
-/// match this, the cache is still valid and the fractal shader is skipped.
+/// geometry inputs match, iteration is skipped and only colour may be redone.
-struct RenderedState {
+struct IterState {
     uniforms: Uniforms,
     generation: u64,
     width: u32,
     height: u32,
 }
 
+/// What the colour texture was last computed with. If the next frame's colour
+/// inputs (and size) match and iteration did not re-run, colourise is skipped.
+struct ColorState {
+    uniforms: Uniforms,
+    width: u32,
+    height: u32,
+}
+
 pub struct FractalRenderer {
-    pipeline: wgpu::RenderPipeline,
+    /// Iteration pass: perturbation iterate → data texture (`fs_data`).
+    iterate_pipeline: wgpu::RenderPipeline,
+    /// Combined iterate + colour in one pass (`fs_color`), used only by export.
+    export_pipeline: wgpu::RenderPipeline,
     bind_group_layout: wgpu::BindGroupLayout,
     uniform_buffer: wgpu::Buffer,
     ref_buffer: wgpu::Buffer,
@@ -81,14 +130,20 @@ pub struct FractalRenderer {
     /// Generation of the reference orbit currently uploaded to `ref_buffer`.
     uploaded_generation: u64,
 
-    /// Blit pipeline + resources that copy the cache texture to egui's surface.
+    /// Colourise pass: data texture → colour texture (palette mapping).
+    colorize_pipeline: wgpu::RenderPipeline,
+    colorize_bind_group_layout: wgpu::BindGroupLayout,
+
+    /// Blit pipeline + resources that copy the colour texture to egui's surface.
     blit_pipeline: wgpu::RenderPipeline,
     blit_bind_group_layout: wgpu::BindGroupLayout,
     blit_sampler: wgpu::Sampler,
-    /// The offscreen cache; `None` until the first frame sizes it.
+    /// The offscreen textures; `None` until the first frame sizes them.
     cache: Option<CacheTarget>,
-    /// What the cache currently holds; `None` forces a re-render.
+    /// What the data texture holds; `None` forces re-iteration.
-    rendered: Option<RenderedState>,
+    iterated: Option<IterState>,
+    /// What the colour texture holds; `None` forces a recolour.
+    colored: Option<ColorState>,
 }
 
 impl FractalRenderer {
@@ -159,8 +214,9 @@ impl FractalRenderer {
             immediate_size: 0,
         });
 
-        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+        // Iteration pass: perturbation iterate → data texture (color-independent).
-            label: Some("fractal pipeline"),
+        let iterate_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+            label: Some("fractal iterate pipeline"),
             layout: Some(&pipeline_layout),
             vertex: wgpu::VertexState {
                 module: &shader,
@@ -170,6 +226,97 @@ impl FractalRenderer {
             },
             fragment: Some(wgpu::FragmentState {
                 module: &shader,
+                entry_point: Some("fs_data"),
+                targets: &[Some(wgpu::ColorTargetState {
+                    format: DATA_FORMAT,
+                    blend: None,
+                    write_mask: wgpu::ColorWrites::ALL,
+                })],
+                compilation_options: Default::default(),
+            }),
+            primitive: wgpu::PrimitiveState::default(),
+            depth_stencil: None,
+            multisample: wgpu::MultisampleState::default(),
+            multiview_mask: None,
+            cache: None,
+        });
+
+        // Combined iterate + colour in one pass — for PNG export only.
+        let export_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+            label: Some("fractal export pipeline"),
+            layout: Some(&pipeline_layout),
+            vertex: wgpu::VertexState {
+                module: &shader,
+                entry_point: Some("vs_main"),
+                buffers: &[],
+                compilation_options: Default::default(),
+            },
+            fragment: Some(wgpu::FragmentState {
+                module: &shader,
+                entry_point: Some("fs_color"),
+                targets: &[Some(wgpu::ColorTargetState {
+                    format: target_format,
+                    blend: None,
+                    write_mask: wgpu::ColorWrites::ALL,
+                })],
+                compilation_options: Default::default(),
+            }),
+            primitive: wgpu::PrimitiveState::default(),
+            depth_stencil: None,
+            multisample: wgpu::MultisampleState::default(),
+            multiview_mask: None,
+            cache: None,
+        });
+
+        // Colourise pass: data texture + colour uniforms → colour texture.
+        let colorize_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+            label: Some("colorize"),
+            source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/colorize.wgsl").into()),
+        });
+        let colorize_bind_group_layout =
+            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                label: Some("colorize bind group layout"),
+                entries: &[
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 0,
+                        visibility: wgpu::ShaderStages::FRAGMENT,
+                        ty: wgpu::BindingType::Buffer {
+                            ty: wgpu::BufferBindingType::Uniform,
+                            has_dynamic_offset: false,
+                            min_binding_size: None,
+                        },
+                        count: None,
+                    },
+                    wgpu::BindGroupLayoutEntry {
+                        binding: 1,
+                        visibility: wgpu::ShaderStages::FRAGMENT,
+                        ty: wgpu::BindingType::Texture {
+                            // Nearest only: iteration data must not be filtered.
+                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
+                            view_dimension: wgpu::TextureViewDimension::D2,
+                            multisampled: false,
+                        },
+                        count: None,
+                    },
+                ],
+            });
+        let colorize_pipeline_layout =
+            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
+                label: Some("colorize pipeline layout"),
+                bind_group_layouts: &[Some(&colorize_bind_group_layout)],
+                immediate_size: 0,
+            });
+        let colorize_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
+            label: Some("colorize pipeline"),
+            layout: Some(&colorize_pipeline_layout),
+            vertex: wgpu::VertexState {
+                module: &colorize_shader,
+                entry_point: Some("vs_main"),
+                buffers: &[],
+                compilation_options: Default::default(),
+            },
+            fragment: Some(wgpu::FragmentState {
+                module: &colorize_shader,
                 entry_point: Some("fs_main"),
                 targets: &[Some(wgpu::ColorTargetState {
                     format: target_format,
@@ -254,18 +401,22 @@ impl FractalRenderer {
         });
 
         Self {
-            pipeline,
+            iterate_pipeline,
+            export_pipeline,
             bind_group_layout,
             uniform_buffer,
             ref_buffer,
             bind_group,
             target_format,
             uploaded_generation: u64::MAX,
+            colorize_pipeline,
+            colorize_bind_group_layout,
             blit_pipeline,
             blit_bind_group_layout,
             blit_sampler,
             cache: None,
-            rendered: None,
+            iterated: None,
+            colored: None,
         }
     }
 
@@ -279,13 +430,29 @@ impl FractalRenderer {
             return;
         }
 
-        let texture = device.create_texture(&wgpu::TextureDescriptor {
+        let extent = wgpu::Extent3d {
-            label: Some("fractal cache"),
-            size: wgpu::Extent3d {
             width,
             height,
             depth_or_array_layers: 1,
-            },
+        };
+
+        // Iteration-data texture (color-independent escape data).
+        let data_texture = device.create_texture(&wgpu::TextureDescriptor {
+            label: Some("fractal data"),
+            size: extent,
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: wgpu::TextureDimension::D2,
+            format: DATA_FORMAT,
+            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
+            view_formats: &[],
+        });
+        let data_view = data_texture.create_view(&wgpu::TextureViewDescriptor::default());
+
+        // Colour texture (colourise output; blit source).
+        let color_texture = device.create_texture(&wgpu::TextureDescriptor {
+            label: Some("fractal color cache"),
+            size: extent,
             mip_level_count: 1,
             sample_count: 1,
             dimension: wgpu::TextureDimension::D2,
@@ -293,7 +460,22 @@ impl FractalRenderer {
             usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
             view_formats: &[],
         });
-        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
+        let color_view = color_texture.create_view(&wgpu::TextureViewDescriptor::default());
+
+        let colorize_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: Some("colorize bind group"),
+            layout: &self.colorize_bind_group_layout,
+            entries: &[
+                wgpu::BindGroupEntry {
+                    binding: 0,
+                    resource: self.uniform_buffer.as_entire_binding(),
+                },
+                wgpu::BindGroupEntry {
+                    binding: 1,
+                    resource: wgpu::BindingResource::TextureView(&data_view),
+                },
+            ],
+        });
 
         let blit_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
             label: Some("blit bind group"),
@@ -301,7 +483,7 @@ impl FractalRenderer {
             entries: &[
                 wgpu::BindGroupEntry {
                     binding: 0,
-                    resource: wgpu::BindingResource::TextureView(&view),
+                    resource: wgpu::BindingResource::TextureView(&color_view),
                 },
                 wgpu::BindGroupEntry {
                     binding: 1,
@@ -311,13 +493,16 @@ impl FractalRenderer {
         });
 
         self.cache = Some(CacheTarget {
-            view,
+            data_view,
+            color_view,
+            colorize_bind_group,
             blit_bind_group,
             width,
             height,
         });
-        // New texture → old render is gone.
+        // New textures → old renders are gone.
-        self.rendered = None;
+        self.iterated = None;
+        self.colored = None;
     }
 
     /// Handles needed to build a standalone [`ExportRender`] off the UI thread:
@@ -325,7 +510,7 @@ impl FractalRenderer {
     /// format. Cloned so the caller can drop the render-state lock before use.
     pub fn export_handles(&self) -> (wgpu::RenderPipeline, wgpu::BindGroupLayout, wgpu::TextureFormat) {
         (
-            self.pipeline.clone(),
+            self.export_pipeline.clone(),
             self.bind_group_layout.clone(),
             self.target_format,
         )
@@ -582,8 +767,9 @@ pub fn encode_png_with_progress(
 
 /// A per-frame paint callback. Carries this frame's uniforms plus a reference to
 /// the current reference orbit (cheap `Arc` clone). The orbit is only re-uploaded
-/// to the GPU when its `generation` changes, and the fractal is only re-rendered
+/// when its `generation` changes; the expensive iteration pass re-runs only when
-/// into the cache when the uniforms, generation, or `size_px` change.
+/// a geometry input changes, and colour-only changes re-run just the cheap
+/// colourise pass (see `prepare`).
 pub struct FractalCallback {
     pub uniforms: Uniforms,
     pub reference: Arc<Vec<[f32; 2]>>,
@@ -619,28 +805,35 @@ impl egui_wgpu::CallbackTrait for FractalCallback {
             renderer.uploaded_generation = self.generation;
         }
 
-        // Re-render the cache only when what it depends on changed.
+        // Iteration (expensive) re-runs only when the geometry inputs change;
-        let dirty = renderer.rendered.as_ref().is_none_or(|r| {
+        // colourise (cheap) re-runs when it did, or when only a colour changed —
+        // so palette / colour-scale / offset tweaks (e.g. colour cycling) skip
+        // the perturbation entirely.
+        let iter_dirty = renderer.iterated.as_ref().is_none_or(|r| {
             r.generation != self.generation
                 || r.width != width
                 || r.height != height
-                || bytemuck::bytes_of(&r.uniforms) != bytemuck::bytes_of(&self.uniforms)
+                || geom_differs(&r.uniforms, &self.uniforms)
         });
-        if !dirty {
+        let color_dirty = iter_dirty
-            return Vec::new();
+            || renderer.colored.as_ref().is_none_or(|c| {
+                c.width != width || c.height != height || color_differs(&c.uniforms, &self.uniforms)
+            });
+
+        if !color_dirty {
+            return Vec::new(); // cache still valid; paint() just blits it
         }
 
-        queue.write_buffer(
+        // Both passes read the uniform buffer; refresh it once.
-            &renderer.uniform_buffer,
+        queue.write_buffer(&renderer.uniform_buffer, 0, bytemuck::bytes_of(&self.uniforms));
-            0,
-            bytemuck::bytes_of(&self.uniforms),
-        );
 
         if let Some(cache) = &renderer.cache {
+            if iter_dirty {
+                // Iteration pass: perturbation iterate → data texture.
                 let mut pass = egui_encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
-                label: Some("fractal cache pass"),
+                    label: Some("fractal iterate pass"),
                     color_attachments: &[Some(wgpu::RenderPassColorAttachment {
-                    view: &cache.view,
+                        view: &cache.data_view,
                         depth_slice: None,
                         resolve_target: None,
                         ops: wgpu::Operations {
@@ -653,17 +846,46 @@ impl egui_wgpu::CallbackTrait for FractalCallback {
                     occlusion_query_set: None,
                     multiview_mask: None,
                 });
-            pass.set_pipeline(&renderer.pipeline);
+                pass.set_pipeline(&renderer.iterate_pipeline);
                 pass.set_bind_group(0, &renderer.bind_group, &[]);
                 pass.draw(0..3, 0..1);
             }
 
-        renderer.rendered = Some(RenderedState {
+            // Colourise pass: data texture → colour texture.
+            let mut pass = egui_encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+                label: Some("fractal colorize pass"),
+                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
+                    view: &cache.color_view,
+                    depth_slice: None,
+                    resolve_target: None,
+                    ops: wgpu::Operations {
+                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
+                        store: wgpu::StoreOp::Store,
+                    },
+                })],
+                depth_stencil_attachment: None,
+                timestamp_writes: None,
+                occlusion_query_set: None,
+                multiview_mask: None,
+            });
+            pass.set_pipeline(&renderer.colorize_pipeline);
+            pass.set_bind_group(0, &cache.colorize_bind_group, &[]);
+            pass.draw(0..3, 0..1);
+        }
+
+        if iter_dirty {
+            renderer.iterated = Some(IterState {
                 uniforms: self.uniforms,
                 generation: self.generation,
                 width,
                 height,
             });
+        }
+        renderer.colored = Some(ColorState {
+            uniforms: self.uniforms,
+            width,
+            height,
+        });
         Vec::new()
     }
 

src/shaders/colorize.wgsl

@@ -0,0 +1,78 @@
+// Colourise pass: map the iteration pass's per-pixel escape data (from
+// `mandelbrot.wgsl`'s `fs_data`) through the palette. This is the only
+// color-dependent step, so changing the palette / colour scale / offset (e.g.
+// colour cycling) re-runs just this cheap pass — the expensive perturbation
+// iteration in the data texture is reused untouched.
+//
+// The data texture holds, per texel: R = ci (palette parameter), G = DE
+// darkening factor, B = interior fraction (for boundary anti-aliasing). It is
+// the same resolution as this pass's target, so we read it with `textureLoad`
+// at the fragment's integer pixel coordinate (nearest — iteration data must not
+// be linearly filtered across escape boundaries).
+
+// Must match `Uniforms` in mandelbrot.wgsl / the Rust `Uniforms` struct.
+struct Uniforms {
+    span: vec2<f32>,
+    max_iter: u32,
+    ref_len: u32,
+    color_offset: f32,
+    color_scale: f32,
+    bailout_sq: f32,
+    is_julia: u32,
+    palette_id: u32,
+    aa_level: u32,
+    kind: u32,
+    power: u32,
+    dc_offset: vec2<f32>,
+    phoenix_p: vec2<f32>,
+    de_coloring: u32,
+};
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var data_tex: texture_2d<f32>;
+
+// Smooth cyclic palettes (Inigo Quilez cosine palettes). Must match the palette
+// in mandelbrot.wgsl.
+fn palette(id: u32, t: f32) -> vec3<f32> {
+    if (id == 4u) {
+        return vec3<f32>(t, t, t); // grayscale
+    }
+    let a = vec3<f32>(0.5, 0.5, 0.5);
+    let b = vec3<f32>(0.5, 0.5, 0.5);
+    var c = vec3<f32>(1.0, 1.0, 1.0);
+    var d = vec3<f32>(0.00, 0.10, 0.20); // 0: amber / blue
+    if (id == 1u) {
+        d = vec3<f32>(0.00, 0.33, 0.67); // rainbow
+    } else if (id == 2u) {
+        d = vec3<f32>(0.30, 0.20, 0.20); // warm ember
+    } else if (id == 3u) {
+        c = vec3<f32>(1.0, 1.0, 0.5);
+        d = vec3<f32>(0.80, 0.90, 0.30); // lime / magenta
+    }
+    return a + b * cos(6.28318530718 * (c * t + d));
+}
+
+@vertex
+fn vs_main(@builtin(vertex_index) idx: u32) -> @builtin(position) vec4<f32> {
+    var verts = array<vec2<f32>, 3>(
+        vec2<f32>(-1.0, -1.0),
+        vec2<f32>(3.0, -1.0),
+        vec2<f32>(-1.0, 3.0),
+    );
+    return vec4<f32>(verts[idx], 0.0, 1.0);
+}
+
+@fragment
+fn fs_main(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
+    let d = textureLoad(data_tex, vec2<i32>(i32(pos.x), i32(pos.y)), 0);
+    let ci = d.r;
+    let de = d.g;
+    let interior_frac = d.b;
+
+    let t = fract(ci * u.color_scale + u.color_offset);
+    var col = palette(u.palette_id, t) * de;
+    // Anti-alias the set boundary: fade toward black by the fraction of the
+    // pixel's sub-samples that landed in the interior.
+    col = col * (1.0 - interior_frac);
+    return vec4<f32>(col, 1.0);
+}

src/shaders/mandelbrot.wgsl

@@ -193,12 +193,21 @@ fn palette(id: u32, t: f32) -> vec3<f32> {
     return a + b * cos(6.28318530718 * (c * t + d));
 }
 
-// Perturbation iterate + color a single sample. `offset` is the per-pixel
+// Escape data for one sample: `ci` is the (color-independent) palette parameter,
-// offset in complex units. For Mandelbrot it is the c-plane offset added every
+// `de` the distance-estimate darkening factor in [0,1], `escaped` false for the
-// step (delta starts at 0); for Julia it is the z-plane offset that seeds the
+// interior of the set. Splitting iteration from coloring lets a colour change be
-// initial delta (c is fixed, so nothing is added per step). Interior pixels
+// remapped cheaply (see the colourise pass) without re-iterating.
-// return black.
+struct Sample {
-fn shade(offset: vec2<f32>, px: f32) -> vec3<f32> {
+    ci: f32,
+    de: f32,
+    escaped: bool,
+};
+
+// Perturbation iterate a single sample. `offset` is the per-pixel offset in
+// complex units. For Mandelbrot it is the c-plane offset added every step (delta
+// starts at 0); for Julia it is the z-plane offset that seeds the initial delta
+// (c is fixed, so nothing is added per step).
+fn iterate_sample(offset: vec2<f32>, px: f32) -> Sample {
     let z0 = ref_orbit[0]; // reference start (0 for Mandelbrot, center for Julia)
 
     var step_add = offset;
@@ -282,7 +291,7 @@ fn shade(offset: vec2<f32>, px: f32) -> vec3<f32> {
     }
 
     if (!escaped) {
-        return vec3<f32>(0.0, 0.0, 0.0); // interior of the set
+        return Sample(0.0, 1.0, false); // interior of the set
     }
 
     let z2 = dot(z, z);
@@ -295,9 +304,8 @@ fn shade(offset: vec2<f32>, px: f32) -> vec3<f32> {
     // sqrt compresses the huge iteration counts of deep zooms so the palette
     // varies smoothly instead of aliasing into speckle.
     let ci = sqrt(max(smooth_i, 0.0));
-    let t = fract(ci * u.color_scale + u.color_offset);
-    var col = palette(u.palette_id, t);
 
+    var de = 1.0;
     if (u.de_coloring != 0u) {
         // Exterior distance estimate (complex-plane units): |z|·ln|z| / |dz|.
         // Divided by the pixel footprint it becomes a distance in pixels; we
@@ -306,37 +314,77 @@ fn shade(offset: vec2<f32>, px: f32) -> vec3<f32> {
         // boundary simply reads as dark, which is the correct limit.
         let zmag = sqrt(max(z2, 1.0));
         let dzmag = sqrt(max(dot(dz, dz), 1e-20));
-        let de = zmag * log(zmag) / dzmag;
+        let d = zmag * log(zmag) / dzmag;
-        let de_px = de / max(px, 1e-30);
+        de = clamp(d / max(px, 1e-30), 0.0, 1.0);
-        col = col * clamp(de_px, 0.0, 1.0);
     }
-    return col;
+    return Sample(ci, de, true);
 }
 
-@fragment
+// Map a sample's escape data through the palette (+ DE darkening). This is the
-fn fs_main(in: VsOut) -> @location(0) vec4<f32> {
+// only color-dependent step, so it can be redone without re-iterating. Interior
-    let base = in.centered * u.span + u.dc_offset;
+// samples are black.
+fn color_sample(s: Sample) -> vec3<f32> {
+    if (!s.escaped) {
+        return vec3<f32>(0.0, 0.0, 0.0);
+    }
+    let t = fract(s.ci * u.color_scale + u.color_offset);
+    return palette(u.palette_id, t) * s.de;
+}
 
-    // Screen-space complex-units-per-pixel. Derivatives must be evaluated in
+// Iteration pass: write per-pixel escape data (color-independent) so a colour
-    // uniform control flow, so take them here; used to place sub-pixel AA
+// change is remapped by the cheap colourise pass without re-iterating.
-    // samples and to convert the distance estimate into pixels.
+//   R = ci (palette parameter), G = DE factor, B = interior fraction (for AA).
+// AA is grid-supersampled here; the interior fraction lets the colourise pass
+// anti-alias the set boundary (blend toward black) after the fact.
+@fragment
+fn fs_data(in: VsOut) -> @location(0) vec4<f32> {
+    let base = in.centered * u.span + u.dc_offset;
     let dx = dpdx(base);
     let dy = dpdy(base);
-    let px = length(abs(dx) + abs(dy)); // ~ complex units per pixel (footprint)
+    let px = length(abs(dx) + abs(dy));
 
     let aa = max(u.aa_level, 1u);
-    if (aa <= 1u) {
-        return vec4<f32>(shade(base, px), 1.0);
-    }
-
-    var acc = vec3<f32>(0.0, 0.0, 0.0);
     let inv = 1.0 / f32(aa);
+    var ci_sum = 0.0;
+    var de_sum = 0.0;
+    var escaped_n = 0u;
     for (var sy: u32 = 0u; sy < aa; sy = sy + 1u) {
         for (var sx: u32 = 0u; sx < aa; sx = sx + 1u) {
-            // Sample centers evenly spread across the pixel, jitter in (-0.5, 0.5).
             let jx = (f32(sx) + 0.5) * inv - 0.5;
             let jy = (f32(sy) + 0.5) * inv - 0.5;
-            acc = acc + shade(base + jx * dx + jy * dy, px);
+            let s = iterate_sample(base + jx * dx + jy * dy, px);
+            if (s.escaped) {
+                ci_sum = ci_sum + s.ci;
+                de_sum = de_sum + s.de;
+                escaped_n = escaped_n + 1u;
+            }
+        }
+    }
+    let total = f32(aa * aa);
+    let ci_avg = select(0.0, ci_sum / f32(escaped_n), escaped_n > 0u);
+    let de_avg = select(1.0, de_sum / f32(escaped_n), escaped_n > 0u);
+    let interior_frac = 1.0 - f32(escaped_n) / total;
+    return vec4<f32>(ci_avg, de_avg, interior_frac, 1.0);
+}
+
+// Combined iterate + colour in a single pass, for PNG export (which never needs
+// incremental recolouring). The interactive path uses fs_data + the colourise
+// pass so colour changes skip iteration.
+@fragment
+fn fs_color(in: VsOut) -> @location(0) vec4<f32> {
+    let base = in.centered * u.span + u.dc_offset;
+    let dx = dpdx(base);
+    let dy = dpdy(base);
+    let px = length(abs(dx) + abs(dy));
+
+    let aa = max(u.aa_level, 1u);
+    let inv = 1.0 / f32(aa);
+    var acc = vec3<f32>(0.0, 0.0, 0.0);
+    for (var sy: u32 = 0u; sy < aa; sy = sy + 1u) {
+        for (var sx: u32 = 0u; sx < aa; sx = sx + 1u) {
+            let jx = (f32(sx) + 0.5) * inv - 0.5;
+            let jy = (f32(sy) + 0.5) * inv - 0.5;
+            acc = acc + color_sample(iterate_sample(base + jx * dx + jy * dy, px));
         }
     }
     return vec4<f32>(acc / f32(aa * aa), 1.0);

tests/shader_valid.rs

@@ -25,6 +25,11 @@ fn mandelbrot_shader_is_valid() {
     );
 }
 
+#[test]
+fn colorize_shader_is_valid() {
+    validate("colorize.wgsl", include_str!("../src/shaders/colorize.wgsl"));
+}
+
 #[test]
 fn blit_shader_is_valid() {
     validate("blit.wgsl", include_str!("../src/shaders/blit.wgsl"));