stuff

9 files changed, 220 insertions, 104 deletions

diff --git a/vgcodec/Cargo.lock b/vgcodec/Cargo.lock
index 06a117c..4c7eacf 100644
--- a/vgcodec/Cargo.lock
+++ b/vgcodec/Cargo.lock
@@ -113,6 +113,20 @@ name = "bytemuck"
 version = "1.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "aaa3a8d9a1ca92e282c96a32d6511b695d7d994d1d102ba85d279f9b2756947f"
+dependencies = [
+ "bytemuck_derive",
+]
+
+[[package]]
+name = "bytemuck_derive"
+version = "1.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5fe233b960f12f8007e3db2d136e3cb1c291bfd7396e384ee76025fc1a3932b4"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
 
 [[package]]
 name = "byteorder"
diff --git a/vgcodec/Cargo.toml b/vgcodec/Cargo.toml
index 5b15036..6680e56 100644
--- a/vgcodec/Cargo.toml
+++ b/vgcodec/Cargo.toml
@@ -8,6 +8,6 @@ wgpu = "0.14.0"
 log = "0.4.17"
 env_logger = "0.9.3"
 pollster = "0.2.5"
-bytemuck = "1.12.3"
+bytemuck = { version = "1.12.3", features = ["derive"] }
 futures-intrusive = "0.5.0"
 image = "0.24.5"
diff --git a/vgcodec/src/app.rs b/vgcodec/src/app.rs
index 0c063a5..b51284c 100644
--- a/vgcodec/src/app.rs
+++ b/vgcodec/src/app.rs
@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use wgpu::{Adapter, Device, Instance, Queue};
+use wgpu::{Adapter, Device, Extent3d, ImageCopyTexture, Instance, Origin3d, Queue, Texture};
 
 pub struct App {
     pub instance: Instance,
@@ -34,4 +34,26 @@ impl App {
             queue,
         })
     }
+
+    pub fn copy_texture(&self, source: &Texture, destination: &Texture, size: Extent3d) {
+        let mut encoder = self
+            .device
+            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
+        encoder.copy_texture_to_texture(
+            ImageCopyTexture {
+                aspect: wgpu::TextureAspect::All,
+                mip_level: 0,
+                origin: Origin3d::ZERO,
+                texture: source,
+            },
+            ImageCopyTexture {
+                aspect: wgpu::TextureAspect::All,
+                mip_level: 0,
+                origin: Origin3d::ZERO,
+                texture: destination,
+            },
+            size,
+        );
+        self.queue.submit(Some(encoder.finish()));
+    }
 }
