start voronoi

src/desktop.rs

@@ -17,16 +17,17 @@ impl DesktopEntries {
     pub fn count(&self) -> usize {
         self.entries.len()
     }
+    pub fn colors(&self) -> impl Iterator<Item = Oklab> + ExactSizeIterator {
+        self.entries.iter().map(|entry| entry.avg_icon_color)
+    }
     pub fn find_entry(&self, color: Oklab) -> Option<&DesktopEntry> {
-        self.entries.iter().min_by(|x, y| {
+        self.entries
-            f32::total_cmp(
+            .iter()
-                &diff_color(x.avg_icon_color, color),
+            .min_by_key(|x| OrdFloat(diff_color(x.avg_icon_color, color)))
-                &diff_color(y.avg_icon_color, color),
-            )
-        })
     }
 }
 
+// keep it in sync with the gpu implementation
 fn diff_color(icon: Oklab, color: Oklab) -> f32 {
     icon.distance_squared(color)
 }
@@ -98,9 +99,14 @@ pub(crate) fn find_desktop_files() -> Result<DesktopEntries> {
         .wrap_err_with(|| format!("{}", base.display()))?;
     }
 
-    Ok(DesktopEntries {
+    let mut entries = results.into_values().collect::<Vec<_>>();
-        entries: results.into_values().collect(),
+
-    })
+    entries.sort_by_key(|entry| {
+        let (l, a, b) = entry.avg_icon_color.into_components();
+        (OrdFloat(a), OrdFloat(b), OrdFloat(l))
+    });
+
+    Ok(DesktopEntries { entries })
 }
 
 fn average_color(image: &image::DynamicImage) -> palette::Oklab {
@@ -131,3 +137,22 @@ fn average_color(image: &image::DynamicImage) -> palette::Oklab {
         b: total_b / count,
     }
 }
+
+struct OrdFloat(f32);
+
+impl PartialEq for OrdFloat {
+    fn eq(&self, other: &Self) -> bool {
+        self.cmp(other).is_eq()
+    }
+}
+impl Eq for OrdFloat {}
+impl PartialOrd for OrdFloat {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(&other))
+    }
+}
+impl Ord for OrdFloat {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.0.total_cmp(&other.0)
+    }
+}

src/gpu.rs

@@ -1,6 +1,7 @@
-use std::ptr::NonNull;
+use std::{mem::offset_of, ptr::NonNull};
 
 use eyre::{Context, Result};
+use palette::Oklab;
 use raw_window_handle::{
     RawDisplayHandle, RawWindowHandle, WaylandDisplayHandle, WaylandWindowHandle,
 };
@@ -9,27 +10,42 @@ use wgpu::util::DeviceExt;
 
 pub struct AppGpuState {
     instance: wgpu::Instance,
-    adapter: wgpu::Adapter,
     device: wgpu::Device,
     queue: wgpu::Queue,
     render_pipeline: wgpu::RenderPipeline,
     screen_size_bind_group_layout: wgpu::BindGroupLayout,
+    desktop_colors_bind_group: wgpu::BindGroup,
 }
 
 pub struct SurfaceGpuState {
     surface: wgpu::Surface<'static>,
-    screen_size_buffer: wgpu::Buffer,
+    width: u32,
+    height: u32,
+    input_buffer: wgpu::Buffer,
     screen_size_bind_group: wgpu::BindGroup,
 }
 
 #[repr(C)]
 #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
-struct ScreenSizeUniform {
+struct InputUniform {
     size: [f32; 2], // width, height
+    voronoi_progress: f32,
+    _pad: f32,
+}
+
+#[repr(C)]
+#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
+struct DesktopColorsStorage {
+    l: f32,
+    a: f32,
+    b: f32,
+    _pad: f32,
 }
 
 impl AppGpuState {
-    pub fn new() -> Result<Self> {
+    pub fn new(
+        desktop_colors: impl IntoIterator<Item = Oklab> + ExactSizeIterator,
+    ) -> Result<Self> {
         let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
 
         let adapter =
@@ -54,10 +70,28 @@ impl AppGpuState {
                     count: None,
                 }],
             });
+        let desktop_colors_bind_group_layout =
+            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+                label: Some("desktop_colors_bind_group_layout"),
+                entries: &[wgpu::BindGroupLayoutEntry {
+                    binding: 0,
+                    visibility: wgpu::ShaderStages::FRAGMENT,
+                    ty: wgpu::BindingType::Buffer {
+                        ty: wgpu::BufferBindingType::Storage { read_only: true },
+                        has_dynamic_offset: false,
+                        min_binding_size: None,
+                    },
+                    count: None,
+                }],
+            });
+
         let render_pipeline_layout =
             device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                 label: Some("Render Pipeline Layout"),
