Skip to main content

ursus_materials/pack/
aos.rs

1use crate::field::{FieldDesc, write_field_bytes};
2use crate::layout::MaterialLayout;
3use crate::material::MaterialHandle;
4use crate::strategy::ShadingStrategy;
5use crate::value::MaterialValue;
6use std::collections::HashMap;
7
8/// Computes std430-compatible byte offsets for `fields`, in order, and the
9/// total size of the packed struct, rounded up to a 16-byte multiple as
10/// std430 requires for a struct used as an SSBO array element.
11pub fn compute_layout(fields: &[FieldDesc]) -> (Vec<usize>, usize) {
12    let mut offsets = Vec::with_capacity(fields.len());
13    let mut cursor = 0usize;
14
15    for field in fields {
16        let align = field.ty.std430_align();
17        cursor = cursor.next_multiple_of(align);
18        offsets.push(cursor);
19        cursor += field.ty.size();
20    }
21
22    let total = cursor.next_multiple_of(16);
23    (offsets, total)
24}
25
26/// Packs `values` (in the same order as `fields`) into a single contiguous
27/// byte buffer using std430 layout rules (array-of-structs).
28///
29/// Returns `None` if `values.len() != fields.len()` or if any value's type
30/// doesn't match its field's declared type.
31pub fn pack(fields: &[FieldDesc], values: &[MaterialValue]) -> Option<Vec<u8>> {
32    if fields.len() != values.len() {
33        return None;
34    }
35
36    let (offsets, total_size) = compute_layout(fields);
37    let mut bytes = vec![0u8; total_size];
38
39    for ((field, value), &offset) in fields.iter().zip(values.iter()).zip(offsets.iter()) {
40        let end = offset + field.ty.size();
41        write_field_bytes(field.ty, value, &mut bytes[offset..end])?;
42    }
43
44    Some(bytes)
45}
46
47/// Per-stride bucket of packed material bytes: every material sharing a
48/// strategy with the same packed size lands in the same bucket, indexed by
49/// its position in `order`.
50#[derive(Default)]
51struct Bucket {
52    stride: usize,
53    order: Vec<MaterialHandle>,
54    index_of: HashMap<MaterialHandle, usize>,
55    bytes: Vec<u8>,
56}
57
58impl Bucket {
59    fn new(stride: usize) -> Self {
60        Self { stride, order: Vec::new(), index_of: HashMap::new(), bytes: Vec::new() }
61    }
62
63    // TODO: if a material switches strategy to one with a different stride,
64    // its old slot in this bucket is never reclaimed - the bytes stay
65    // allocated and orphaned.
66    fn write(&mut self, handle: MaterialHandle, data: &[u8]) {
67        debug_assert_eq!(data.len(), self.stride);
68
69        let index = *self.index_of.entry(handle).or_insert_with(|| {
70            let i = self.order.len();
71            self.order.push(handle);
72            self.bytes.resize(self.bytes.len() + self.stride, 0);
73            i
74        });
75
76        let start = index * self.stride;
77        self.bytes[start..start + self.stride].copy_from_slice(data);
78    }
79}
80
81/// `MaterialLayout` implementation that packs each material into a single
82/// contiguous byte blob (array-of-structs) via `pack()`, grouping materials
83/// by their packed size into separate buckets - materials using different
84/// strategies with different field sets don't have to share a stride.
85///
86/// Holds plain bytes in memory; turning those bytes into an actual GPU
87/// buffer (e.g. a Vulkan SSBO) is the caller's responsibility.
88#[derive(Default)]
89pub struct AosLayout {
90    buckets: HashMap<usize, Bucket>,
91}
92
93impl AosLayout {
94    pub fn new() -> Self {
95        Self::default()
96    }
97
98    /// Raw bytes currently stored for the bucket at `stride`, if any
99    /// material has been uploaded at that size yet.
100    pub fn bucket_bytes(&self, stride: usize) -> Option<&[u8]> {
101        self.buckets.get(&stride).map(|b| b.bytes.as_slice())
102    }
103
104    /// Index of `handle` within its stride's bucket - this is the
105    /// `material_id` a shader would use to index into that bucket's buffer.
106    pub fn index_of(&self, handle: MaterialHandle, stride: usize) -> Option<usize> {
107        self.buckets.get(&stride)?.index_of.get(&handle).copied()
108    }
109
110    /// All strides currently in use, for callers that need to allocate one
111    /// GPU buffer per stride.
112    pub fn strides(&self) -> impl Iterator<Item = usize> + '_ {
113        self.buckets.keys().copied()
114    }
115}
116
117impl MaterialLayout for AosLayout {
118    fn register_strategy(&mut self, strategy: &dyn ShadingStrategy) {
119        let (_, stride) = compute_layout(strategy.fields());
120        self.buckets.entry(stride).or_insert_with(|| Bucket::new(stride));
121    }
122
123    fn upload(&mut self, handle: MaterialHandle, strategy: &dyn ShadingStrategy, values: &[MaterialValue]) {
124        let Some(bytes) = pack(strategy.fields(), values) else {
125            log::error!(
126                "AosLayout: failed to pack material {:?} for strategy '{}' - value count or types don't match fields",
127                handle,
128                strategy.name()
129            );
130            return;
131        };
132
133        let stride = bytes.len();
134        let bucket = self.buckets.entry(stride).or_insert_with(|| Bucket::new(stride));
135        bucket.write(handle, &bytes);
136    }
137}