rafx-framework 0.0.16

use super::dyn_resources;
use super::pipeline_cache;
use super::resource_lookup;
use crate::{
    DescriptorSetAllocatorProvider, DescriptorSetAllocatorRef, DynResourceAllocatorSet,
    GraphicsPipelineCache, MAX_FRAMES_IN_FLIGHT,
};

use crate::graph::RenderGraphCache;
use crate::render_features::RenderRegistry;
use crate::resources::builtin_pipelines::BuiltinPipelines;
use crate::resources::descriptor_sets::DescriptorSetAllocatorManager;
use crate::resources::dyn_commands::DynCommandPoolAllocator;
use crate::resources::dyn_resources::{
    DynResourceAllocatorSetManager, DynResourceAllocatorSetProvider,
};
use crate::resources::resource_lookup::ResourceLookupSet;
use rafx_api::{RafxDeviceContext, RafxResult};
use std::sync::Arc;

//TODO: Support descriptors that can be different per-view
//TODO: Support dynamic descriptors tied to command buffers?
//TODO: Support data inheritance for descriptors

#[derive(Debug)]
pub struct ResourceManagerMetrics {
    pub dyn_resource_metrics: dyn_resources::ResourceMetrics,
    pub resource_metrics: resource_lookup::ResourceMetrics,
    pub graphics_pipeline_cache_metrics: pipeline_cache::GraphicsPipelineCacheMetrics,
}

struct ResourceContextInner {
    descriptor_set_allocator_provider: DescriptorSetAllocatorProvider,
    dyn_resources_allocator_provider: DynResourceAllocatorSetProvider,
    dyn_command_pool_allocator: DynCommandPoolAllocator,
    resources: ResourceLookupSet,
    graphics_pipeline_cache: GraphicsPipelineCache,
    render_graph_cache: RenderGraphCache,
    builtin_pipelines: BuiltinPipelines,
}

#[derive(Clone)]
pub struct ResourceContext {
    inner: Arc<ResourceContextInner>,
}

impl ResourceContext {
    pub fn device_context(&self) -> &RafxDeviceContext {
        self.inner.resources.device_context()
    }

    pub fn resources(&self) -> &ResourceLookupSet {
        &self.inner.resources
    }

    pub fn graphics_pipeline_cache(&self) -> &GraphicsPipelineCache {
        &self.inner.graphics_pipeline_cache
    }

    pub fn render_graph_cache(&self) -> &RenderGraphCache {
        &self.inner.render_graph_cache
    }

    pub fn create_dyn_command_pool_allocator(&self) -> DynCommandPoolAllocator {
        self.inner.dyn_command_pool_allocator.clone()
    }

    pub fn create_dyn_resource_allocator_set(&self) -> DynResourceAllocatorSet {
        self.inner.dyn_resources_allocator_provider.get_allocator()
    }

    pub fn create_descriptor_set_allocator(&self) -> DescriptorSetAllocatorRef {
        self.inner.descriptor_set_allocator_provider.get_allocator()
    }

    pub fn builtin_pipelines(&self) -> &BuiltinPipelines {
        &self.inner.builtin_pipelines
    }
}

pub struct ResourceManager {
    render_registry: RenderRegistry,
    dyn_resource_allocators: DynResourceAllocatorSetManager,
    dyn_command_pool_allocator: DynCommandPoolAllocator,
    resources: ResourceLookupSet,
    render_graph_cache: RenderGraphCache,
    descriptor_set_allocator: DescriptorSetAllocatorManager,
    graphics_pipeline_cache: GraphicsPipelineCache,

    // This is Some() until ResourceManager is dropped
    builtin_pipelines: Option<BuiltinPipelines>,
}

impl ResourceManager {
    pub fn new(
        device_context: &RafxDeviceContext,
        render_registry: &RenderRegistry,
    ) -> Self {
        let resources = ResourceLookupSet::new(device_context, MAX_FRAMES_IN_FLIGHT as u32);
        let builtin_pipelines =
            BuiltinPipelines::new(&resources).expect("Failed to load a built-in resource");

        //DX12TODO: Disable resource reuse for DX12 for now
        // vulkan lets us transition from COMMON but dx12 doesn't. The graph
        // is trying to reuse an image from previous frame transitioning it from COMMON
        let reuse_resources = !device_context.is_dx12();

        ResourceManager {
            render_registry: render_registry.clone(),
            dyn_command_pool_allocator: DynCommandPoolAllocator::new(MAX_FRAMES_IN_FLIGHT as u32),
            dyn_resource_allocators: DynResourceAllocatorSetManager::new(
                device_context,
                MAX_FRAMES_IN_FLIGHT as u32,
            ),
            resources: resources.clone(),
            render_graph_cache: RenderGraphCache::new(MAX_FRAMES_IN_FLIGHT as u32, reuse_resources),
            descriptor_set_allocator: DescriptorSetAllocatorManager::new(device_context),
            graphics_pipeline_cache: GraphicsPipelineCache::new(render_registry, resources),
            builtin_pipelines: Some(builtin_pipelines),
        }
    }

