cyborg/editor/src/render.rs

use crate::model;
use crate::ui::{ObjectWidget, ViewportWidget};
use crate::wgpu;
use cyborg::camera::Camera;
use cyborg::storage::mesh::MeshHandle;
use cyborg::viewport::{Viewport, ViewportInfo, ViewportViews};
use std::sync::Arc;

pub struct OffscreenTextures {
    pub width: u32,
    pub height: u32,
    pub output_texture: wgpu::Texture,
    pub output_view: wgpu::TextureView,
    pub depth_texture: wgpu::Texture,
    pub depth_view: wgpu::TextureView,
}

impl OffscreenTextures {
    pub const DEPTH_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Depth32Float;

    pub fn new(
        device: &wgpu::Device,
        width: u32,
        height: u32,
        output_format: wgpu::TextureFormat,
    ) -> Self {
        let size = wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        };

        let mut tex_desc = wgpu::TextureDescriptor {
            label: None,
            size,
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: output_format,
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
        };

        tex_desc.label = Some("Offscreen Output Texture");
        let output_texture = device.create_texture(&tex_desc);

        tex_desc.label = Some("Offscreen Depth Texture");
        tex_desc.format = Self::DEPTH_FORMAT;
        let depth_texture = device.create_texture(&tex_desc);

        let view_desc = wgpu::TextureViewDescriptor::default();
        let output_view = output_texture.create_view(&view_desc);
        let depth_view = depth_texture.create_view(&view_desc);

        Self {
            width,
            height,
            output_texture,
            output_view,
            depth_texture,
            depth_view,
        }
    }
}

pub struct OffscreenViewport {
    pub device: Arc<wgpu::Device>,
    pub queue: Arc<wgpu::Queue>,
    pub output_format: wgpu::TextureFormat,
    pub textures: OffscreenTextures,
    pub egui_texture: egui::TextureId,
    pub camera: Camera,
}

impl OffscreenViewport {
    pub fn new(
        device: Arc<wgpu::Device>,
        queue: Arc<wgpu::Queue>,
        output_format: wgpu::TextureFormat,
        render_pass: &mut egui_wgpu_backend::RenderPass,
    ) -> Self {
        let textures = OffscreenTextures::new(&device, 640, 480, output_format);

        let egui_texture = render_pass.egui_texture_from_wgpu_texture(
            &device,
            &textures.output_view,
            wgpu::FilterMode::Nearest,
        );

        Self {
            device,
            queue,
            output_format,
            textures,
            egui_texture,
            camera: Camera {
                eye: [1.0, 1.0, 1.0, 0.0],
                vp: Default::default(),
            },
        }
    }

    pub fn get_camera(&self) -> Camera {
        self.camera.clone()
    }
}

impl Viewport for OffscreenViewport {
    fn get_info(&self) -> ViewportInfo {
        ViewportInfo {
            output_format: self.output_format,
            depth_format: OffscreenTextures::DEPTH_FORMAT,
        }
    }

    fn get_queue(&self) -> &wgpu::Queue {
        &self.queue
    }

    fn get_views(&self) -> ViewportViews {
        ViewportViews {
            output: &self.textures.output_view,
            depth: &self.textures.depth_view,
        }
    }
}

pub struct ViewportStore {
    pub device: Arc<wgpu::Device>,
    pub output_format: wgpu::TextureFormat,
    pub viewport: OffscreenViewport,
}

impl ViewportStore {
    pub fn new(
        device: Arc<wgpu::Device>,
        queue: Arc<wgpu::Queue>,
        output_format: wgpu::TextureFormat,
        render_pass: &mut egui_wgpu_backend::RenderPass,
    ) -> Self {
        Self {
            device: device.clone(),
            output_format,
            viewport: OffscreenViewport::new(device, queue, output_format, render_pass),
        }
    }
}

impl cyborg::legion::RenderCallbacks for ViewportStore {
    fn get_viewports(&mut self) -> Vec<(&dyn Viewport, Camera)> {
        vec![(&self.viewport, self.viewport.get_camera())]
    }

    fn present(&mut self) {}
}

pub struct RenderState {
    pub world: legion::World,
    pub resources: legion::Resources,
    pub render_schedule: legion::Schedule,
}

