zenyx-engine-telemetry/engine/src/core/render/ctx.rs

use std::sync::Arc;
use std::time::Instant;
use std::{backtrace::Backtrace, borrow::Cow};

use cgmath::{Matrix4, Point3, Rad, Vector3, perspective};
use futures::executor::block_on;
use thiserror::Error;
use tracing::{error, trace};
use wgpu::TextureUsages;
use wgpu::{Backends, InstanceDescriptor, util::DeviceExt};
use winit::window::Window;
#[derive(Debug, Error)]
#[error(transparent)]
pub enum ContextErrorKind {
    #[error("Surface creation failed")]
    SurfaceCreation,
    #[error("Surface configuration failed")]
    SurfaceConfiguration,
    #[error("Adapter request failed")]
    AdapterRequest,
    #[error("Device request failed")]
    DeviceRequest,
    #[error("Surface texture acquisition failed")]
    SurfaceTexture,
}

#[derive(Debug, Error)]
pub struct RenderContextError {
    kind: ContextErrorKind,
    label: Option<Cow<'static, str>>,
    #[source]
    source: Option<Box<dyn std::error::Error + Send + Sync>>,
}

impl RenderContextError {
    pub fn new(
        kind: ContextErrorKind,
        label: impl Into<Option<Cow<'static, str>>>,
        source: impl Into<Option<Box<dyn std::error::Error + Send + Sync>>>,
    ) -> Self {
        Self {
            kind,
            label: label.into(),
            source: source.into(),
        }
    }

    pub fn with_label(mut self, label: impl Into<Cow<'static, str>>) -> Self {
        self.label = Some(label.into());
        self
    }
}

impl std::fmt::Display for RenderContextError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(label) = &self.label {
            writeln!(f, "[{}] {}", label, self.kind)?;
        } else {
            writeln!(f, "{}", self.kind)?;
        }

        if let Some(source) = &self.source {
            fn fmt_chain(
                err: &(dyn std::error::Error + 'static),
                indent: usize,
                f: &mut std::fmt::Formatter<'_>,
            ) -> std::fmt::Result {
                let indent_str = "  ".repeat(indent);
                writeln!(f, "{}{}", indent_str, err)?;
                if let Some(next) = err.source() {
                    writeln!(f, "{}Caused by:", indent_str)?;
                    fmt_chain(next, indent + 1, f)?;
                }
                Ok(())
            }
            writeln!(f, "Caused by:")?;
            fmt_chain(source.as_ref(), 1, f)?;
        }
        Ok(())
    }
}

trait IntoRenderContextError<T> {
    fn ctx_err(
        self,
        kind: ContextErrorKind,
        label: impl Into<Cow<'static, str>>,
    ) -> Result<T, RenderContextError>;
}

impl<T, E> IntoRenderContextError<T> for Result<T, E>
where
    E: std::error::Error + Send + Sync + 'static,
{
    fn ctx_err(
        self,
        kind: ContextErrorKind,
        label: impl Into<Cow<'static, str>>,
    ) -> Result<T, RenderContextError> {
        self.map_err(|e| {
            RenderContextError::new(
                kind,
                Some(label.into()),
                Some(Box::new(e) as Box<dyn std::error::Error + Send + Sync>),
            )
        })
    }
}

impl From<wgpu::CreateSurfaceError> for RenderContextError {
    fn from(err: wgpu::CreateSurfaceError) -> Self {
        RenderContextError::new(
            ContextErrorKind::SurfaceCreation,
            Some("Surface creation".into()),
            Some(Box::new(err) as Box<dyn std::error::Error + Send + Sync>),
        )
    }
}

impl From<wgpu::RequestDeviceError> for RenderContextError {
    fn from(err: wgpu::RequestDeviceError) -> Self {
        RenderContextError::new(
            ContextErrorKind::DeviceRequest,
            Some("Device setup".into()),
            Some(Box::new(err) as Box<dyn std::error::Error + Send + Sync>),
        )
    }
}

const CUBE_SHADER: &str = r"
struct Uniforms {
    mvp: mat4x4<f32>,
};

@group(0) @binding(0)
var<uniform> u: Uniforms;

