1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163

//! Container for an untyped blob of data.

use core::mem;

#[spirv(buffer_load_intrinsic)]
// HACK(eddyb) try to prevent MIR inlining from breaking our intrinsics.
#[inline(never)]
#[spirv_std_macros::gpu_only]
unsafe fn buffer_load_intrinsic<T>(
    buffer: &[u32],
    // FIXME(eddyb) should be `usize`.
    offset: u32,
) -> T {
    // NOTE(eddyb) this doesn't work with `rustc_codegen_spirv` and is only here
    // for explanatory purposes, and to cause some kind of verbose error if
    // `#[spirv(buffer_load_intrinsic)]` fails to replace calls to this function.
    buffer
        .as_ptr()
        .cast::<u8>()
        .add(offset as usize)
        .cast::<T>()
        .read()
}

#[spirv(buffer_store_intrinsic)]
// HACK(eddyb) try to prevent MIR inlining from breaking our intrinsics.
#[inline(never)]
#[spirv_std_macros::gpu_only]
unsafe fn buffer_store_intrinsic<T>(
    buffer: &mut [u32],
    // FIXME(eddyb) should be `usize`.
    offset: u32,
    value: T,
) {
    // NOTE(eddyb) this doesn't work with `rustc_codegen_spirv` and is only here
    // for explanatory purposes, and to cause some kind of verbose error if
    // `#[spirv(buffer_store_intrinsic)]` fails to replace calls to this function.
    buffer
        .as_mut_ptr()
        .cast::<u8>()
        .add(offset as usize)
        .cast::<T>()
        .write(value);
}

/// `ByteAddressableBuffer` is a view to an untyped blob of data, allowing
/// loads and stores of arbitrary basic data types at arbitrary indices.
///
/// # Alignment
/// All data must be aligned to size 4, each element within the data (e.g.
/// struct fields) must have a size and alignment of a multiple of 4, and the
/// `byte_index` passed to load and store must be a multiple of 4. Technically
/// it is not a *byte* addressable buffer, but rather a *word* buffer, but this
/// naming and behavior was inherited from HLSL (where it's UB to pass in an
/// index not a multiple of 4).
///
/// # Safety
/// Using these functions allows reading a different type from the buffer than
/// was originally written (by a previous `store()` or the host API), allowing
/// all sorts of safety guarantees to be bypassed, making it effectively a
/// transmute.
#[repr(transparent)]
pub struct ByteAddressableBuffer<T> {
    /// The underlying array of bytes, able to be directly accessed.
    pub data: T,
}

fn bounds_check<T>(data: &[u32], byte_index: u32) {
    let sizeof = mem::size_of::<T>() as u32;
    if byte_index % 4 != 0 {
        panic!("`byte_index` should be a multiple of 4");
    }
    if byte_index + sizeof > data.len() as u32 {
        let last_byte = byte_index + sizeof;
        panic!(
            "index out of bounds: the len is {} but loading {} bytes at `byte_index` {} reads until {} (exclusive)",
            data.len(),
            sizeof,
            byte_index,
            last_byte,
        );
    }
}

impl<'a> ByteAddressableBuffer<&'a [u32]> {
    /// Creates a `ByteAddressableBuffer` from the untyped blob of data.
    #[inline]
    pub fn from_slice(data: &'a [u32]) -> Self {
        Self { data }
    }

    /// Loads an arbitrary type from the buffer. `byte_index` must be a
    /// multiple of 4.
    ///
    /// # Safety
    /// See [`Self`].
    pub unsafe fn load<T>(&self, byte_index: u32) -> T {
        bounds_check::<T>(self.data, byte_index);
        buffer_load_intrinsic(self.data, byte_index)
    }

    /// Loads an arbitrary type from the buffer. `byte_index` must be a
    /// multiple of 4.
    ///
    /// # Safety
    /// See [`Self`]. Additionally, bounds or alignment checking is not performed.
    pub unsafe fn load_unchecked<T>(&self, byte_index: u32) -> T {
        buffer_load_intrinsic(self.data, byte_index)
    }
}

impl<'a> ByteAddressableBuffer<&'a mut [u32]> {
    /// Creates a `ByteAddressableBuffer` from the untyped blob of data.
    #[inline]
    pub fn from_mut_slice(data: &'a mut [u32]) -> Self {
        Self { data }
    }

    /// Create a non-mutable `ByteAddressableBuffer` from this mutable one.
    #[inline]
    pub fn as_ref(&self) -> ByteAddressableBuffer<&[u32]> {
        ByteAddressableBuffer { data: self.data }
    }

    /// Loads an arbitrary type from the buffer. `byte_index` must be a
    /// multiple of 4.
    ///
    /// # Safety
    /// See [`Self`].
    #[inline]
    pub unsafe fn load<T>(&self, byte_index: u32) -> T {
        self.as_ref().load(byte_index)
    }

    /// Loads an arbitrary type from the buffer. `byte_index` must be a
    /// multiple of 4.
    ///
    /// # Safety
    /// See [`Self`]. Additionally, bounds or alignment checking is not performed.
    #[inline]
    pub unsafe fn load_unchecked<T>(&self, byte_index: u32) -> T {
        self.as_ref().load_unchecked(byte_index)
    }

    /// Stores an arbitrary type into the buffer. `byte_index` must be a
    /// multiple of 4.
    ///
    /// # Safety
    /// See [`Self`].
    pub unsafe fn store<T>(&mut self, byte_index: u32, value: T) {
        bounds_check::<T>(self.data, byte_index);
        buffer_store_intrinsic(self.data, byte_index, value);
    }

    /// Stores an arbitrary type into the buffer. `byte_index` must be a
    /// multiple of 4.
    ///
    /// # Safety
    /// See [`Self`]. Additionally, bounds or alignment checking is not performed.
    pub unsafe fn store_unchecked<T>(&mut self, byte_index: u32, value: T) {
        buffer_store_intrinsic(self.data, byte_index, value);
    }
}