bottomless_pit/
shader.rs

1//! Contains everything you need to make some cool shaders including the Shader
2//! and UniformData types.
3
4use std::error::Error;
5use std::fmt::Display;
6use std::marker::PhantomData;
7use std::path::Path;
8use std::string::FromUtf8Error;
9
10use encase::private::WriteInto;
11use encase::ShaderType;
12use wgpu::util::DeviceExt;
13use wgpu::{include_wgsl, TextureFormat};
14
15use crate::context::{GraphicsContext, WgpuClump};
16use crate::engine_handle::Engine;
17use crate::render::Renderer;
18use crate::resource::{self, InProgressResource, LoadingOp, ResourceId, ResourceType};
19use crate::texture::{SamplerType, UniformTexture};
20use crate::vectors::Vec2;
21use crate::vertex::Vertex;
22use crate::{layouts, render};
23
24/// An internal representation of an WGSL Shader. Under the hood this creates
25/// a new pipeline with or without the support for any extra uniforms. To be utilze
26/// the shader it must be added to a material
27#[derive(Debug)]
28pub struct Shader {
29    pub(crate) pipeline: wgpu::RenderPipeline,
30    options: FinalShaderOptions,
31}
32
33impl Shader {
34    /// Attempts to create a shader from a file. Requires [ShaderOptions]
35    /// which implents Defualt for quick and easy creation.
36    pub fn new<P: AsRef<Path>, T>(
37        path: P,
38        options: ShaderOptions<T>,
39        engine: &mut Engine,
40        loading_op: LoadingOp,
41    ) -> ResourceId<Shader> {
42        let typed_id = resource::generate_id::<Shader>();
43        let id = typed_id.get_id();
44        let path = path.as_ref();
45        let ip_resource =
46            InProgressResource::new(path, id, ResourceType::Shader(options.into()), loading_op);
47
48        engine.loader.blocking_load(ip_resource, engine.get_proxy());
49
50        typed_id
51    }
52
53    pub(crate) fn from_resource_data(
54        data: &[u8],
55        options: FinalShaderOptions,
56        engine: &Engine,
57    ) -> Result<Self, FromUtf8Error> {
58        let context = engine.context.as_ref().unwrap();
59
60        let string = String::from_utf8(data.to_vec())?;
61        let shader = context
62            .wgpu
63            .device
64            .create_shader_module(wgpu::ShaderModuleDescriptor {
65                label: Some("User Shader Module"),
66                source: wgpu::ShaderSource::Wgsl(string.into()),
67            });
68
69        let optional_layout = options.make_layout(&context.wgpu.device);
70        let pipeline = if let Some(layout) = optional_layout {
71            render::make_pipeline(
72                &context.wgpu.device,
73                wgpu::PrimitiveTopology::TriangleList,
74                &[
75                    &layouts::create_texture_layout(&context.wgpu.device),
76                    &layouts::create_camera_layout(&context.wgpu.device),
77                    &layout,
78                ],
79                &[Vertex::desc()],
80                &shader,
81                context.get_texture_format(),
82                Some("User Shader Pipeline"),
83            )
84        } else {
85            render::make_pipeline(
86                &context.wgpu.device,
87                wgpu::PrimitiveTopology::TriangleList,
88                &[
89                    &layouts::create_texture_layout(&context.wgpu.device),
90                    &layouts::create_camera_layout(&context.wgpu.device),
91                ],
92                &[Vertex::desc()],
93                &shader,
94                context.get_texture_format(),
95                Some("User Shader Pipeline"),
96            )
97        };
98
99        Ok(Self { pipeline, options })
100    }
101
102    pub(crate) fn from_pipeline(pipeline: wgpu::RenderPipeline) -> Self {
103        Self {
104            pipeline,
105            options: FinalShaderOptions::EMPTY,
106        }
107    }
108
109    pub(crate) fn defualt(wgpu: &WgpuClump, texture_format: wgpu::TextureFormat) -> Self {
110        let shader_descriptor = include_wgsl!("shaders/shader.wgsl");
111        let shader = wgpu.device.create_shader_module(shader_descriptor);
112        let pipeline = render::make_pipeline(
113            &wgpu.device,
114            wgpu::PrimitiveTopology::TriangleList,
115            &[
116                &layouts::create_texture_layout(&wgpu.device),
117                &layouts::create_camera_layout(&wgpu.device),
118            ],
119            &[Vertex::desc()],
120            &shader,
121            texture_format,
122            Some("Defualt Shader From Error"),
123        );
124
125        Self {
126            pipeline,
127            options: FinalShaderOptions::EMPTY,
128        }
129    }
130
131    pub(crate) fn update_uniform_data<T: ShaderType + WriteInto>(
132        &self,
133        data: &T,
134        engine: &Engine,
135    ) -> Result<(), UniformError> {
136        self.options.update_uniform_data(data, engine)
137    }
138
139    pub(crate) fn update_uniform_texture(
140        &mut self,
141        texture: &mut UniformTexture,
142        wgpu: &WgpuClump,
143        format: TextureFormat,
144    ) -> Result<(), UniformError> {
145        self.options.update_uniform_texture(texture, wgpu, format)
146    }
147
148    pub(crate) fn set_active<'o>(&'o self, renderer: &mut Renderer<'o, '_>) {
149        renderer.pass.set_pipeline(&self.pipeline);
150
151        if let Some(bind_group) = &self.options.bind_group {
152            renderer.pass.set_bind_group(2, bind_group, &[]);
153        }
154    }
155}
156
157/// `UniformData` contains the byte data of any struct that implements
158/// [ShaderType](https://docs.rs/encase/latest/encase/trait.ShaderType.html) which
159/// can be derived. This data needs to be added to [ShaderOptions].
160#[derive(Debug)]
161pub struct UniformData<T> {
162    initial_data: Vec<u8>,
163    _marker: PhantomData<T>,
164}
165
166impl<T: ShaderType + WriteInto> UniformData<T> {
167    pub fn new(data: &T) -> Self {
168        let mut buffer = encase::UniformBuffer::new(Vec::new());
169        buffer.write(&data).unwrap();
170        let byte_array = buffer.into_inner();
171
172        Self {
173            initial_data: byte_array,
174            _marker: PhantomData,
175        }
176    }
177}
178
179/// `ShaderOptions` controlls the layout of your shader and wether or not you want
180/// extra uniforms like a texture or a uniform buffer
181#[derive(Debug)]
182pub struct ShaderOptions<T> {
183    uniform_data: Option<Vec<u8>>,
184    uniform_texture: Option<(SamplerType, SamplerType, Vec2<u32>)>,
185    _marker: PhantomData<T>,
186}
187
188// switch texture with booleans and only store a layout
189// figure out how to make view later?
190
191impl<T> ShaderOptions<T> {
192    pub const EMPTY: Self = Self {
193        uniform_data: None,
194        uniform_texture: None,
195        _marker: PhantomData,
196    };
197
198    /// This will create a shader with a layout of:
199    /// ```wgsl
200    /// struct MousePos {
201    ///     stuff: vec2<f32>,
202    ///     _junk: vec2<f32>,
203    /// }
204    ///
205    /// @group(2) @binding(0)
206    /// var<uniform> mouse: MousePos;
207    /// ```
208    pub fn with_uniform_data(data: &UniformData<T>) -> Self {
209        // dont like doing this bc clone expensive but we need
210        // to make this before wgpu initlizes
211        let starting_buffer = data.initial_data.clone();
212
213        Self {
214            uniform_data: Some(starting_buffer),
215            uniform_texture: None,
216            _marker: PhantomData,
217        }
218    }
219
220    /// This will create a shader with a layout of:
221    /// ```wgsl
222    /// @group(2) @binding(0)
223    /// var light_map: texture_2d<f32>;
224    /// @group(2) @binding(1)
225    /// var light_map_sampler: sampler;
226    /// ```
227    pub fn with_uniform_texture(texture: &UniformTexture) -> Self {
228        let (mag, min) = texture.get_sampler_info();
229        let size = texture.get_size();
230
231        Self {
232            uniform_data: None,
233            uniform_texture: Some((mag, min, size)),
234            _marker: PhantomData,
235        }
236    }
237
238    /// This will create a shader with a layout of:
239    /// ```wgsl
240    /// @group(2) binding(0)
241    /// var<uniform> value: SomeStruct;
242    /// @group(2) @binding(1)
243    /// var light_map: texture_2d<f32>;
244    /// @group(2) @binding(2)
245    /// var light_map_sampler: sampler;
246    /// ```
247    pub fn with_all(data: &UniformData<T>, texture: &UniformTexture) -> Self {
248        let starting_buffer = data.initial_data.clone();
249        let (mag, min) = texture.get_sampler_info();
250        let size = texture.get_size();
251
252        Self {
253            uniform_data: Some(starting_buffer),
254            uniform_texture: Some((mag, min, size)),
255            _marker: PhantomData,
256        }
257    }
258}
259
260#[derive(Debug)]
261pub(crate) struct IntermediateOptions {
262    uniform_data: Option<Vec<u8>>,
263    uniform_texture: Option<(SamplerType, SamplerType, Vec2<u32>)>,
264}
265
266impl IntermediateOptions {
267    pub(crate) fn check_has(&self) -> (bool, bool) {
268        (self.uniform_data.is_some(), self.uniform_texture.is_some())
269    }
270}
271
272impl<T> From<ShaderOptions<T>> for IntermediateOptions {
273    fn from(value: ShaderOptions<T>) -> Self {
274        Self {
275            uniform_data: value.uniform_data,
276            uniform_texture: value.uniform_texture,
277        }
278    }
279}
280
281#[derive(Debug)]
282pub(crate) struct FinalShaderOptions {
283    uniform_data: Option<wgpu::Buffer>,
284    uniform_texture: Option<(wgpu::TextureView, wgpu::Sampler)>,
285    bind_group: Option<wgpu::BindGroup>,
286}
287
288impl FinalShaderOptions {
289    pub(crate) const EMPTY: Self = Self {
290        uniform_data: None,
291        uniform_texture: None,
292        bind_group: None,
293    };
294
295    pub(crate) fn from_intermediate(
296        options: IntermediateOptions,
297        context: &GraphicsContext,
298    ) -> Self {
299        let wgpu = &context.wgpu;
300        let format = context.get_texture_format();
301
302        let buffer = options.uniform_data.map(|d| {
303            wgpu.device
304                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
305                    label: Some("User Uniform Data Buffer"),
306                    contents: &d,
307                    usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
308                })
309        });
310
311        let uniform_texture = options.uniform_texture.map(|(mag, min, size)| {
312            let sampler = wgpu.device.create_sampler(&wgpu::SamplerDescriptor {
313                address_mode_u: wgpu::AddressMode::Repeat,
314                address_mode_v: wgpu::AddressMode::Repeat,
315                address_mode_w: wgpu::AddressMode::ClampToEdge,
316                // what do when give less or more than 1 pixel to sample
317                // linear interprelates between all of them nearest gives the closet colour
318                mag_filter: mag.into(),
319                min_filter: min.into(),
320                mipmap_filter: wgpu::FilterMode::Nearest,
321                ..Default::default()
322            });
323
324            let texture = wgpu.device.create_texture(&wgpu::TextureDescriptor {
325                label: Some("Uniform Texture"),
326                size: wgpu::Extent3d {
327                    width: size.x,
328                    height: size.y,
329                    depth_or_array_layers: 1,
330                },
331                dimension: wgpu::TextureDimension::D2,
332                mip_level_count: 1,
333                sample_count: 1,
334                format,
335                usage: wgpu::TextureUsages::RENDER_ATTACHMENT
336                    | wgpu::TextureUsages::TEXTURE_BINDING,
337                view_formats: &[],
338            });
339
340            let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
341
342            (view, sampler)
343        });
344
345        let bind_group =
346            make_layout_internal(&wgpu.device, buffer.is_some(), uniform_texture.is_some()).map(
347                |layout| {
348                    let entires = make_entries(
349                        buffer.as_ref(),
350                        uniform_texture.as_ref().map(|(a, b)| (a, b)),
351                        None,
352                    );
353
354                    wgpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
355                        label: Some("User Shader Option Bind Group"),
356                        layout: &layout,
357                        entries: &entires,
358                    })
359                },
360            );
361
362        Self {
363            uniform_data: buffer,
364            uniform_texture,
365            bind_group,
366        }
367    }
368
369    pub(crate) fn make_layout(&self, device: &wgpu::Device) -> Option<wgpu::BindGroupLayout> {
370        make_layout_internal(
371            device,
372            self.uniform_data.is_some(),
373            self.uniform_texture.is_some(),
374        )
375    }
376
377    fn update_uniform_data<H: ShaderType + WriteInto>(
378        &self,
379        data: &H,
380        engine: &Engine,
381    ) -> Result<(), UniformError> {
382        match &self.uniform_data {
383            Some(buffer) => {
384                let wgpu = &engine
385                    .context
386                    .as_ref()
387                    .expect("NEED CONTEXT BEFORE UPDATING UNIFORM DATA")
388                    .wgpu;
389                let mut uniform_buffer = encase::UniformBuffer::new(Vec::new());
390                uniform_buffer.write(&data).unwrap();
391                let byte_array = uniform_buffer.into_inner();
392
393                wgpu.queue.write_buffer(buffer, 0, &byte_array);
394                Ok(())
395            }
396            None => Err(UniformError::DoesntHaveUniformBuffer),
397        }
398    }
399
400    fn update_uniform_texture(
401        &mut self,
402        texture: &mut UniformTexture,
403        wgpu: &WgpuClump,
404        format: wgpu::TextureFormat,
405    ) -> Result<(), UniformError> {
406        if !texture.needs_update() {
407            Ok(())?;
408        }
409
410        texture.updated();
411
412        match &mut self.uniform_texture {
413            Some(view) => {
414                view.0 = texture.make_view(wgpu, format);
415
416                self.bind_group = Some(wgpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
417                    label: Some("Shader Options BindGroup"),
418                    entries: &make_entries(
419                        self.uniform_data.as_ref(),
420                        self.uniform_texture.as_ref().map(|(a, b)| (a, b)),
421                        Some(texture.get_sampler()),
422                    ),
423                    layout: &self.make_layout(&wgpu.device).unwrap(),
424                }));
425
426                Ok(())
427            }
428            None => Err(UniformError::DoesntHaveUniformTexture),
429        }
430    }
431}
432
433#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
434pub enum UniformError {
435    NotLoadedYet,
436    DoesntHaveUniformBuffer,
437    DoesntHaveUniformTexture,
438}
439
440impl Display for UniformError {
441    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
442        match self {
443            Self::NotLoadedYet => write!(f, "The shader has not loaded yet please try again later"),
444            Self::DoesntHaveUniformBuffer => write!(f, "The shader does not have a uniform buffer"),
445            Self::DoesntHaveUniformTexture => {
446                write!(f, "The shader does not have a unform texture")
447            }
448        }
449    }
450}
451
452impl Error for UniformError {}
453
454fn make_layout_internal(
455    device: &wgpu::Device,
456    has_buffer: bool,
457    has_texture: bool,
458) -> Option<wgpu::BindGroupLayout> {
459    match (has_buffer, has_texture) {
460        (true, true) => Some(layouts::create_texture_uniform_layout(device)),
461        (true, false) => Some(layouts::create_uniform_layout(device)),
462        (false, true) => Some(layouts::create_texture_layout(device)),
463        (false, false) => None,
464    }
465}
466
467fn make_entries<'a>(
468    uniform_data: Option<&'a wgpu::Buffer>,
469    uniform_texture: Option<(&'a wgpu::TextureView, &'a wgpu::Sampler)>,
470    other_sampler: Option<&'a wgpu::Sampler>,
471) -> Vec<wgpu::BindGroupEntry<'a>> {
472    let mut entries = Vec::with_capacity(3);
473
474    if let Some(buffer) = uniform_data {
475        entries.push(wgpu::BindGroupEntry {
476            binding: entries.len() as u32,
477            resource: wgpu::BindingResource::Buffer(buffer.as_entire_buffer_binding()),
478        });
479    }
480
481    if let Some((view, sampler)) = uniform_texture {
482        entries.push(wgpu::BindGroupEntry {
483            binding: entries.len() as u32,
484            resource: wgpu::BindingResource::TextureView(view),
485        });
486        entries.push(wgpu::BindGroupEntry {
487            binding: entries.len() as u32,
488            resource: wgpu::BindingResource::Sampler(other_sampler.unwrap_or(sampler)),
489        });
490    }
491
492    entries
493}