1
use crate::camera::Camera;
2
use crate::mesh::{Mesh, Vertex};
3
use anyhow::Result;
4
use bytemuck::{Pod, Zeroable};
5
use wgpu::util::DeviceExt;
6
 
7
#[repr(C)]
8
#[derive(Copy, Clone, Debug, Pod, Zeroable)]
9
struct Uniforms {
10
    view_proj: [[f32; 4]; 4],
11
    model: [[f32; 4]; 4],
12
}
13
 
14
pub struct Renderer {
15
    device: wgpu::Device,
16
    queue: wgpu::Queue,
17
    pipeline: wgpu::RenderPipeline,
18
    uniform_bind_group_layout: wgpu::BindGroupLayout,
19
    width: u32,
20
    height: u32,
21
}
22
 
23
impl Renderer {
24
    pub async fn new(width: u32, height: u32) -> Result<Self> {
25
        // Initialize WGPU
26
        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
27
        let adapter = instance
28
            .request_adapter(&wgpu::RequestAdapterOptions {
29
                power_preference: wgpu::PowerPreference::HighPerformance,
30
                force_fallback_adapter: false,
31
                compatible_surface: None,
32
            })
33
            .await
34
            .ok_or_else(|| anyhow::anyhow!("Failed to find a suitable GPU adapter"))?;
35
 
36
        let (device, queue) = adapter
37
            .request_device(
38
                &wgpu::DeviceDescriptor {
39
                    label: Some("Device"),
40
                    required_features: wgpu::Features::empty(),
41
                    required_limits: wgpu::Limits::default(),
42
                    memory_hints: wgpu::MemoryHints::default(),
43
                },
44
                None,
45
            )
46
            .await?;
47
 
48
        // Create bind group layout for uniforms
49
        let uniform_bind_group_layout =
50
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
51
                label: Some("Uniform Bind Group Layout"),
52
                entries: &[wgpu::BindGroupLayoutEntry {
53
                    binding: 0,
54
                    visibility: wgpu::ShaderStages::VERTEX,
55
                    ty: wgpu::BindingType::Buffer {
56
                        ty: wgpu::BufferBindingType::Uniform,
57
                        has_dynamic_offset: false,
58
                        min_binding_size: None,
59
                    },
60
                    count: None,
61
                }],
62
            });
63
 
64
        // Create vertex shader module
65
        let vertex_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
66
            label: Some("Vertex Shader"),
67
            source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(include_str!(
68
                "../shaders/vertex.wgsl"
69
            ))),
70
        });
71
 
72
        // Create fragment shader module
73
        let fragment_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
74
            label: Some("Fragment Shader"),
75
            source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(include_str!(
76
                "../shaders/fragment.wgsl"
77
            ))),
78
        });
79
 
80
        // Create render pipeline
81
        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
82
            label: Some("Render Pipeline Layout"),
83
            bind_group_layouts: &[&uniform_bind_group_layout],
84
            push_constant_ranges: &[],
85
        });
86
 
87
        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
88
            label: Some("Render Pipeline"),
89
            layout: Some(&pipeline_layout),
90
            vertex: wgpu::VertexState {
91
                module: &vertex_shader,
92
                entry_point: Some("vs_main"),
93
                buffers: &[Vertex::desc()],
94
                compilation_options: wgpu::PipelineCompilationOptions::default(),
95
            },
96
            fragment: Some(wgpu::FragmentState {
97
                module: &fragment_shader,
98
                entry_point: Some("fs_main"),
99
                targets: &[Some(wgpu::ColorTargetState {
100
                    format: wgpu::TextureFormat::Rgba8Unorm,
101
                    blend: Some(wgpu::BlendState::REPLACE),
102
                    write_mask: wgpu::ColorWrites::ALL,
103
                })],
104
                compilation_options: wgpu::PipelineCompilationOptions::default(),
105
            }),
106
            primitive: wgpu::PrimitiveState {
107
                topology: wgpu::PrimitiveTopology::TriangleList,
108
                strip_index_format: None,
109
                front_face: wgpu::FrontFace::Ccw,
110
                cull_mode: Some(wgpu::Face::Back),
111
                polygon_mode: wgpu::PolygonMode::Fill,
112
                unclipped_depth: false,
113
                conservative: false,
114
            },
115
            depth_stencil: Some(wgpu::DepthStencilState {
116
                format: wgpu::TextureFormat::Depth32Float,
117
                depth_write_enabled: true,
118
                depth_compare: wgpu::CompareFunction::Less,
119
                stencil: wgpu::StencilState::default(),
120
                bias: wgpu::DepthBiasState::default(),
121
            }),
122
            multisample: wgpu::MultisampleState {
123
                count: 1,
124
                mask: !0,
125
                alpha_to_coverage_enabled: false,
126
            },
127
            multiview: None,
128
            cache: None,
129
        });
130
 
131
        Ok(Self {
132
            device,
133
            queue,
134
            pipeline,
135
            uniform_bind_group_layout,
136
            width,
137
            height,
138
        })
139
    }
140
 