struct VertexInput {
    @location(0) position: vec3<f32>,
    @location(1) normal: vec3<f32>,
};

struct VertexOutput {
    @builtin(position) clip_position: vec4<f32>,
    @location(0) normal: vec3<f32>,
};

@vertex
fn vs_main(input: VertexInput) -> VertexOutput {
    var output: VertexOutput;
    output.clip_position = u.mvp * vec4<f32>(input.position, 1.0);
    output.normal = input.normal;
    return output;
}


@fragment
fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> {
    let ambient: f32 = 0.2;
    let light_dir = normalize(vec3<f32>(0.5, 1.0, 0.5));
    let diffuse = clamp(dot(normalize(input.normal), light_dir), 0.0, 1.0);
    // Mix ambient light to ensure no face is completely dark.
    let brightness = ambient + (1.0 - ambient) * diffuse;
    return vec4<f32>(0.7 * brightness, 0.7 * brightness, 0.9 * brightness, 1.0);
}
";
#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct Vertex {
    position: [f32; 3],
    normal: [f32; 3],
}

impl Vertex {
    const ATTRIBS: [wgpu::VertexAttribute; 2] = [
        wgpu::VertexAttribute {
            offset: 0,
            shader_location: 0,
            format: wgpu::VertexFormat::Float32x3,
        },
        wgpu::VertexAttribute {
            offset: std::mem::size_of::<[f32; 3]>() as wgpu::BufferAddress,
            shader_location: 1,
            format: wgpu::VertexFormat::Float32x3,
        },
    ];

    fn desc<'a>() -> wgpu::VertexBufferLayout<'a> {
        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &Self::ATTRIBS,
        }
    }
}

static CUBE_VERTICES: &[Vertex] = &[
    Vertex {
        position: [-0.5, -0.5, 0.5],
        normal: [0.0, 0.0, 1.0],
    },
    Vertex {
        position: [0.5, -0.5, 0.5],
        normal: [0.0, 0.0, 1.0],
    },
    Vertex {
        position: [0.5, 0.5, 0.5],
        normal: [0.0, 0.0, 1.0],
    },
    Vertex {
        position: [0.5, 0.5, 0.5],
        normal: [0.0, 0.0, 1.0],
    },
    Vertex {
        position: [-0.5, 0.5, 0.5],
        normal: [0.0, 0.0, 1.0],
    },
    Vertex {
        position: [-0.5, -0.5, 0.5],
        normal: [0.0, 0.0, 1.0],
    },
    Vertex {
        position: [0.5, -0.5, -0.5],
        normal: [0.0, 0.0, -1.0],
    },
    Vertex {
        position: [-0.5, -0.5, -0.5],
        normal: [0.0, 0.0, -1.0],
    },
    Vertex {
        position: [-0.5, 0.5, -0.5],
        normal: [0.0, 0.0, -1.0],
    },
    Vertex {
        position: [-0.5, 0.5, -0.5],
        normal: [0.0, 0.0, -1.0],
    },
    Vertex {
        position: [0.5, 0.5, -0.5],
        normal: [0.0, 0.0, -1.0],
    },
    Vertex {
        position: [0.5, -0.5, -0.5],
        normal: [0.0, 0.0, -1.0],
    },
    Vertex {
        position: [0.5, -0.5, 0.5],
        normal: [1.0, 0.0, 0.0],
    },
    Vertex {
        position: [0.5, -0.5, -0.5],
        normal: [1.0, 0.0, 0.0],
    },
    Vertex {
        position: [0.5, 0.5, -0.5],
        normal: [1.0, 0.0, 0.0],
    },
    Vertex {
        position: [0.5, 0.5, -0.5],
        normal: [1.0, 0.0, 0.0],
    },
    Vertex {
        position: [0.5, 0.5, 0.5],
        normal: [1.0, 0.0, 0.0],
    },
    Vertex {
        position: [0.5, -0.5, 0.5],
        normal: [1.0, 0.0, 0.0],
    },
    Vertex {
        position: [-0.5, -0.5, -0.5],
        normal: [-1.0, 0.0, 0.0],
    },
    Vertex {
        position: [-0.5, -0.5, 0.5],
        normal: [-1.0, 0.0, 0.0],
    },
    Vertex {
        position: [-0.5, 0.5, 0.5],
        normal: [-1.0, 0.0, 0.0],
    },
    Vertex {
        position: [-0.5, 0.5, 0.5],
        normal: [-1.0, 0.0, 0.0],
    },
    Vertex {
        position: [-0.5, 0.5, -0.5],
        normal: [-1.0, 0.0, 0.0],
    },
    Vertex {
        position: [-0.5, -0.5, -0.5],
        normal: [-1.0, 0.0, 0.0],
    },
    Vertex {
        position: [-0.5, 0.5, 0.5],
        normal: [0.0, 1.0, 0.0],
    },
    Vertex {
        position: [0.5, 0.5, 0.5],
        normal: [0.0, 1.0, 0.0],
    },
    Vertex {
        position: [0.5, 0.5, -0.5],
        normal: [0.0, 1.0, 0.0],
    },
    Vertex {
        position: [0.5, 0.5, -0.5],
        normal: [0.0, 1.0, 0.0],
    },
    Vertex {
        position: [-0.5, 0.5, -0.5],
        normal: [0.0, 1.0, 0.0],
    },
    Vertex {
        position: [-0.5, 0.5, 0.5],
        normal: [0.0, 1.0, 0.0],
    },
    Vertex {
        position: [-0.5, -0.5, -0.5],
        normal: [0.0, -1.0, 0.0],
    },
    Vertex {
        position: [0.5, -0.5, -0.5],
        normal: [0.0, -1.0, 0.0],
    },
    Vertex {
        position: [0.5, -0.5, 0.5],
        normal: [0.0, -1.0, 0.0],
    },
    Vertex {
        position: [0.5, -0.5, 0.5],
        normal: [0.0, -1.0, 0.0],
    },
    Vertex {
        position: [-0.5, -0.5, 0.5],
        normal: [0.0, -1.0, 0.0],
    },
    Vertex {
        position: [-0.5, -0.5, -0.5],
        normal: [0.0, -1.0, 0.0],
    },
];