diff --git a/vgcodec/src/approximate.rs b/vgcodec/src/approximate.rs
new file mode 100644
index 0000000..0192973
--- /dev/null
+++ b/vgcodec/src/approximate.rs
@@ -0,0 +1,119 @@
+use std::sync::Arc;
+
+use image::EncodableLayout;
+use log::info;
+use wgpu::{Extent3d, ImageCopyTexture, Origin3d, Texture, TextureUsages};
+
+use crate::{
+    app::App,
+    diff::Differ,
+    export::Exporter,
+    helper::write_texture,
+    paint::{PaintUniforms, Painter},
+};
+
+pub struct Approximator {
+    app: Arc<App>,
+    size: Extent3d,
+    differ: Differ,
+    painter: Painter,
+    exporter: Exporter,
+
+    tex_approx: Texture,
+    tex_savestate: Texture,
+    tex_target: Texture,
+}
+
+impl Approximator {
+    pub fn new(app: &Arc<App>, tex_target: Texture, size: Extent3d) -> Self {
+        let App { device, queue, .. } = app.as_ref();
+        let tex_approx = device.create_texture(&wgpu::TextureDescriptor {
+            size,
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: wgpu::TextureDimension::D2,
+            format: wgpu::TextureFormat::Rgba8Unorm,
+            usage: TextureUsages::COPY_DST
+                | TextureUsages::TEXTURE_BINDING
+                | TextureUsages::STORAGE_BINDING
+                | TextureUsages::COPY_SRC,
+            label: None,
+        });
+        let tex_savestate = device.create_texture(&wgpu::TextureDescriptor {
+            size,
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: wgpu::TextureDimension::D2,
+            format: wgpu::TextureFormat::Rgba8Unorm,
+            usage: TextureUsages::COPY_DST | TextureUsages::COPY_SRC,
+            label: None,
+        });
+
+        let img_initial = image::open("a/initial.png").unwrap().into_rgba8();
+        write_texture(queue, &tex_approx, img_initial.as_bytes(), size);
+
+        let differ = Differ::new(&app, size, &tex_approx, &tex_target);
+        let painter = Painter::new(&app, size, &tex_approx);
+        let exporter = Exporter::new(&app, size);
+
+        Approximator {
+            app: app.clone(),
+            size,
+            differ,
+            tex_savestate,
+            tex_approx,
+            tex_target,
+            painter,
+            exporter,
+        }
+    }
+
+    pub fn save(&self) {
+        self.app
+            .copy_texture(&self.tex_approx, &self.tex_savestate, self.size);
+    }
+    pub fn restore(&self) {
+        self.app
+            .copy_texture(&self.tex_savestate, &self.tex_approx, self.size);
+    }
+
+    pub async fn run(&mut self) {
+        let mut current_diff = self.differ.run().await;
+        let mut params = PaintUniforms {
+            x: 128.0,
+            y: 128.0,
+            r: 0.5,
+            g: 0.5,
+            b: 0.5,
+            radius: 64.0,
+        };
+
+        self.optimize_param(&mut current_diff, &mut params, "", |p| p.b += 0.1);
+
+        self.exporter.run(&self.tex_approx, "a/approx.png").await;
+    }
+
+    pub fn optimize_param<F>(
+        &self,
+        current_diff: &mut u32,
+        params: &mut PaintUniforms,
+        label: &'static str,
+        f: F,
+    ) where
+        F: Fn(&mut PaintUniforms) -> (),
+    {
+        let mut p = params.clone();
+        loop {
+            self.save();
+            info!("apply '{label:?}'");
+            f(&mut p);
+            self.painter.run(p);
+            let diff = pollster::block_on(self.differ.run());
+            self.restore();
+            if diff >= *current_diff {
+                break;
+            }
+            *current_diff = diff;
+        }
+    }
+}
diff --git a/vgcodec/src/diff.rs b/vgcodec/src/diff.rs
index 44efabe..b10c4e6 100644
--- a/vgcodec/src/diff.rs
+++ b/vgcodec/src/diff.rs
@@ -78,7 +78,7 @@ impl Differ {
         }
     }
 