-                bind_group_layouts: &[&screen_size_bind_group_layout],
+                bind_group_layouts: &[
+                    &screen_size_bind_group_layout,
+                    &desktop_colors_bind_group_layout,
+                ],
                 immediate_size: 0,
             });
 
@@ -73,11 +107,7 @@ impl AppGpuState {
             fragment: Some(wgpu::FragmentState {
                 module: &shader,
                 entry_point: Some("fs_main"),
-                targets: &[Some(wgpu::ColorTargetState {
+                targets: &[Some(wgpu::TextureFormat::Bgra8UnormSrgb.into())],
-                    format: wgpu::TextureFormat::Rgba8UnormSrgb,
-                    blend: Some(wgpu::BlendState::REPLACE),
-                    write_mask: wgpu::ColorWrites::ALL,
-                })],
                 compilation_options: wgpu::PipelineCompilationOptions::default(),
             }),
             primitive: wgpu::PrimitiveState {
@@ -99,13 +129,42 @@ impl AppGpuState {
             cache: None,
         });
 
+        let desktop_colors = desktop_colors
+            .into_iter()
+            .map(|color| DesktopColorsStorage {
+                l: color.l,
+                a: color.a,
+                b: color.b,
+                _pad: 0.0,
+            })
+            .collect::<Vec<_>>();
+
+        let desktop_colors_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
+            label: Some("desktop_colors_buffer"),
+            contents: bytemuck::cast_slice::<DesktopColorsStorage, u8>(&desktop_colors),
+            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
+        });
+
+        let desktop_colors_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            layout: &desktop_colors_bind_group_layout,
+            entries: &[wgpu::BindGroupEntry {
+                binding: 0,
+                resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
+                    buffer: &desktop_colors_buffer,
+                    offset: 0,
+                    size: None,
+                }),
+            }],
+            label: Some("desktop_colors_bind_group"),
+        });
+
         Ok(Self {
             instance,
-            adapter,
             device,
             queue,
             render_pipeline,
             screen_size_bind_group_layout,
+            desktop_colors_bind_group,
         })
     }
 }
@@ -135,7 +194,11 @@ impl SurfaceGpuState {
                 .device
                 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                     label: Some("Screen Size Uniform Buffer"),
-                    contents: bytemuck::bytes_of(&ScreenSizeUniform { size: [0.0, 0.0] }),
+                    contents: bytemuck::bytes_of(&InputUniform {
+                        size: [0.0, 0.0],
+                        voronoi_progress: 0.0,
+                        _pad: 0.0,
+                    }),
                     usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                 });
 
@@ -157,38 +220,63 @@ impl SurfaceGpuState {
 
         Ok(Self {
             surface,
-            screen_size_buffer,
+            input_buffer: screen_size_buffer,
             screen_size_bind_group,
+            width: 0,
+            height: 0,
         })
     }
 