impl RenderState {
    pub fn new(
        device: Arc<wgpu::Device>,
        queue: Arc<wgpu::Queue>,
        output_format: wgpu::TextureFormat,
        render_pass: &mut egui_wgpu_backend::RenderPass,
    ) -> Self {
        use cyborg::shader::{ShaderStore, ShaderWatcher};

        let mut world = legion::World::default();
        let mut resources = legion::Resources::default();
        let viewport_store =
            ViewportStore::new(device.clone(), queue.clone(), output_format, render_pass);

        let renderer = cyborg::Renderer::new(device.clone(), queue.clone());
        resources.insert(renderer);

        let viewport_info = ViewportInfo {
            output_format,
            depth_format: OffscreenTextures::DEPTH_FORMAT,
        };
        resources.insert(viewport_info);

        let shader_store = Arc::new(ShaderStore::new(device.clone()));
        let shaders_dir = std::env::current_dir().unwrap();
        let shaders_dir = shaders_dir.join("shaders/");
        let shader_watcher = ShaderWatcher::new(shader_store.to_owned(), shaders_dir).unwrap();

        let mesh_forward = shader_watcher.add_file("mesh_forward.wgsl").unwrap();
        let mesh_skinning = shader_watcher.add_file("mesh_skinning.wgsl").unwrap();

        let mesh_shaders = cyborg::pass::mesh::ShaderInfo {
            store: shader_store.clone(),
            forward: mesh_forward,
            skinning: mesh_skinning,
        };

        resources.insert(mesh_shaders);

        let mut render_schedule = legion::Schedule::builder();
        cyborg::legion::build_renderer(viewport_store, &mut resources, &mut render_schedule);
        let render_schedule = render_schedule.build();

        world.push((
            cyborg::scene::Transform {
                transform: Default::default(),
            },
            cyborg::scene::DebugDrawList {
                vertices: vec![
                    cyborg::scene::DebugVertex {
                        color: [1.0, 0.0, 0.0],
                        position: [0.0, 0.0, -1.0],
                    },
                    cyborg::scene::DebugVertex {
                        color: [0.0, 0.0, 1.0],
                        position: [0.0, 0.0, 1.0],
                    },
                ],
                indices: vec![0, 1],
            },
        ));

        Self {
            world,
            resources,
            render_schedule,
        }
    }

    pub fn update_viewport(
        &mut self,
        egui_rp: &mut egui_wgpu_backend::RenderPass,
        widget: &mut ViewportWidget,
    ) {
        let mut store = self.resources.get_mut::<ViewportStore>().unwrap();
        let mut viewport = &mut store.viewport;
        viewport.camera = widget.flycam.get_camera();

        if viewport.textures.width != widget.width || viewport.textures.height != widget.height {
            viewport.textures = OffscreenTextures::new(
                &viewport.device,
                widget.width,
                widget.height,
                viewport.output_format,
            );

            egui_rp
                .update_egui_texture_from_wgpu_texture(
                    &viewport.device,
                    &viewport.textures.output_view,
                    wgpu::FilterMode::Nearest,
                    viewport.egui_texture,
                )
                .unwrap();
        }
    }

    pub fn render(&mut self) {
        self.render_schedule
            .execute(&mut self.world, &mut self.resources);
    }

    pub fn load_model(&mut self, model: &model::Model) -> Vec<ObjectWidget> {
        // TODO use model name?
        let mut objects = Vec::new();
        for object in model.objects.iter() {
            objects.push(self.load_object(object));
        }

        objects
    }

    pub fn load_object(&mut self, object: &model::Object) -> ObjectWidget {
        let transform = cyborg::scene::Transform {
            transform: object.transform.to_mat4(),
        };

        let mut entities = Vec::new();
        for mesh in object.meshes.iter() {
            let mesh = cyborg::scene::Mesh {
                mesh: self.load_mesh(mesh),
            };

            let entity = self.world.push((mesh, transform.clone()));
            entities.push(entity);
        }

        let mut children = Vec::new();
        for child in object.children.iter() {
            children.push(self.load_object(child));
        }

        ObjectWidget {
            name: object.name.clone(),
            transform: object.transform.clone(),
            entities,
            children,
            dirty: false,
            children_dirty: false,
        }
    }

    pub fn load_mesh(&mut self, data: &model::Mesh) -> MeshHandle {
        use cyborg::pass::{
            mesh::{MeshPass, Vertex},
            RenderPassBox,
        };
        use cyborg::storage::mesh::{AttrBuffer, MeshBuffer};

        let mesh_pass = self.resources.get::<RenderPassBox<MeshPass>>().unwrap();
        let attributes = mesh_pass.get_attributes();

        let vertices: Vec<_> = data
            .vertices
            .iter()
            .map(|v| Vertex {
                position: v.position.to_array(),
                tan_frame: 0, // TODO encode tangents
            })
            .collect();

        let vertices = AttrBuffer {
            id: attributes.vertex,
            count: vertices.len(),
            data: bytemuck::cast_slice(&vertices).to_vec(),
        };

        let indices = AttrBuffer {
            id: attributes.index,
            count: data.indices.len(),
            data: bytemuck::cast_slice(&data.indices).to_vec(),
        };

        let mut mesh = MeshBuffer::default();
        mesh.attributes.push(vertices);
        mesh.attributes.push(indices);

        mesh_pass.get_mesh_pool().load(mesh).unwrap()
    }
}