-    pub async fn run(&mut self) -> u32 {
+    pub async fn run(&self) -> u32 {
         let App { device, queue, .. } = self.app.as_ref();
 
         let mut encoder =
diff --git a/vgcodec/src/helper.rs b/vgcodec/src/helper.rs
new file mode 100644
index 0000000..3aca81d
--- /dev/null
+++ b/vgcodec/src/helper.rs
@@ -0,0 +1,20 @@
+use wgpu::{Queue, Texture, Extent3d};
+
+
+pub fn write_texture(queue: &Queue, target: &Texture, data: &[u8], size: Extent3d) {
+    queue.write_texture(
+        wgpu::ImageCopyTexture {
+            texture: &target,
+            mip_level: 0,
+            origin: wgpu::Origin3d::ZERO,
+            aspect: wgpu::TextureAspect::All,
+        },
+        &data,
+        wgpu::ImageDataLayout {
+            offset: 0,
+            bytes_per_row: Some(std::num::NonZeroU32::try_from((size.width * 4) as u32).unwrap()),
+            rows_per_image: None,
+        },
+        size,
+    );
+}
diff --git a/vgcodec/src/main.rs b/vgcodec/src/main.rs
index 9450894..7fb0750 100644
--- a/vgcodec/src/main.rs
+++ b/vgcodec/src/main.rs
@@ -1,14 +1,15 @@
 pub mod app;
+pub mod approximate;
 pub mod diff;
 pub mod export;
+pub mod helper;
 pub mod paint;
 
 use app::App;
-use diff::Differ;
+use approximate::Approximator;
+use helper::write_texture;
 use image::EncodableLayout;
-use wgpu::{Extent3d, Queue, Texture, TextureUsages};
-
-use crate::export::Exporter;
+use wgpu::TextureUsages;
 
 fn main() {
     env_logger::init_from_env("LOG");
@@ -21,56 +22,25 @@ async fn run() {
     let App { device, queue, .. } = app.as_ref();
 
     let img_target = image::open("a/a.png").unwrap().into_rgba8();
-    let img_initial = image::open("a/initial.png").unwrap().into_rgba8();
 
     let size = wgpu::Extent3d {
-        width: img_initial.width(),
-        height: img_initial.height(),
+        width: img_target.width(),
+        height: img_target.height(),
         depth_or_array_layers: 1,
     };
 
-    let create_texture = || {
-        device.create_texture(&wgpu::TextureDescriptor {
-            size,
-            mip_level_count: 1,
-            sample_count: 1,
-            dimension: wgpu::TextureDimension::D2,
-            format: wgpu::TextureFormat::Rgba8Unorm,
-            usage: TextureUsages::COPY_DST
-                | TextureUsages::TEXTURE_BINDING
-                | TextureUsages::COPY_SRC,
-            label: None,
-        })
-    };
-
-    let tex_target = create_texture();
-    let tex_approx = create_texture();
+    let tex_target = device.create_texture(&wgpu::TextureDescriptor {
+        size,
+        mip_level_count: 1,
+        sample_count: 1,
+        dimension: wgpu::TextureDimension::D2,
+        format: wgpu::TextureFormat::Rgba8Unorm,
+        usage: TextureUsages::COPY_DST | TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_SRC,
+        label: None,
+    });
 
     write_texture(queue, &tex_target, img_target.as_bytes(), size);
-    write_texture(queue, &tex_approx, img_initial.as_bytes(), size);
-
-    let mut differ = Differ::new(&app, size, &tex_approx, &tex_target);
-    let exporter = Exporter::new(&app, size);
 
-    println!("{}", differ.run().await);
-
-    exporter.run(&tex_approx, "a/approx.png").await;
-}
-
-pub fn write_texture(queue: &Queue, target: &Texture, data: &[u8], size: Extent3d) {
-    queue.write_texture(
-        wgpu::ImageCopyTexture {
-            texture: &target,
-            mip_level: 0,
-            origin: wgpu::Origin3d::ZERO,
-            aspect: wgpu::TextureAspect::All,
-        },
-        &data,
-        wgpu::ImageDataLayout {
-            offset: 0,
-            bytes_per_row: Some(std::num::NonZeroU32::try_from((size.width * 4) as u32).unwrap()),
-            rows_per_image: None,
-        },
-        size,
-    );
+    let mut a = Approximator::new(&app, tex_target, size);
+    a.run().await;
 }
diff --git a/vgcodec/src/paint.rs b/vgcodec/src/paint.rs
index 3da0647..f25b749 100644
--- a/vgcodec/src/paint.rs
+++ b/vgcodec/src/paint.rs
@@ -1,3 +1,4 @@
+use bytemuck::{Pod, Zeroable};
 use std::{borrow::Cow, sync::Arc};
 use wgpu::{util::DeviceExt, BindGroup, Buffer, ComputePipeline, Extent3d, Texture};
 
@@ -10,32 +11,33 @@ pub struct Painter {
     pipeline: ComputePipeline,
     bind_group: BindGroup,
 
-    uniform_buffer_size: u64,
     uniform_buffer: Buffer,
-    uniform_staging_buffer: Buffer,
+}
+
+#[repr(C)]
+#[derive(Pod, Zeroable, Clone, Copy)]
+pub struct PaintUniforms {
+    pub x: f32,
+    pub y: f32,
+    pub radius: f32,
+    pub r: f32,
+    pub g: f32,
+    pub b: f32,
 }
 
 impl Painter {
-    pub fn new(app: &Arc<App>, extent: Extent3d, tex_a: &Texture, tex_b: &Texture) -> Self {
+    pub fn new(app: &Arc<App>, extent: Extent3d, texture: &Texture) -> Self {
         let App { device, .. } = app.as_ref();
         let cs_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
             label: None,
             source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("paint.wgsl"))),
         });
 
-        let uniform_buffer_size = std::mem::size_of::<u32>() as wgpu::BufferAddress;
-        let uniform_staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
-            label: None,
-            size: uniform_buffer_size,
-            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
-            mapped_at_creation: false,
-        });
+        // let uniform_buffer_size = std::mem::size_of::<PaintUniforms>() as wgpu::BufferAddress;
         let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
             label: None,
