a

2 files changed, 84 insertions, 23 deletions

diff --git a/flowy/src/main.rs b/flowy/src/main.rs
index d225c00..19084ef 100644
--- a/flowy/src/main.rs
+++ b/flowy/src/main.rs
@@ -44,7 +44,7 @@ fn main() {
         depth_or_array_layers: 1,
     };
 
-    let [texa, texb] = [(), ()].map(|_| {
+    let textures = [(), ()].map(|_| {
         device.create_texture(&TextureDescriptor {
             label: None,
             size: extent,
@@ -99,27 +99,46 @@ fn main() {
                 },
                 visibility: ShaderStages::COMPUTE,
             },
-        ],
-    });
-
-    let bind_goup = device.create_bind_group(&BindGroupDescriptor {
-        label: None,
-        layout: &bind_group_layout,
-        entries: &[
-            BindGroupEntry {
-                binding: 0,
-                resource: wgpu::BindingResource::TextureView(
-                    &out_tex.create_view(&Default::default()),
-                ),
-            },
-            BindGroupEntry {
-                binding: 1,
-                resource: wgpu::BindingResource::TextureView(
-                    &texa.create_view(&Default::default()),
-                ),
+            BindGroupLayoutEntry {
+                binding: 2,
+                count: None,
+                ty: BindingType::Texture {
+                    sample_type: TextureSampleType::Float { filterable: false },
+                    view_dimension: TextureViewDimension::D2,
+                    multisampled: false,
+                },
+                visibility: ShaderStages::COMPUTE,
             },
         ],
     });
+    let bind_groups = [false, true].map(|i| {
+        device.create_bind_group(&BindGroupDescriptor {
+            label: None,
+            layout: &bind_group_layout,
+            entries: &[
+                BindGroupEntry {
+                    binding: 0,
+                    resource: wgpu::BindingResource::TextureView(
+                        &out_tex.create_view(&Default::default()),
+                    ),
+                },
+                BindGroupEntry {
+                    binding: 1,
+                    resource: wgpu::BindingResource::TextureView(
+                        &(if i { &textures[0] } else { &textures[1] })
+                            .create_view(&Default::default()),
+                    ),
+                },
+                BindGroupEntry {
+                    binding: 2,
+                    resource: wgpu::BindingResource::TextureView(
+                        &(if i { &textures[1] } else { &textures[0] })
+                            .create_view(&Default::default()),
+                    ),
+                },
+            ],
+        })
+    });
 
     let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
         label: None,
@@ -135,6 +154,7 @@ fn main() {
         entry_point: "main",
     });
 
+    let mut i = 0;
     loop {
         eprintln!("read");
         stdin().read_exact(&mut buffer).unwrap();
@@ -145,7 +165,7 @@ fn main() {
                 aspect: wgpu::TextureAspect::All,
                 mip_level: 0,
                 origin: Origin3d::ZERO,
-                texture: &texa,
+                texture: &textures[i],
             },
             &buffer,
             wgpu::ImageDataLayout {
@@ -165,7 +185,7 @@ fn main() {
                 timestamp_writes: None,
             });
             cpass.set_pipeline(&pipeline);
-            cpass.set_bind_group(0, &bind_goup, &[]);
+            cpass.set_bind_group(0, &bind_groups[i], &[]);
             cpass.dispatch_workgroups(width as u32, height as u32, 1);
         }
 
@@ -203,5 +223,7 @@ fn main() {
         download_buffer.unmap();
         eprintln!("write");
         stdout().write_all(&buffer).unwrap();
+        i += 1;
+        i %= 2;
     }
 }
diff --git a/flowy/src/shader.wgsl b/flowy/src/shader.wgsl
index 1597ff6..bfa4176 100644
--- a/flowy/src/shader.wgsl
+++ b/flowy/src/shader.wgsl
@@ -1,11 +1,50 @@
 
 @group(0) @binding(0) var output: texture_storage_2d<bgra8unorm, write>;
 @group(0) @binding(1) var texa: texture_2d<f32>;
+@group(0) @binding(2) var texb: texture_2d<f32>;
 
+const SDIST: i32 = 4;
+
+const NN: vec2<i32> = vec2(-1,-1);
+const NP: vec2<i32> = vec2(-1,1);
+const PN: vec2<i32> = vec2(1,-1);
+const PP: vec2<i32> = vec2(1,1);
 
 @compute @workgroup_size(1) fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-    let p = textureLoad(texa, global_id.xy, 0);
-    textureStore(output, global_id.xy, p);
+    let uv = vec2<i32>(global_id.xy);
+    let orig = textureLoad(texa, uv, 0);
+    let orignn = textureLoad(texa, uv+NN, 0);
+    let orignp = textureLoad(texa, uv+NP, 0);
+    let origpn = textureLoad(texa, uv+PN, 0);
+    let origpp = textureLoad(texa, uv+PP, 0);
+
+    var best_sim = 1000.;
+    var best_coord = vec2(0,0);
+    for (var x = -SDIST; x <= SDIST; x++) {
+    for (var y = -SDIST; y <= SDIST; y++) {
+        let ov = uv + vec2(x,y);
+        let samp = textureLoad(texb, ov, 0);
+        let sampnn = textureLoad(texb, ov+NN, 0);
+        let sampnp = textureLoad(texb, ov+NP, 0);
+        let samppn = textureLoad(texb, ov+PN, 0);
+        let samppp = textureLoad(texb, ov+PP, 0);
+        let d = distance(samp.rgb, orig.rgb)
+            + distance(sampnn.rgb, orignn.rgb)
+            + distance(sampnp.rgb, orignp.rgb)
+            + distance(samppn.rgb, origpn.rgb)
+            + distance(samppp.rgb, origpp.rgb);
+        if d < best_sim {
+            best_sim = d;
+            best_coord = vec2(x,y);
+        }
+    }}
+
+    let col = vec4(colormap_vec(vec2<f32>(best_coord) / f32(SDIST)), 1.);
+
+    textureStore(output, global_id.xy, col);
 }
 
+fn colormap_vec(v: vec2<f32>) -> vec3<f32> {
+    return vec3(v.y, v.x - 0.5 * v.y, -v.x - 0.5 * v.y);
+}