1use sha2::{Digest, Sha256};
2use std::path::{Path, PathBuf};
3
4pub const MAGIC: [u8; 8] = *b"PITOPLC1";
5
6#[derive(Clone, Debug, PartialEq, Eq)]
7pub enum CacheLoad {
8 Loaded { bytes: usize },
9 Fresh,
10 Ignored(String),
11 Unsupported,
12}
13
14pub struct PipelineCache {
15 cache: wgpu::PipelineCache,
16 path: PathBuf,
17 key: String,
18}
19
20pub fn cache_key(info: &wgpu::AdapterInfo) -> Option<String> {
21 wgpu::util::pipeline_cache_key(info).map(|backend| {
22 format!(
23 "{backend}|{}|{}|{}|pito-engine {}",
24 info.name,
25 info.driver,
26 info.driver_info,
27 env!("CARGO_PKG_VERSION")
28 )
29 })
30}
31
32pub fn wrap(key: &str, payload: &[u8]) -> Vec<u8> {
33 let mut bytes = Vec::with_capacity(MAGIC.len() + 4 + key.len() + 32 + 8 + payload.len());
34 bytes.extend_from_slice(&MAGIC);
35 bytes.extend_from_slice(&(key.len() as u32).to_le_bytes());
36 bytes.extend_from_slice(key.as_bytes());
37 bytes.extend_from_slice(&Sha256::digest(payload));
38 bytes.extend_from_slice(&(payload.len() as u64).to_le_bytes());
39 bytes.extend_from_slice(payload);
40 bytes
41}
42
43pub fn unwrap<'a>(key: &str, bytes: &'a [u8]) -> Result<&'a [u8], String> {
44 let rest = bytes
45 .strip_prefix(&MAGIC)
46 .ok_or("not a pito-engine pipeline cache")?;
47 let (length, rest) = split(rest, 4)?;
48 let length = u32::from_le_bytes(length.try_into().expect("four bytes")) as usize;
49 let (stored, rest) = split(rest, length)?;
50 if stored != key.as_bytes() {
51 return Err(format!(
52 "made for {}, this is {key}",
53 String::from_utf8_lossy(stored)
54 ));
55 }
56 let (digest, rest) = split(rest, 32)?;
57 let (size, payload) = split(rest, 8)?;
58 let size = u64::from_le_bytes(size.try_into().expect("eight bytes"));
59 if payload.len() as u64 != size {
60 return Err(format!(
61 "payload is {} bytes, header says {size}",
62 payload.len()
63 ));
64 }
65 if Sha256::digest(payload).as_slice() != digest {
66 return Err("payload checksum does not match".into());
67 }
68 Ok(payload)
69}
70
71fn split(bytes: &[u8], at: usize) -> Result<(&[u8], &[u8]), String> {
72 if bytes.len() < at {
73 return Err("truncated".into());
74 }
75 Ok(bytes.split_at(at))
76}
77
78impl PipelineCache {
79 pub fn open(
80 device: &wgpu::Device,
81 info: &wgpu::AdapterInfo,
82 path: &Path,
83 ) -> (Option<Self>, CacheLoad) {
84 let Some(key) =
85 cache_key(info).filter(|_| device.features().contains(wgpu::Features::PIPELINE_CACHE))
86 else {
87 return (None, CacheLoad::Unsupported);
88 };
89 let read = std::fs::read(path);
90 let (data, load) = match &read {
91 Ok(bytes) => match unwrap(&key, bytes) {
92 Ok(payload) => (
93 Some(payload),
94 CacheLoad::Loaded {
95 bytes: payload.len(),
96 },
97 ),
98 Err(reason) => (None, CacheLoad::Ignored(reason)),
99 },
100 Err(error) if error.kind() == std::io::ErrorKind::NotFound => (None, CacheLoad::Fresh),
101 Err(error) => (None, CacheLoad::Ignored(error.to_string())),
102 };
103 let cache = unsafe {
104 device.create_pipeline_cache(&wgpu::PipelineCacheDescriptor {
105 label: Some("live pipeline cache"),
106 data,
107 fallback: true,
108 })
109 };
110 (
111 Some(Self {
112 cache,
113 path: path.to_path_buf(),
114 key,
115 }),
116 load,
117 )
118 }
119
120 pub fn cache(&self) -> &wgpu::PipelineCache {
121 &self.cache
122 }
123
124 pub fn save(&self) -> Result<usize, String> {
125 let payload = self
126 .cache
127 .get_data()
128 .ok_or("the backend gave no pipeline cache data")?;
129 if let Some(parent) = self.path.parent() {
130 std::fs::create_dir_all(parent).map_err(|error| error.to_string())?;
131 }
132 let mut partial = self.path.clone().into_os_string();
133 partial.push(".partial");
134 let partial = PathBuf::from(partial);
135 std::fs::write(&partial, wrap(&self.key, &payload)).map_err(|error| error.to_string())?;
136 std::fs::rename(&partial, &self.path).map_err(|error| error.to_string())?;
137 Ok(payload.len())
138 }
139}
140
141#[cfg(test)]
142mod tests {
143 use super::*;
144
145 const KEY: &str =
146 "wgpu_pipeline_cache_vulkan_4098_30016|AMD Radeon RX 9060 XT|radv|Mesa|pito-engine 0.18.18";
147
148 #[test]
149 fn a_wrapped_payload_unwraps_to_the_same_bytes() {
150 let payload: Vec<u8> = (0..4096u32).map(|i| (i * 31 % 251) as u8).collect();
151 let wrapped = wrap(KEY, &payload);
152 assert_eq!(unwrap(KEY, &wrapped).unwrap(), payload.as_slice());
153 assert_eq!(unwrap(KEY, &wrap(KEY, &[])).unwrap(), &[] as &[u8]);
154 }
155
156 #[test]
157 fn another_adapter_driver_or_engine_version_is_refused() {
158 let wrapped = wrap(KEY, b"pipelines");
159 for other in [
160 KEY.replace("30016", "29695"),
161 KEY.replace("Mesa", "Mesa 26.1"),
162 KEY.replace("0.18.18", "0.18.19"),
163 ] {
164 let refused = unwrap(&other, &wrapped).unwrap_err();
165 assert!(refused.starts_with("made for"), "{refused}");
166 }
167 }
168
169 #[test]
170 fn a_damaged_or_foreign_file_is_refused() {
171 let wrapped = wrap(KEY, b"pipelines of the opaque pass");
172 assert!(unwrap(KEY, b"").is_err());
173 assert!(unwrap(KEY, b"WGPUPLCH and more").is_err());
174 for cut in [4, 12, wrapped.len() - 40, wrapped.len() - 1] {
175 assert!(unwrap(KEY, &wrapped[..cut]).is_err(), "cut at {cut}");
176 }
177 let mut flipped = wrapped.clone();
178 let last = flipped.len() - 1;
179 flipped[last] ^= 1;
180 assert_eq!(
181 unwrap(KEY, &flipped).unwrap_err(),
182 "payload checksum does not match"
183 );
184 let mut longer = wrapped.clone();
185 longer.push(0);
186 assert!(unwrap(KEY, &longer).is_err());
187 }
188}