lyra-engine/crates/lyra-game/src/render/graph/render_target.rs

use std::sync::Arc;

use tracing::debug;

use crate::math;

enum RenderTargetInner {
    Surface {
        /// The surface that will be rendered to.
        /// 
        /// You can create a new surface with a `'static` lifetime if you have an `Arc<Window>`:
        /// ```nobuild
        /// let window = Arc::new(window);
        /// let surface = instance.create_surface(Arc::clone(&window))?;
        /// ```
        surface: wgpu::Surface<'static>,
        /// the configuration of the surface render target..
        config: wgpu::SurfaceConfiguration,
    },
    Texture {
        /// The texture that will be rendered to.
        texture: Arc<wgpu::Texture>,
    }
}

/// A render target that is a surface or a texture.
#[repr(transparent)]
pub struct RenderTarget(RenderTargetInner);

impl From<wgpu::Texture> for RenderTarget {
    fn from(value: wgpu::Texture) -> Self {
        Self(RenderTargetInner::Texture { texture: Arc::new(value) })
    }
}

impl RenderTarget {
    pub fn from_surface(surface: wgpu::Surface<'static>, config: wgpu::SurfaceConfiguration) -> Self {
        Self(RenderTargetInner::Surface { surface, config })
    }

    pub fn new_texture(device: &wgpu::Device, format: wgpu::TextureFormat, size: math::UVec2) -> Self {
        let tex = device.create_texture(&wgpu::TextureDescriptor {
            label: None,
            size: wgpu::Extent3d {
                width: size.x,
                height: size.y,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_DST | wgpu::TextureUsages::COPY_SRC,
            view_formats: &[],
        });

        Self(RenderTargetInner::Texture { texture: Arc::new(tex) })
    }

    pub fn format(&self) -> wgpu::TextureFormat {
        match &self.0 {
            RenderTargetInner::Surface { config, .. } => config.format,
            RenderTargetInner::Texture { texture } => texture.format(),
        }
    }

    pub fn size(&self) -> math::UVec2 {
        match &self.0 {
            RenderTargetInner::Surface { config, .. } => math::UVec2::new(config.width, config.height),
            RenderTargetInner::Texture { texture } => {
                let s = texture.size();
                math::UVec2::new(s.width, s.height)
            },
        }
    }

    /// Get the frame texture of the [`RenderTarget`]
    /// 
    /// If this is target is a surface and the frame texture was already retrieved from the swap
    /// chain, a [`wgpu::SurfaceError`] error will be returned.
    pub fn frame_texture(&self) -> Result<FrameTexture, wgpu::SurfaceError> {
        match &self.0 {
            RenderTargetInner::Surface { surface, .. } => Ok(FrameTexture::Surface(surface.get_current_texture()?)),
            RenderTargetInner::Texture { texture } => Ok(FrameTexture::Texture(texture.clone())),
        }
    }

    pub fn resize(&mut self, device: &wgpu::Device, new_size: math::UVec2) {
        match &mut self.0 {
            RenderTargetInner::Surface { surface, config } => {
                config.width = new_size.x;
                config.height = new_size.y;
                surface.configure(device, config);
            },
            RenderTargetInner::Texture { texture } => {
                let format = texture.format();
                let size = self.size();

                *self = Self::new_texture(device, format, size);
            },
        }
    }

    /// Create the frame of the RenderTarget.
    /// 
    /// If this is target is a surface and the frame texture was already retrieved from the
    /// swap chain, a [`wgpu::SurfaceError`] error will be returned.
    pub fn create_frame(&self) -> Frame {
        let texture = self.frame_texture()
            .expect("failed to create frame texture"); // TODO: should be returned to the user
        let size = self.size();

        Frame {
            size,
            texture,
        }
    }
}

pub enum FrameTexture {
    Surface(wgpu::SurfaceTexture),
    Texture(Arc<wgpu::Texture>),
}

/// Represents the current frame that is being rendered to.
//#[allow(dead_code)]
pub struct Frame {
    pub(crate) size: math::UVec2,
    pub(crate) texture: FrameTexture,
}

impl Frame {
    pub fn texture(&self) -> &wgpu::Texture {
        match &self.texture {
            FrameTexture::Surface(s) => &s.texture,
            FrameTexture::Texture(t) => t,
        }
    }

    /// Present the frame
    /// 
    /// If this frame is from a surface, it will be present, else nothing will happen.
    pub fn present(self) {
        match self.texture {
            FrameTexture::Surface(s) => s.present(),
            FrameTexture::Texture(_) => {},
        }
    }

    /// The size of the frame
    pub fn size(&self) -> math::UVec2 {
        self.size
    }
}

/// Stores the current frame, and the render target it came from.
pub struct FrameTarget {
    pub render_target: RenderTarget,
    /// None when a frame has not been created yet
    pub frame: Option<Frame>,
    /// The view to use to render to the frame.
    pub frame_view: Option<wgpu::TextureView>,
}

impl FrameTarget {
    pub fn new(render_target: RenderTarget) -> Self {
        Self {
            render_target,
            frame: None,
            frame_view: None,
        }
    }

    /// Returns the size of the [`RenderTarget`].
    pub fn size(&self) -> math::UVec2 {
        self.render_target.size()
    }

    /// Returns the [`wgpu::TextureFormat`] of the [`RenderTarget`].
    pub fn format(&self) -> wgpu::TextureFormat {
        self.render_target.format()
    }
    
