Merge pull request #12 from pemattern/11-add-option-for-custominput-data-to-tuishaderbackend

feat: user data can be added

Author: pemattern

src/backend/cpu.rs

@@ -0,0 +1,64 @@
+use super::{NoUserData, TuiShaderBackend};
+
+pub struct CpuBackend<T> {
+    callback: Box<dyn CpuShaderCallback<T>>,
+}
+
+impl CpuBackend<NoUserData> {
+    pub fn new<F>(callback: F) -> Self
+    where
+        F: Fn(u16, u16) -> [u8; 4] + 'static,
+    {
+        Self {
+            callback: Box::new(CpuShaderCallbackWithoutUserData(callback)),
+        }
+    }
+}
+
+impl<T> CpuBackend<T> {
+    pub fn new_with_user_data<F>(callback: F) -> Self
+    where
+        F: Fn(u16, u16, &T) -> [u8; 4] + 'static,
+    {
+        Self {
+            callback: Box::new(callback),
+        }
+    }
+}
+
+impl<T> TuiShaderBackend<T> for CpuBackend<T> {
+    fn execute(&mut self, width: u16, height: u16, user_data: &T) -> Vec<[u8; 4]> {
+        let mut pixels = Vec::new();
+        for y in 0..height {
+            for x in 0..width {
+                let value = self.callback.call(x, y, user_data);
+                pixels.push(value);
+            }
+        }
+        pixels
+    }
+}
+
+pub trait CpuShaderCallback<T> {
+    fn call(&self, x: u16, y: u16, user_data: &T) -> [u8; 4];
+}
+
+impl<T, F> CpuShaderCallback<T> for F
+where
+    F: Fn(u16, u16, &T) -> [u8; 4],
+{
+    fn call(&self, x: u16, y: u16, user_data: &T) -> [u8; 4] {
+        self(x, y, user_data)
+    }
+}
+
+// NewType required to avoid conflicting implementations
+struct CpuShaderCallbackWithoutUserData<F>(F);
+impl<F> CpuShaderCallback<NoUserData> for CpuShaderCallbackWithoutUserData<F>
+where
+    F: Fn(u16, u16) -> [u8; 4],
+{
+    fn call(&self, x: u16, y: u16, _user_data: &NoUserData) -> [u8; 4] {
+        self.0(x, y)
+    }
+}

src/backend/mod.rs

@@ -1,5 +1,10 @@
+pub mod cpu;
 pub mod wgpu;
 