141
    pub async fn render(&self, mesh: &Mesh, camera: &Camera) -> Result<Vec<u8>> {
142
        // Create render texture
143
        let texture_extent = wgpu::Extent3d {
144
            width: self.width,
145
            height: self.height,
146
            depth_or_array_layers: 1,
147
        };
148
 
149
        let texture = self.device.create_texture(&wgpu::TextureDescriptor {
150
            label: Some("Render Texture"),
151
            size: texture_extent,
152
            mip_level_count: 1,
153
            sample_count: 1,
154
            dimension: wgpu::TextureDimension::D2,
155
            format: wgpu::TextureFormat::Rgba8Unorm,
156
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
157
            view_formats: &[],
158
        });
159
 
160
        let texture_view = texture.create_view(&wgpu::TextureViewDescriptor::default());
161
 
162
        // Create depth texture
163
        let depth_texture = self.device.create_texture(&wgpu::TextureDescriptor {
164
            label: Some("Depth Texture"),
165
            size: texture_extent,
166
            mip_level_count: 1,
167
            sample_count: 1,
168
            dimension: wgpu::TextureDimension::D2,
169
            format: wgpu::TextureFormat::Depth32Float,
170
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
171
            view_formats: &[],
172
        });
173
 
174
        let depth_view = depth_texture.create_view(&wgpu::TextureViewDescriptor::default());
175
 
176
        // Create vertex buffer
177
        let vertex_buffer = self
178
            .device
179
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
180
                label: Some("Vertex Buffer"),
181
                contents: bytemuck::cast_slice(&mesh.vertices),
182
                usage: wgpu::BufferUsages::VERTEX,
183
            });
184
 
185
        // Create index buffer
186
        let index_buffer = self
187
            .device
188
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
189
                label: Some("Index Buffer"),
190
                contents: bytemuck::cast_slice(&mesh.indices),
191
                usage: wgpu::BufferUsages::INDEX,
192
            });
193
 
194
        // Create uniform buffer with camera and model matrices
195
        let view_proj = camera.get_view_proj_matrix();
196
 
197
        // Identity matrix for the model transform - can be changed for more complex scenarios
198
        let model = glam::Mat4::IDENTITY;
199
 
200
        let uniforms = Uniforms {
201
            view_proj: view_proj.to_cols_array_2d(),
202
            model: model.to_cols_array_2d(),
203
        };
204
 
205
        let uniform_buffer = self
206
            .device
207
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
208
                label: Some("Uniform Buffer"),
209
                contents: bytemuck::cast_slice(&[uniforms]),
210
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
211
            });
212
 
213
        // Create bind group for uniforms
214
        let uniform_bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
215
            label: Some("Uniform Bind Group"),
216
            layout: &self.uniform_bind_group_layout,
217
            entries: &[wgpu::BindGroupEntry {
218
                binding: 0,
219
                resource: uniform_buffer.as_entire_binding(),
220
            }],
221
        });
222
 
223
        // Create output buffer to copy the rendered texture
224
        let output_buffer_size = wgpu::BufferAddress::from(self.width * self.height * 4);
225
        let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
226
            label: Some("Output Buffer"),
227
            size: output_buffer_size,
228
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
229
            mapped_at_creation: false,
230
        });
231
 
232
        // Start rendering
233
        let mut encoder = self
234
            .device
235
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
236
                label: Some("Render Encoder"),
237
            });
238
 
239
        {
240
            let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
241
                label: Some("Render Pass"),
242
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
243
                    view: &texture_view,
244
                    resolve_target: None,
245
                    ops: wgpu::Operations {
246
                        load: wgpu::LoadOp::Clear(wgpu::Color {
247
                            r: 0.5, // Light gray for better visibility
248
                            g: 0.5,
249
                            b: 0.5,
250
                            a: 1.0,
251
                        }),
252
                        store: wgpu::StoreOp::Store,
253
                    },
254
                })],
255
                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
256
                    view: &depth_view,
257
                    depth_ops: Some(wgpu::Operations {
258
                        load: wgpu::LoadOp::Clear(1.0),
259
                        store: wgpu::StoreOp::Store,
260
                    }),
261
                    stencil_ops: None,
262
                }),
263
                occlusion_query_set: None,
264
                timestamp_writes: None,
265
            });
266
 
267
            render_pass.set_pipeline(&self.pipeline);
268
            render_pass.set_bind_group(0, &uniform_bind_group, &[]);
269
            render_pass.set_vertex_buffer(0, vertex_buffer.slice(..));
270
            render_pass.set_index_buffer(index_buffer.slice(..), wgpu::IndexFormat::Uint32);
271
            render_pass.draw_indexed(0..u32::try_from(mesh.indices.len()).unwrap(), 0, 0..1);
272
        }
273
 
274
        // Copy the texture to the output buffer
275
        encoder.copy_texture_to_buffer(
276
            wgpu::TexelCopyTextureInfo {
277
                texture: &texture,
278
                mip_level: 0,
279
                origin: wgpu::Origin3d::ZERO,
280
                aspect: wgpu::TextureAspect::All,
281
            },
282
            wgpu::TexelCopyBufferInfo {
283
                buffer: &output_buffer,
284
                layout: wgpu::TexelCopyBufferLayout {
285
                    offset: 0,
286
                    bytes_per_row: Some(self.width * 4),
287
                    rows_per_image: Some(self.height),
288
                },
289
            },
290
            texture_extent,
291
        );
292
 
293
        // Submit the command buffer and await completion
294
        self.queue.submit(std::iter::once(encoder.finish()));
295
 
296
        // Read the output buffer
297
        let buffer_slice = output_buffer.slice(..);
298
        let (tx, rx) = std::sync::mpsc::channel();
299
        buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
300
            tx.send(result).unwrap();
301
        });
302
        self.device.poll(wgpu::Maintain::Wait);
303
 
304
        rx.recv().unwrap()?;
305
 
306
        // Get the buffer data
307
        let data = buffer_slice.get_mapped_range();
308
        let result = data.to_vec();
309
        drop(data);
310
        output_buffer.unmap();
311
 
312
        Ok(result)
313
    }
314
}