    pub fn device_context(&self) -> &RafxDeviceContext {
        self.resources.device_context()
    }

    pub fn resource_context(&self) -> ResourceContext {
        let inner = ResourceContextInner {
            descriptor_set_allocator_provider: self
                .descriptor_set_allocator
                .create_allocator_provider(),
            dyn_resources_allocator_provider: self
                .dyn_resource_allocators
                .create_allocator_provider(),
            dyn_command_pool_allocator: self.dyn_command_pool_allocator.clone(),
            resources: self.resources.clone(),
            graphics_pipeline_cache: self.graphics_pipeline_cache.clone(),
            render_graph_cache: self.render_graph_cache.clone(),
            builtin_pipelines: self.builtin_pipelines.as_ref().unwrap().clone(),
        };

        ResourceContext {
            inner: Arc::new(inner),
        }
    }

    pub fn resources(&self) -> &ResourceLookupSet {
        &self.resources
    }

    pub fn graphics_pipeline_cache(&self) -> &GraphicsPipelineCache {
        &self.graphics_pipeline_cache
    }

    pub fn dyn_command_pool_allocator(&self) -> &DynCommandPoolAllocator {
        &self.dyn_command_pool_allocator
    }

    pub fn create_dyn_resource_allocator_set(&self) -> DynResourceAllocatorSet {
        self.dyn_resource_allocators.get_allocator()
    }

    pub fn create_dyn_resource_allocator_provider(&self) -> DynResourceAllocatorSetProvider {
        self.dyn_resource_allocators.create_allocator_provider()
    }

    pub fn create_descriptor_set_allocator(&self) -> DescriptorSetAllocatorRef {
        self.descriptor_set_allocator.get_allocator()
    }

    pub fn create_descriptor_set_allocator_provider(&self) -> DescriptorSetAllocatorProvider {
        self.descriptor_set_allocator.create_allocator_provider()
    }

    pub fn render_registry(&self) -> &RenderRegistry {
        &self.render_registry
    }

    pub fn metrics(&self) -> ResourceManagerMetrics {
        let dyn_resource_metrics = self.dyn_resource_allocators.metrics();
        let resource_metrics = self.resources.metrics();
        let graphics_pipeline_cache_metrics = self.graphics_pipeline_cache.metrics();

        ResourceManagerMetrics {
            dyn_resource_metrics,
            resource_metrics,
            graphics_pipeline_cache_metrics,
        }
    }

    #[profiling::function]
    pub fn on_frame_complete(&mut self) -> RafxResult<()> {
        self.render_graph_cache.on_frame_complete();
        self.graphics_pipeline_cache.on_frame_complete();
        self.resources.on_frame_complete()?;
        self.dyn_command_pool_allocator.on_frame_complete()?;
        self.dyn_resource_allocators.on_frame_complete()?;
        self.descriptor_set_allocator.on_frame_complete();
        Ok(())
    }
}

impl Drop for ResourceManager {
    fn drop(&mut self) {
        log::info!("Cleaning up resource manager");
        log::trace!("Resource Manager Metrics:\n{:#?}", self.metrics());

        self.builtin_pipelines = None;

        // Wipe caches to ensure we don't keep anything alive
        self.render_graph_cache.clear();
        self.graphics_pipeline_cache.clear_all_pipelines();

        // Drop all descriptors. These bind to raw resources, so we need to drop them before
        // dropping resources
        self.descriptor_set_allocator.destroy().unwrap();

        // Now drop all resources with a zero ref count and warn for any resources that remain
        self.resources.destroy().unwrap();
        self.dyn_resource_allocators.destroy().unwrap();

        log::info!("Dropping resource manager");
        log::trace!("Resource Manager Metrics:\n{:#?}", self.metrics());
    }
}