pub struct WgpuCtx<'window> {
    device: wgpu::Device,
    queue: wgpu::Queue,
    surface: wgpu::Surface<'window>,
    surface_config: wgpu::SurfaceConfiguration,
    adapter: wgpu::Adapter,
    render_pipeline: wgpu::RenderPipeline,
    uniform_buffer: wgpu::Buffer,
    vertex_buffer: wgpu::Buffer,
    start_time: Instant,
    bg_color: wgpu::Color,
    last_frame_instant: Instant,
    frame_count: u32,
}

impl<'window> WgpuCtx<'window> {
    pub async fn new(window: Arc<Window>) -> Result<WgpuCtx<'window>, RenderContextError> {
        let instance = wgpu::Instance::new(&InstanceDescriptor {
            backends: Backends::from_comma_list("dx12,metal,opengl,webgpu"),
            ..Default::default()
        });

        let surface = instance
            .create_surface(Arc::clone(&window))
            .ctx_err(ContextErrorKind::SurfaceCreation, "Surface initialization")?;

        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::default(),
                compatible_surface: Some(&surface),
                ..Default::default()
            })
            .await
            .ok_or_else(|| {
                RenderContextError::new(
                    ContextErrorKind::AdapterRequest,
                    Some("Adapter selection".into()),
                    None,
                )
            })?;
        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor::default(), None)
            .await
            .ctx_err(ContextErrorKind::DeviceRequest, "Device configuration")?;
        let size = window.inner_size();
        let width = size.width.max(1);
        let height = size.height.max(1);
        let surface_config = wgpu::SurfaceConfiguration {
            width: width.max(1),
            height: height.max(1),
            format: wgpu::TextureFormat::Rgba8UnormSrgb,
            present_mode: wgpu::PresentMode::AutoNoVsync,
            alpha_mode: wgpu::CompositeAlphaMode::Auto,
            view_formats: Vec::new(),
            usage: TextureUsages::RENDER_ATTACHMENT,
            desired_maximum_frame_latency: 3,
        };
        surface.configure(&device, &surface_config);
        let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("Uniform Buffer"),
            size: 64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Cube Vertex Buffer"),
            contents: bytemuck::cast_slice(CUBE_VERTICES),
            usage: wgpu::BufferUsages::VERTEX,
        });
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Cube Shader"),
            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(CUBE_SHADER)),
        });
        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("Uniform Bind Group Layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: wgpu::BufferSize::new(64),
                },
                count: None,
            }],
        });
        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Cube Pipeline Layout"),
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[],
        });
        let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("Cube Render Pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[Vertex::desc()],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_config.format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: wgpu::PipelineCompilationOptions::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview: None,
            cache: None,
        });
        Ok(WgpuCtx {
            device,
            queue,
            surface,
            surface_config,
            adapter,
            render_pipeline,
            uniform_buffer,
            vertex_buffer,
            bg_color: wgpu::Color {
                r: 0.1,
                g: 0.1,
                b: 0.1,
                a: 1.0,
            },
            start_time: Instant::now(),
            last_frame_instant: Instant::now(),
            frame_count: 0,
        })
    }

    pub fn new_blocking(window: Arc<Window>) -> Result<WgpuCtx<'window>, RenderContextError> {
        block_on(Self::new(window))
    }

    pub fn resize(&mut self, new_size: (u32, u32)) {
        let (width, height) = new_size;
        self.surface_config.width = width.max(1);
        self.surface_config.height = height.max(1);
        self.surface.configure(&self.device, &self.surface_config);
    }

    pub fn draw(&mut self) {
        let elapsed = self.start_time.elapsed().as_secs_f32() * 0.80f32;
        let model = Matrix4::from_angle_x(Rad(elapsed)) * Matrix4::from_angle_y(Rad(elapsed));
        let view = Matrix4::look_at_rh(
            Point3::new(0.0, 0.0, 3.0),
            Point3::new(0.0, 0.0, 0.0),
            Vector3::unit_y(),
        );
        let aspect = self.surface_config.width as f32 / self.surface_config.height as f32;
        let proj = perspective(Rad(std::f32::consts::FRAC_PI_4), aspect, 0.1, 100.0);
        let mvp = proj * view * model;
        let mvp_array: [[f32; 4]; 4] = [
            [mvp.x.x, mvp.x.y, mvp.x.z, mvp.x.w],
            [mvp.y.x, mvp.y.y, mvp.y.z, mvp.y.w],
            [mvp.z.x, mvp.z.y, mvp.z.z, mvp.z.w],
            [mvp.w.x, mvp.w.y, mvp.w.z, mvp.w.w],
        ];
        self.queue
            .write_buffer(&self.uniform_buffer, 0, bytemuck::bytes_of(&mvp_array));
        let surface_texture = self
            .surface
            .get_current_texture()
            .expect("Failed to get surface texture");
        let view_texture = surface_texture
            .texture
            .create_view(&wgpu::TextureViewDescriptor::default());
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("Cube Command Encoder"),
            });
        {
            let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("Cube Render Pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &view_texture,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(self.bg_color),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
            });
            render_pass.set_pipeline(&self.render_pipeline);
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("Uniform Bind Group"),
                layout: &self.render_pipeline.get_bind_group_layout(0),
                entries: &[wgpu::BindGroupEntry {
                    binding: 0,
                    resource: self.uniform_buffer.as_entire_binding(),
                }],
            });
            render_pass.set_bind_group(0, &bind_group, &[]);
            render_pass.set_vertex_buffer(0, self.vertex_buffer.slice(..));
            render_pass.draw(0..36, 0..1);
        }
        self.queue.submit(Some(encoder.finish()));
        surface_texture.present();

        self.frame_count += 1;
        let elapsed_secs = self.last_frame_instant.elapsed().as_secs_f32();
        if elapsed_secs >= 1.0 {
            let fps = self.frame_count as f32 / elapsed_secs;
            trace!("FPS: {:.2}", fps);
            self.frame_count = 0;
            self.last_frame_instant = Instant::now();
        }
    }

    pub fn change_bg_color(&mut self, color: wgpu::Color) {
        self.bg_color = color;
    }
    pub fn bg_color(&self) -> wgpu::Color {
        self.bg_color.clone()
    }
}