    /// Create the frame using the inner [`RenderTarget`].
    pub fn create_frame(&mut self) -> &mut Frame {
        self.frame = Some(self.render_target.create_frame());
        self.frame.as_mut().unwrap()
    }

    /// Create the [`wgpu::TextureView`] for the [`Frame`], storing it in self and returning a reference to it.
    pub fn create_frame_view(&mut self) -> &wgpu::TextureView {
        let frame = self.frame.as_ref().expect("frame was not created, cannot create view");

        self.frame_view = Some(frame.texture().create_view(&wgpu::TextureViewDescriptor::default()));
        self.frame_view.as_ref().unwrap()
    }
}

pub struct TargetViewChain<'a> {
    pub source: &'a mut FrameTarget,
    pub dest: &'a mut FrameTarget,
}

struct ViewChain {
    source: FrameTarget,
    dest: FrameTarget,
    /// tracks the target that is currently being presented
    active: u8,
}

impl ViewChain {
    /// Returns the currently active [`FrameTarget`].
    fn active(&self) -> &FrameTarget {
        if self.active == 0 {
            &self.source
        } else if self.active == 1 {
            &self.dest
        } else {
            panic!("active chain index became invalid! ({})", self.active);
        }
    }
}

pub struct ViewTarget {
    device: Arc<wgpu::Device>,
    /// The primary RenderTarget, likely a Surface
    pub primary: FrameTarget,
    chain: Option<ViewChain>,
}

impl ViewTarget {
    pub fn new(device: Arc<wgpu::Device>, primary: RenderTarget) -> Self {
        let mut s = Self {
            device,
            primary: FrameTarget::new(primary),
            chain: None,
        };

        s.create_chain(s.primary.format(), s.primary.size());
        s
    }

    /// Returns the size of the target.
    pub fn size(&self) -> math::UVec2 {
        self.primary.size()
    }

    /// Returns the [`wgpu::TextureFormat`]
    pub fn format(&self) -> wgpu::TextureFormat {
        self.primary.format()
    }

    /// Resize all the targets, causes the chain to be recreated.
    pub fn resize(&mut self, device: &wgpu::Device, size: math::UVec2) {
        if size != self.primary.size() {
            self.primary.render_target.resize(device, size);
            self.create_chain(self.primary.format(), size);
        }
    }

    fn create_chain(&mut self, format: wgpu::TextureFormat, size: math::UVec2) {
        debug!("Creating chain with {:?} format and {:?} size", format, size);
        
        let mut source = FrameTarget::new(RenderTarget::new_texture(&self.device, format, size));
        source.create_frame();
        source.create_frame_view();
        
        let mut dest = FrameTarget::new(RenderTarget::new_texture(&self.device, format, size));
        dest.create_frame();
        dest.create_frame_view();

        self.chain = Some(ViewChain {
            source,
            dest,
            active: 0,
        });
    }

    /// Cycle the target view chain, storing it in self, and returning a mutable borrow to it.
    pub fn get_chain(&mut self) -> TargetViewChain {
        let format = self.primary.format();
        let size = self.primary.size();

        if let Some(chain) = &self.chain {
            // check if the chain needs to be recreated
            if chain.source.format() != format || chain.source.size() != size {
                self.create_chain(format, size);
            }
        } else {
            self.create_chain(format, size);
        }

        let chain = self.chain.as_mut().unwrap();

        if chain.active == 0 {
            chain.active = 1;
            TargetViewChain {
                source: &mut chain.source,
                dest: &mut chain.dest,
            }
        } else if chain.active == 1 {
            chain.active = 0;
            TargetViewChain {
                source: &mut chain.dest,
                dest: &mut chain.source,
            }
        } else {
            panic!("active chain index became invalid! ({})", chain.active);
        }
    }

    /// Get the [`wgpu::TextureView`] to render to.
    pub fn render_view(&self) -> &wgpu::TextureView {
        let chain = self.chain.as_ref().unwrap();
        chain.active().frame_view.as_ref().unwrap()
    }

    /// Copy the chain target to the primary target
    /// 
    /// The primary target must have `wgpu::TextureUsages::COPY_DST`. This also resets the active
    /// chain texture.
    pub fn copy_to_primary(&mut self, encoder: &mut wgpu::CommandEncoder) {
        let chain = self.chain.as_mut().unwrap();
        let active_tex = chain.active().frame.as_ref().unwrap().texture();

        let active_copy = wgpu::ImageCopyTexture {
            texture: active_tex,
            mip_level: 0,
            origin: wgpu::Origin3d::ZERO,
            aspect: wgpu::TextureAspect::All,
        };

        let dest_tex = self.primary.frame.as_ref().unwrap().texture();
        let dest_copy = wgpu::ImageCopyTexture {
            texture: dest_tex,
            mip_level: 0,
            origin: wgpu::Origin3d::ZERO,
            aspect: wgpu::TextureAspect::All,
        };

        let size = self.primary.size();
        let size = wgpu::Extent3d {
            width: size.x,
            height: size.y,
            depth_or_array_layers: 1,
        };

        encoder.copy_texture_to_texture(active_copy, dest_copy, size);
        
        // reset active texture after a render
        // must get the chain again because of the borrow checker
        let chain = self.chain.as_mut().unwrap();
        chain.active = 0;
    }
}