-            contents: bytemuck::cast_slice(&[0u32]),
-            usage: wgpu::BufferUsages::STORAGE
-                | wgpu::BufferUsages::COPY_DST
-                | wgpu::BufferUsages::COPY_SRC,
+            contents: bytemuck::cast_slice(&[PaintUniforms::zeroed()]),
+            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
         });
 
         let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
@@ -52,17 +54,11 @@ impl Painter {
                 wgpu::BindGroupEntry {
                     binding: 0,
                     resource: wgpu::BindingResource::TextureView(
-                        &tex_a.create_view(&wgpu::TextureViewDescriptor::default()),
+                        &texture.create_view(&wgpu::TextureViewDescriptor::default()),
                     ),
                 },
                 wgpu::BindGroupEntry {
                     binding: 1,
-                    resource: wgpu::BindingResource::TextureView(
-                        &tex_b.create_view(&wgpu::TextureViewDescriptor::default()),
-                    ),
-                },
-                wgpu::BindGroupEntry {
-                    binding: 2,
                     resource: uniform_buffer.as_entire_binding(),
                 },
             ],
@@ -71,16 +67,16 @@ impl Painter {
             app: app.clone(),
             size: extent,
             pipeline,
-            uniform_buffer_size,
             uniform_buffer,
             bind_group,
-            uniform_staging_buffer,
         }
     }
 
-    pub async fn run(&mut self) -> u32 {
+    pub fn run(&self, params: PaintUniforms) {
         let App { device, queue, .. } = self.app.as_ref();
 
+        queue.write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[params]));
+
         let mut encoder =
             device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
         {
@@ -90,31 +86,6 @@ impl Painter {
             cpass.set_bind_group(0, &self.bind_group, &[]);
             cpass.dispatch_workgroups(self.size.width, self.size.height, 1);
         }
-
-        encoder.copy_buffer_to_buffer(
-            &self.uniform_buffer,
-            0,
-            &self.uniform_staging_buffer,
-            0,
-            self.uniform_buffer_size,
-        );
-
         queue.submit(Some(encoder.finish()));
-
-        let buffer_slice = self.uniform_staging_buffer.slice(..);
-        let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
-        buffer_slice.map_async(wgpu::MapMode::Read, move |v| {
-            sender.send(v.unwrap()).unwrap()
-        });
-        device.poll(wgpu::Maintain::Wait);
-        receiver.receive().await;
-
-        let data = buffer_slice.get_mapped_range();
-        let result: u32 = bytemuck::cast_slice(&data).to_vec()[0];
-
-        drop(data);
-        self.uniform_staging_buffer.unmap();
-
-        result
     }
 }
diff --git a/vgcodec/src/paint.wgsl b/vgcodec/src/paint.wgsl
index 02acd8e..f23e903 100644
--- a/vgcodec/src/paint.wgsl
+++ b/vgcodec/src/paint.wgsl
@@ -7,13 +7,13 @@ struct Uniforms {
     b: f32
 };
 
-@group(0) @binding(0) var<uniform> uniforms: Uniforms;
-@group(0) @binding(1) var tex: texture_storage_2d<rgba8unorm, write>;
+@group(0) @binding(0) var tex: texture_storage_2d<rgba8unorm, write>;
+@group(0) @binding(1) var<uniform> uniforms: Uniforms;
 
 @compute @workgroup_size(1)
 fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-    let coords = vec2<i32>(global_id.xy);
+    let coords = global_id.xy;
     if distance(vec2<f32>(coords), vec2(uniforms.x, uniforms.y)) < uniforms.radius {
-        textureStore(tex, coords.xy, vec4<f32>(vec3<f32>(uniforms.r, uniforms.g, uniforms.b), 1.0));
+        textureStore(tex, vec2<i32>(coords), vec4<f32>(vec3<f32>(uniforms.r, uniforms.g, uniforms.b), 1.0));
     }
 }