Skip to main content

ursus_core/assets/storage/
material.rs

1use crate::render::gfx::descriptor::{DescriptorAllocator, DescriptorSetDesc};
2use crate::render::gfx::types::{DescriptorSetId, ShaderStage};
3use crate::vulkan::resources::growable_buffer::GrowableBuffer;
4use ash::vk;
5use std::collections::HashMap;
6use ursus_materials::MaterialHandle;
7
8/// Render-thread GPU storage for packed material bytes, one `GrowableBuffer`
9/// per stride (materials whose packed size differs get separate buffers -
10/// see `ursus_materials::pack::aos::AosLayout`, which this store mirrors
11/// on the GPU side).
12///
13/// For V1, exactly one stride is expected to be in use at a time (`mesh.frag`
14/// / `depth_prepass.frag` / `shadow.frag` all declare a single
15/// `set = 1, binding = 0` `MaterialBuffer`), so `descriptor_set` binds the
16/// first stride's buffer to that slot. Supporting multiple simultaneous
17/// strides on the GPU side (multiple bindings, or a technique-driven choice
18/// of which to bind) is future work - see the architecture doc's SoA/
19/// multi-strategy discussion.
20pub struct MaterialStore {
21    buffers: HashMap<usize, GrowableBuffer>,
22    bucket_bytes: HashMap<usize, Vec<u8>>,
23    index_of: HashMap<MaterialHandle, (usize, usize)>, // handle -> (stride, index)
24
25    descriptor_set: DescriptorSetId,
26    bound_stride: Option<usize>,
27
28    device: ash::Device,
29    physical_device: vk::PhysicalDevice,
30    instance: ash::Instance,
31}
32
33impl MaterialStore {
34    pub fn new(
35        device: ash::Device,
36        physical_device: vk::PhysicalDevice,
37        instance: ash::Instance,
38        descriptors: &mut DescriptorAllocator,
39    ) -> anyhow::Result<Self> {
40        let desc = DescriptorSetDesc::new().with_storage_buffer::<()>(0, ShaderStage::Fragment);
41        let descriptor_set = descriptors.create_set(desc)?;
42
43        Ok(Self {
44            buffers: HashMap::new(),
45            bucket_bytes: HashMap::new(),
46            index_of: HashMap::new(),
47            descriptor_set,
48            bound_stride: None,
49            device,
50            physical_device,
51            instance,
52        })
53    }
54
55    pub fn descriptor_set(&self) -> DescriptorSetId {
56        self.descriptor_set
57    }
58
59    /// Writes `bytes` (this material's full packed data, `bytes.len()` is
60    /// its stride) for `handle`, growing/creating the stride's buffer as
61    /// needed, re-uploading that stride's whole bucket to the GPU, and - if
62    /// this is the first stride seen, or the bound stride's buffer just
63    /// reallocated - (re)binding `descriptor_set` to point at it.
64    pub fn upload(
65        &mut self,
66        handle: MaterialHandle,
67        bytes: &[u8],
68        descriptors: &DescriptorAllocator,
69    ) -> anyhow::Result<()> {
70        let stride = bytes.len();
71        let bucket = self.bucket_bytes.entry(stride).or_default();
72
73        let index = match self.index_of.get(&handle) {
74            Some(&(existing_stride, existing_index)) if existing_stride == stride => existing_index,
75            _ => {
76                let index = bucket.len() / stride.max(1);
77                bucket.resize(bucket.len() + stride, 0);
78                self.index_of.insert(handle, (stride, index));
79                index
80            }
81        };
82
83        let start = index * stride;
84        bucket[start..start + stride].copy_from_slice(bytes);
85
86        let buffer = match self.buffers.get_mut(&stride) {
87            Some(b) => b,
88            None => {
89                let new_buffer = GrowableBuffer::new(
90                    &self.device,
91                    self.physical_device,
92                    &self.instance,
93                    vk::BufferUsageFlags::STORAGE_BUFFER,
94                    stride.max(64),
95                )?;
96                self.buffers.entry(stride).or_insert(new_buffer)
97            }
98        };
99
100        let reallocated = buffer.upload(bucket)?;
101
102        if self.bound_stride != Some(stride) || reallocated {
103            self.bind_descriptor(stride, descriptors)?;
104        }
105
106        Ok(())
107    }
108
109    fn bind_descriptor(&mut self, stride: usize, descriptors: &DescriptorAllocator) -> anyhow::Result<()> {
110        if self.bound_stride.is_some() && self.bound_stride != Some(stride) {
111            log::warn!(
112                "MaterialStore: binding stride {stride} over previously bound stride {:?} - \
113                 only one stride can be visible to shaders via the current single MaterialBuffer binding",
114                self.bound_stride
115            );
116        }
117
118        let buffer = self.buffers.get(&stride).expect("stride buffer must exist before binding");
119        descriptors.bind_storage_buffer(
120            self.descriptor_set,
121            0,
122            buffer.buffer(),
123            buffer.capacity() as vk::DeviceSize,
124        )?;
125        self.bound_stride = Some(stride);
126        Ok(())
127    }
128
129    pub fn index_of(&self, handle: MaterialHandle) -> Option<(usize, usize)> {
130        self.index_of.get(&handle).copied()
131    }
132
133    pub fn buffer_for_stride(&self, stride: usize) -> Option<vk::Buffer> {
134        self.buffers.get(&stride).map(|b| b.buffer())
135    }
136
137    pub fn strides(&self) -> impl Iterator<Item = usize> + '_ {
138        self.buffers.keys().copied()
139    }
140}