-    pub fn resize(&self, gpu_state: &AppGpuState, width: u32, height: u32) {
+    pub fn resize(&mut self, gpu_state: &AppGpuState, width: u32, height: u32) {
+        self.width = width;
+        self.height = height;
+
         gpu_state.queue.write_buffer(
-            &self.screen_size_buffer,
+            &self.input_buffer,
             0,
-            bytemuck::bytes_of(&ScreenSizeUniform {
+            bytemuck::bytes_of(&InputUniform {
                 size: [width as f32, height as f32],
+                voronoi_progress: 0.0,
+                _pad: 0.0,
             }),
         );
 
-        let cap = self.surface.get_capabilities(&gpu_state.adapter);
+        self.configure(gpu_state);
+    }
+
+    fn configure(&self, gpu_state: &AppGpuState) {
         let surface_config = wgpu::SurfaceConfiguration {
             usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
-            format: cap.formats[0],
+            format: wgpu::TextureFormat::Bgra8UnormSrgb,
-            view_formats: vec![cap.formats[0]],
+            view_formats: vec![wgpu::TextureFormat::Bgra8UnormSrgb],
             alpha_mode: wgpu::CompositeAlphaMode::Auto,
-            width,
+            width: self.width,
-            height,
+            height: self.height,
             desired_maximum_frame_latency: 2,
             // Wayland is inherently a mailbox system.
             present_mode: wgpu::PresentMode::Mailbox,
         };
         self.surface.configure(&gpu_state.device, &surface_config);
+    }
+
+    pub fn set_voronoi_progress(&self, gpu_state: &AppGpuState, voronoi_progress: f32) {
+        gpu_state.queue.write_buffer(
+            &self.input_buffer,
+            offset_of!(InputUniform, voronoi_progress) as u64,
+            bytemuck::bytes_of(&voronoi_progress),
+        );
+    }
+
+    pub fn draw(&self, gpu_state: &AppGpuState) {
+        let surface_texture = match self.surface.get_current_texture() {
+            Ok(texture) => texture,
+            Err(wgpu::SurfaceError::Outdated | wgpu::SurfaceError::Lost) => {
+                self.configure(gpu_state);
+                self.surface.get_current_texture().unwrap()
+            }
+            Err(e) => panic!("failed to acquire next swapchain texture: {e}"),
+        };
 
-        let surface_texture = self
-            .surface
-            .get_current_texture()
-            .expect("failed to acquire next swapchain texture");
         let texture_view: wgpu::TextureView = surface_texture
             .texture
             .create_view(&wgpu::TextureViewDescriptor::default());
@@ -219,7 +307,8 @@ impl SurfaceGpuState {
             });
 
             render_pass.set_pipeline(&gpu_state.render_pipeline);
-            render_pass.set_bind_group(0, Some(&self.screen_size_bind_group), &[]);
+            render_pass.set_bind_group(0, &self.screen_size_bind_group, &[]);
+            render_pass.set_bind_group(1, &gpu_state.desktop_colors_bind_group, &[]);
             render_pass.draw(0..6, 0..1);
         }
 

src/main.rs

@@ -9,7 +9,7 @@ use std::{
 use eyre::{Context, Result, bail, eyre};
 use freedesktop_file_parser::EntryType;
 use log::{error, info, warn};
-use palette::{FromColor, IntoColor, Oklab};
+use palette::Oklab;
 use smithay_client_toolkit::{
     compositor::{CompositorHandler, CompositorState},
     output::{OutputHandler, OutputState},
@@ -18,7 +18,7 @@ use smithay_client_toolkit::{
     registry_handlers,
     seat::{
         SeatHandler, SeatState,
-        pointer::{BTN_LEFT, PointerEventKind, PointerHandler},
+        pointer::{BTN_LEFT, BTN_RIGHT, PointerEventKind, PointerHandler},
     },
     shell::{
         WaylandSurface,
@@ -70,7 +70,7 @@ fn main() -> Result<()> {
         shm: Shm::bind(&globals, qh).wrap_err("failed to bind shm")?,
         seat_state: SeatState::new(&globals, qh),
 
-        gpu: AppGpuState::new()?,
+        gpu: AppGpuState::new(desktop_files.colors())?,
 
         desktop_files,
         pointers: HashMap::new(),
@@ -114,6 +114,7 @@ struct OutputSurface {
     layer_surface: LayerSurface,
     width: u32,
     height: u32,
+    voronoi_progress: f32,
 }
 
 impl ProvidesRegistryState for App {
@@ -165,6 +166,7 @@ impl OutputHandler for App {
                     layer_surface,
                     width: 0,
                     height: 0,
+                    voronoi_progress: 0.0,
                 });
             }
             Err(err) => error!(
@@ -233,6 +235,7 @@ impl CompositorHandler for App {
         _surface: &wayland_client::protocol::wl_surface::WlSurface,
         _time: u32,
     ) {
+        dbg!("yeet");
     }
 
     fn surface_enter(
@@ -293,20 +296,20 @@ impl LayerShellHandler for App {
         surface.height = height;
 
         surface.gpu.resize(&self.gpu, width, height);
+        surface.gpu.draw(&self.gpu);
     }
 }
 
 // keep it in sync with the gpu implementation
-fn color_for_pixel(x: u32, y: u32, width: u32, height: u32) -> palette::Srgb<u8> {
+fn color_for_pixel(x: u32, y: u32, width: u32, height: u32) -> Oklab {
     let xf = x as f32 / width as f32;
     let yf = y as f32 / height as f32;
 
-    palette::Srgb::from_color(palette::Oklab {
+    palette::Oklab {
         l: 0.7,
         a: xf * 0.8 - 0.4,
-        b: yf * 0.8 - 0.4,
+        b: yf * 0.7 - 0.4,
-    })
+    }
-    .into_format::<u8>()
 }
 
 impl ShmHandler for App {
@@ -373,27 +376,25 @@ impl PointerHandler for App {
         events: &[smithay_client_toolkit::seat::pointer::PointerEvent],
     ) {
         for event in events {
-            if let PointerEventKind::Release {
-                button: BTN_LEFT, ..
-            } = event.kind
-            {
             let Some(surface) = self
                 .layer_surfaces
-                    .iter()
+                .iter_mut()
                 .find(|surface| *surface.layer_surface.wl_surface() == event.surface)
             else {
                 return;
             };
 
-                let srgb = color_for_pixel(
+            match event.kind {
+                PointerEventKind::Release {
+                    button: BTN_LEFT, ..
+                } => {
+                    let oklab = color_for_pixel(
                         event.position.0 as u32,
                         event.position.1 as u32,
                         surface.width,
                         surface.height,
                     );
 
-                let oklab: Oklab = srgb.into_format::<f32>().into_color();
-
                     let best_match = self.desktop_files.find_entry(oklab);
 
                     if let Some(best_match) = best_match
@@ -415,6 +416,29 @@ impl PointerHandler for App {
                         }
                     }
                 }
+                PointerEventKind::Press {
+                    button: BTN_RIGHT, ..
+                } => {
+                    surface.voronoi_progress = 1.0;
+
+                    surface
+                        .gpu
+                        .set_voronoi_progress(&self.gpu, surface.voronoi_progress);
+                    surface.gpu.draw(&self.gpu);
+                }
+                PointerEventKind::Release {
+                    button: BTN_RIGHT, ..
+                }
+                | PointerEventKind::Leave { .. } => {
+                    surface.voronoi_progress = 0.0;
+
+                    surface
+                        .gpu
+                        .set_voronoi_progress(&self.gpu, surface.voronoi_progress);
+                    surface.gpu.draw(&self.gpu);
+                }
+                _ => (),
+            }
         }
     }
 }

src/shader.wgsl

@@ -1,14 +1,19 @@
-struct Screen {
+struct Input {
     size: vec2<f32>,
+    voronoi_progress: f32,
 };
 
 @group(0) @binding(0)
-var<uniform> screen: Screen;
+var<uniform> input: Input;
+
+@group(1) @binding(0)
+var<storage, read> desktop_colors: array<vec4f>;
 
 @vertex
 fn vs_main(
     @builtin(vertex_index) in_vertex_index: u32,
 ) -> @builtin(position) vec4<f32> {
+    // full-screen quad
     var pos = array<vec2f, 6>(
         vec2(-1.0, 1.0),
         vec2(-1.0, -1.0),
@@ -23,20 +28,42 @@ fn vs_main(
 
 @fragment
 fn fs_main(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
-    var posf = pos.xy / screen.size;
+    var posf = pos.xy / input.size;
 
-    // keep it in sync with the cpu implementation
+    var color = vec3<f32>(
-    var color = oklab_to_linear_srgb(vec3<f32>(
         0.7,
         posf.x * 0.8 - 0.4,
-        posf.y * 0.8 - 0.4,
+        posf.y * 0.7 - 0.4,
-    ));
+    );
 
-    return vec4<f32>(color.x, color.y, color.z, 1.0);
+    var best = vec3f(0.0, 0.0, 0.0);
+    var best_score = 1000000000000.0;
+    for (var i: u32 = 0; i < arrayLength(&desktop_colors); i++) {
+        var elem = desktop_colors[i].xyz;
+        var score = diff_colors(elem, color);
+        if (score < best_score) {
+            best = elem;
+            best_score = score;
+        }
+    }
+    var voronoi_color = best;
+     
+    color = mix(color, voronoi_color, input.voronoi_progress);
+
+    // keep it in sync with the cpu implementation
+    var srgbcolor = oklab_to_linear_srgb(color);
+
+    return vec4<f32>(srgbcolor.x, srgbcolor.y, srgbcolor.z, 1.0);
 }
 
+// keep it in sync with the cpu implementation
+fn diff_colors(oklab_a: vec3f, oklab_b: vec3f) -> f32 {
+    var diff = oklab_a - oklab_b;
+    var diff_sq = diff * diff;
+    return diff_sq.x + diff_sq.y + diff_sq.z;
+}
 
-fn oklab_to_linear_srgb(oklab: vec3<f32>) -> vec3<f32> {
+fn oklab_to_linear_srgb(oklab: vec3f) -> vec3f {
     let l_ =  0.2158037573 * oklab.z + (0.3963377774 * oklab.y + oklab.x);
     let m_ = -0.0638541728 * oklab.z + (-0.1055613458 * oklab.y + oklab.x);
     let s_ = -1.2914855480 * oklab.z + (-0.0894841775 * oklab.y + oklab.x);
@@ -49,4 +76,3 @@ fn oklab_to_linear_srgb(oklab: vec3<f32>) -> vec3<f32> {
         1.7076147010 * s + (-0.0041960863 * l + -0.7034186147 * m),
     );
 }
-