-pub trait TuiShaderBackend: Eq {
-    fn execute(&mut self, width: u16, height: u16) -> Vec<[u8; 4]>;
+#[repr(C)]
+#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
+pub struct NoUserData;
+
+pub trait TuiShaderBackend<T> {
+    fn execute(&mut self, width: u16, height: u16, user_data: &T) -> Vec<[u8; 4]>;
 }

src/backend/wgpu.rs

@@ -1,12 +1,12 @@
-use std::{fs, time::Instant};
+use std::{fs, marker::PhantomData, time::Instant};
 
 use pollster::FutureExt as _;
 use wgpu::util::DeviceExt;
 
-use super::TuiShaderBackend;
+use super::{NoUserData, TuiShaderBackend};
 
 #[repr(C)]
-#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
+#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
 struct ShaderInput {
     // struct field order matters
     time: f32,
@@ -15,7 +15,10 @@ struct ShaderInput {
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]
-pub struct WgpuBackend {
+pub struct WgpuBackend<T>
+where
+    T: Copy + Clone + bytemuck::Pod + bytemuck::Zeroable,
+{
     device: wgpu::Device,
     queue: wgpu::Queue,
     pipeline: wgpu::RenderPipeline,
@@ -26,9 +29,13 @@ pub struct WgpuBackend {
     bind_group: wgpu::BindGroup,
     width: u16,
     height: u16,
+    _user_data: PhantomData<T>,
 }
 
-impl WgpuBackend {
+impl<T> WgpuBackend<T>
+where
+    T: Copy + Clone + bytemuck::Pod + bytemuck::Zeroable,
+{
     pub fn new(path_to_fragment_shader: &str, entry_point: &str) -> Self {
         Self::new_inner(path_to_fragment_shader, entry_point).block_on()
     }
@@ -70,8 +77,8 @@ impl WgpuBackend {
         let width = 64u16;
         let height = 64u16;
 
-        let texture = WgpuBackend::create_texture(&device, width.into(), height.into());
-        let buffer = WgpuBackend::create_buffer(&device, width.into(), height.into());
+        let texture = WgpuBackend::<T>::create_texture(&device, width.into(), height.into());
+        let output_buffer = WgpuBackend::<T>::create_buffer(&device, width.into(), height.into());
 
         let shader_input = ShaderInput {
             time: creation_time.elapsed().as_secs_f32(),
@@ -150,11 +157,12 @@ impl WgpuBackend {
             pipeline,
             creation_time,
             texture,
-            output_buffer: buffer,
-            width,
-            height,
+            output_buffer,
             shader_input_buffer,
             bind_group,
+            width,
+            height,
+            _user_data: PhantomData,
         }
     }
 
@@ -188,24 +196,23 @@ impl WgpuBackend {
         })
     }
 
-    fn bytes_per_row(width: u16) -> u16 {
+    fn bytes_per_row(&self, width: u16) -> u16 {
         let row_size = width * 4;
         (row_size + 255) & !255
     }
 
-    pub fn row_padding(width: u16) -> u16 {
+    fn row_padding(&self, width: u16) -> u16 {
         let row_size = width * 4;
-        let bytes_per_row = Self::bytes_per_row(width);
+        let bytes_per_row = self.bytes_per_row(width);
         (bytes_per_row - row_size) / 4
     }
 
-    async fn execute_inner(&mut self, width: u16, height: u16) -> Vec<[u8; 4]> {
-        if Self::bytes_per_row(width) != Self::bytes_per_row(self.width) || height != self.height {
+    async fn execute_inner(&mut self, width: u16, height: u16, _user_data: &T) -> Vec<[u8; 4]> {
+        if self.bytes_per_row(width) != self.bytes_per_row(self.width) || height != self.height {
             self.texture = Self::create_texture(&self.device, width.into(), height.into());
             self.output_buffer = Self::create_buffer(&self.device, width.into(), height.into());
         }
-
-        let bytes_per_row = Self::bytes_per_row(width);
+        let bytes_per_row = self.bytes_per_row(width);
 
         let texture_view = self
             .texture
@@ -283,23 +290,34 @@ impl WgpuBackend {
             .await
             .expect("unable to receive message all senders have been dropped")
             .expect("on unexpected error occured");
-        let slice: Vec<[u8; 4]>;
+        let padded_buffer: Vec<[u8; 4]>;
         {
             let view = buffer_slice.get_mapped_range();
-            slice = bytemuck::cast_slice(&view).to_vec();
+            padded_buffer = bytemuck::cast_slice(&view).to_vec();
         }
         self.output_buffer.unmap();
-        slice
+        let mut buffer: Vec<[u8; 4]> = Vec::new();
+        for y in 0..height {
+            for x in 0..width {
+                let index = (y * (width + self.row_padding(width)) + x) as usize;
+                let pixel = padded_buffer[index];
+                buffer.push(pixel);
+            }
+        }
+        buffer
     }
 }
 
-impl TuiShaderBackend for WgpuBackend {
-    fn execute(&mut self, width: u16, height: u16) -> Vec<[u8; 4]> {
-        self.execute_inner(width, height).block_on()
+impl<T> TuiShaderBackend<T> for WgpuBackend<T>
+where
+    T: Copy + Clone + bytemuck::Pod + bytemuck::Zeroable,
+{
+    fn execute(&mut self, width: u16, height: u16, user_data: &T) -> Vec<[u8; 4]> {
+        self.execute_inner(width, height, user_data).block_on()
     }
 }
 
-impl Default for WgpuBackend {
+impl Default for WgpuBackend<NoUserData> {
     fn default() -> Self {
         Self::new("src/shaders/default_fragment.wgsl", "magenta")
     }

src/lib.rs

@@ -31,7 +31,8 @@ mod backend;
 
 use std::marker::PhantomData;
 
-use backend::TuiShaderBackend;
+use backend::cpu::CpuBackend;
+use backend::{NoUserData, TuiShaderBackend};
 use ratatui::layout::{Position, Rect};
 use ratatui::style::{Color, Style};
 use ratatui::widgets::StatefulWidget;
@@ -59,17 +60,17 @@ pub struct ShaderCanvas<T> {
     pub character_rule: CharacterRule,
     pub style_rule: StyleRule,
     pub entry_point: String,
-    _marker: PhantomData<T>,
+    _user_data: PhantomData<T>,
 }
 
-impl<T: TuiShaderBackend> ShaderCanvas<T> {
+impl<T> ShaderCanvas<T> {
     /// Creates a new instance of [`ShaderCanvas`].
     pub fn new() -> Self {
         Self {
             character_rule: CharacterRule::default(),
             style_rule: StyleRule::default(),
             entry_point: String::from("main"),
-            _marker: PhantomData,
+            _user_data: PhantomData,
         }
     }
 
@@ -96,13 +97,13 @@ impl<T: TuiShaderBackend> ShaderCanvas<T> {
     }
 }
 
-impl<T: TuiShaderBackend> Default for ShaderCanvas<T> {
+impl<T> Default for ShaderCanvas<T> {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl<T: TuiShaderBackend> StatefulWidget for ShaderCanvas<T> {
+impl<T> StatefulWidget for ShaderCanvas<T> {
     type State = ShaderCanvasState<T>;
     fn render(
         self,
@@ -112,13 +113,12 @@ impl<T: TuiShaderBackend> StatefulWidget for ShaderCanvas<T> {
     ) {
         let width = area.width;
         let height = area.height;
-
-        let raw_buffer = state.backend.execute(width, height);
+        let samples = state.backend.execute(width, height, &state.user_data);
 
         for y in 0..height {
             for x in 0..width {
-                let index = (y * (width + WgpuBackend::row_padding(width)) + x) as usize;
-                let value = raw_buffer[index];
+                let index = (y * width + x) as usize;
+                let value = samples[index];
                 let position = (x, y);
                 let character = match self.character_rule {
                     CharacterRule::Always(character) => character,
@@ -142,35 +142,69 @@ impl<T: TuiShaderBackend> StatefulWidget for ShaderCanvas<T> {
 }
 
 /// State struct for [`ShaderCanvas`], it holds the [`TuiShaderBackend`].
-#[derive(Debug, Clone, Eq, PartialEq)]
-pub struct ShaderCanvasState<T: TuiShaderBackend> {
-    backend: T,
+pub struct ShaderCanvasState<T> {
+    backend: Box<dyn TuiShaderBackend<T>>,
+    user_data: T,
 }
 
-impl ShaderCanvasState<WgpuBackend> {
+impl ShaderCanvasState<NoUserData> {
     /// Creates a new [`ShaderCanvasState`] using [`WgpuBackend`] as it's
     /// [`TuiShaderBackend`].
-    pub fn wgpu(
+    pub fn wgpu(path_to_fragment_shader: &str, entry_point: &str) -> Self {
+        let backend = Box::new(WgpuBackend::new(path_to_fragment_shader, entry_point));
+        Self {
+            backend,
+            user_data: NoUserData,
+        }
+    }
+}
+
+impl<T> ShaderCanvasState<T>
+where
+    T: Copy + bytemuck::Pod + bytemuck::Zeroable,
+{
+    pub fn wgpu_with_user_data(
         path_to_fragment_shader: &str,
         entry_point: &str,
-    ) -> ShaderCanvasState<WgpuBackend> {
-        let backend = WgpuBackend::new(path_to_fragment_shader, entry_point);
-        ShaderCanvasState { backend }
+        user_data: T,
+    ) -> Self {
+        let backend = Box::new(WgpuBackend::new(path_to_fragment_shader, entry_point));
+        Self { backend, user_data }
     }
 }
 
-impl<T: TuiShaderBackend> ShaderCanvasState<T> {
-    /// Creates a new [`ShaderCanvasState`] instance by passing in the desired [`TuiShaderBackend`].
-    pub fn new(backend: T) -> ShaderCanvasState<T> {
-        ShaderCanvasState { backend }
+impl ShaderCanvasState<NoUserData> {
+    pub fn cpu<F>(callback: F) -> Self
+    where
+        F: Fn(u16, u16) -> [u8; 4] + 'static,
+    {
+        let backend = Box::new(CpuBackend::new(callback));
+        Self {
+            backend,
+            user_data: NoUserData,
+        }
+    }
+}
+
+impl<T> ShaderCanvasState<T>
+where
+    T: 'static,
+{
+    pub fn cpu_with_user_data<F>(callback: F, user_data: T) -> Self
+    where
+        F: Fn(u16, u16, &T) -> [u8; 4] + 'static,
+    {
+        let backend = Box::new(CpuBackend::new_with_user_data(callback));
+        Self { backend, user_data }
     }
 }
 
-impl Default for ShaderCanvasState<WgpuBackend> {
+impl Default for ShaderCanvasState<NoUserData> {
     /// Creates a new [`ShaderCanvasState`] instance with a [`WgpuBackend`].
-    fn default() -> ShaderCanvasState<WgpuBackend> {
+    fn default() -> Self {
         Self {
-            backend: WgpuBackend::default(),
+            backend: Box::new(WgpuBackend::default()),
+            user_data: NoUserData,
         }
     }
 }
@@ -262,14 +296,14 @@ mod tests {
     #[test]
     fn default_wgsl_context() {
         let mut context = WgpuBackend::default();
-        let raw_buffer = context.execute(64, 64);
+        let raw_buffer = context.execute(64, 64, &NoUserData);
         assert!(raw_buffer.iter().all(|pixel| pixel == &[255, 0, 255, 255]));
     }
 
     #[test]
     fn different_entry_points() {
         let mut context = WgpuBackend::new("src/shaders/default_fragment.wgsl", "green");
-        let raw_buffer = context.execute(64, 64);
+        let raw_buffer = context.execute(64, 64, &NoUserData);
         assert!(raw_buffer.iter().all(|pixel| pixel == &[0, 255, 0, 255]));
     }