1use crate::{CameraIntrinsics, ObjectRotation, RenderConfig, RenderOutput};
45use bevy::prelude::Transform;
46use std::collections::VecDeque;
47use std::path::PathBuf;
48
49#[derive(Clone, Debug)]
51pub struct BatchRenderConfig {
52 pub max_batch_size: usize,
54 pub frame_timeout_ms: u32,
56 pub enable_depth_readback: bool,
58 pub enable_asset_caching: bool,
60 pub resource_cleanup_interval: u32,
62}
63
64impl Default for BatchRenderConfig {
65 fn default() -> Self {
66 Self {
67 max_batch_size: 256,
68 frame_timeout_ms: 500,
69 enable_depth_readback: true,
70 enable_asset_caching: true,
71 resource_cleanup_interval: 32,
72 }
73 }
74}
75
76#[derive(Clone, Debug)]
78pub struct BatchRenderRequest {
79 pub object_dir: PathBuf,
81 pub viewpoint: Transform,
83 pub object_rotation: ObjectRotation,
85 pub render_config: RenderConfig,
87}
88
89#[derive(Clone, Debug, Copy, PartialEq, Eq)]
91pub enum RenderStatus {
92 Success,
94 PartialFailure,
96 Failed,
98}
99
100#[derive(Clone, Debug)]
102pub struct BatchRenderOutput {
103 pub request: BatchRenderRequest,
105 pub rgba: Vec<u8>,
107 pub depth: Vec<f64>,
109 pub width: u32,
111 pub height: u32,
113 pub intrinsics: CameraIntrinsics,
115 pub status: RenderStatus,
117 pub error_message: Option<String>,
119}
120
121impl BatchRenderOutput {
122 pub fn to_rgb_image(&self) -> Vec<Vec<[u8; 3]>> {
124 let mut image = Vec::with_capacity(self.height as usize);
125 for y in 0..self.height {
126 let mut row = Vec::with_capacity(self.width as usize);
127 for x in 0..self.width {
128 let idx = ((y * self.width + x) * 4) as usize;
129 if idx + 2 < self.rgba.len() {
130 row.push([self.rgba[idx], self.rgba[idx + 1], self.rgba[idx + 2]]);
131 } else {
132 row.push([0, 0, 0]);
133 }
134 }
135 image.push(row);
136 }
137 image
138 }
139
140 pub fn to_depth_image(&self) -> Vec<Vec<f64>> {
142 let mut image = Vec::with_capacity(self.height as usize);
143 for y in 0..self.height {
144 let mut row = Vec::with_capacity(self.width as usize);
145 for x in 0..self.width {
146 let idx = (y * self.width + x) as usize;
147 if idx < self.depth.len() {
148 row.push(self.depth[idx]);
149 } else {
150 row.push(0.0);
151 }
152 }
153 image.push(row);
154 }
155 image
156 }
157
158 pub fn from_render_output(request: BatchRenderRequest, output: RenderOutput) -> Self {
160 Self {
161 request,
162 rgba: output.rgba,
163 depth: output.depth,
164 width: output.width,
165 height: output.height,
166 intrinsics: output.intrinsics,
167 status: RenderStatus::Success,
168 error_message: None,
169 }
170 }
171}
172
173#[derive(Debug, Clone)]
175pub enum BatchRenderError {
176 PartialFailure { successful: usize, failed: usize },
178 TotalFailure(String),
180 InvalidConfig(String),
182 QueueFull,
184 EmptyQueue,
186 DeviceLost { reason: String, message: String },
195}
196
197impl std::fmt::Display for BatchRenderError {
198 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
199 match self {
200 BatchRenderError::PartialFailure { successful, failed } => {
201 write!(
202 f,
203 "Batch render partial failure: {} succeeded, {} failed",
204 successful, failed
205 )
206 }
207 BatchRenderError::TotalFailure(msg) => write!(f, "Batch render total failure: {}", msg),
208 BatchRenderError::InvalidConfig(msg) => write!(f, "Invalid batch config: {}", msg),
209 BatchRenderError::QueueFull => write!(f, "Batch queue is full"),
210 BatchRenderError::EmptyQueue => write!(f, "No renders queued"),
211 BatchRenderError::DeviceLost { reason, message } => {
212 write!(f, "wgpu device lost ({}): {}", reason, message)
213 }
214 }
215 }
216}
217
218impl std::error::Error for BatchRenderError {}
219
220#[derive(Clone, Copy, Debug, PartialEq, Eq)]
222pub enum BatchState {
223 Idle,
225 LoadingAssets,
227 RenderingFrame,
229 ExtractingResults,
231 Cleanup,
233 Shutdown,
235}
236
237pub struct BatchRenderer {
239 pub pending_requests: VecDeque<BatchRenderRequest>,
241 pub completed_results: Vec<BatchRenderOutput>,
243 pub current_request: Option<BatchRenderRequest>,
245 pub current_output: Option<BatchRenderOutput>,
247 pub frame_count: u32,
249 pub state: BatchState,
251 pub config: BatchRenderConfig,
253 pub renders_processed: usize,
255}
256
257impl BatchRenderer {
258 pub fn new(config: BatchRenderConfig) -> Self {
260 Self {
261 pending_requests: VecDeque::new(),
262 completed_results: Vec::new(),
263 current_request: None,
264 current_output: None,
265 frame_count: 0,
266 state: BatchState::Idle,
267 config,
268 renders_processed: 0,
269 }
270 }
271
272 pub fn queue_request(&mut self, request: BatchRenderRequest) -> Result<(), BatchRenderError> {
274 if self.pending_requests.len() >= self.config.max_batch_size {
275 return Err(BatchRenderError::QueueFull);
276 }
277 self.pending_requests.push_back(request);
278 Ok(())
279 }
280
281 pub fn pending_count(&self) -> usize {
283 self.pending_requests.len()
284 }
285
286 pub fn completed_count(&self) -> usize {
288 self.completed_results.len()
289 }
290
291 pub fn take_completed(&mut self) -> Vec<BatchRenderOutput> {
293 std::mem::take(&mut self.completed_results)
294 }
295
296 pub fn is_finished(&self) -> bool {
298 self.pending_requests.is_empty() && self.current_request.is_none()
299 }
300}
301
302#[cfg(test)]
303mod tests {
304 use super::*;
305
306 #[test]
307 fn test_batch_config_defaults() {
308 let config = BatchRenderConfig::default();
309 assert_eq!(config.max_batch_size, 256);
310 assert_eq!(config.frame_timeout_ms, 500);
311 assert!(config.enable_depth_readback);
312 assert!(config.enable_asset_caching);
313 }
314
315 #[test]
316 fn test_batch_renderer_creation() {
317 let config = BatchRenderConfig::default();
318 let renderer = BatchRenderer::new(config);
319 assert_eq!(renderer.state, BatchState::Idle);
320 assert_eq!(renderer.pending_count(), 0);
321 assert_eq!(renderer.completed_count(), 0);
322 assert!(renderer.is_finished());
323 }
324
325 #[test]
326 fn test_queue_request() {
327 let mut renderer = BatchRenderer::new(BatchRenderConfig::default());
328 let request = BatchRenderRequest {
329 object_dir: "/tmp/test".into(),
330 viewpoint: Transform::default(),
331 object_rotation: ObjectRotation::identity(),
332 render_config: RenderConfig::tbp_default(),
333 };
334 assert!(renderer.queue_request(request).is_ok());
335 assert_eq!(renderer.pending_count(), 1);
336 }
337
338 #[test]
339 fn test_queue_full() {
340 let config = BatchRenderConfig {
341 max_batch_size: 1,
342 ..BatchRenderConfig::default()
343 };
344 let mut renderer = BatchRenderer::new(config);
345
346 let request = BatchRenderRequest {
347 object_dir: "/tmp/test".into(),
348 viewpoint: Transform::default(),
349 object_rotation: ObjectRotation::identity(),
350 render_config: RenderConfig::tbp_default(),
351 };
352
353 assert!(renderer.queue_request(request.clone()).is_ok());
354 assert!(matches!(
355 renderer.queue_request(request),
356 Err(BatchRenderError::QueueFull)
357 ));
358 }
359
360 #[test]
361 fn test_batch_render_output_rgb_conversion() {
362 let request = BatchRenderRequest {
363 object_dir: "/tmp/test".into(),
364 viewpoint: Transform::default(),
365 object_rotation: ObjectRotation::identity(),
366 render_config: RenderConfig::tbp_default(),
367 };
368
369 let mut rgba = vec![0u8; 2 * 2 * 4];
371 rgba[0] = 255;
373 rgba[1] = 0;
374 rgba[2] = 0;
375 rgba[3] = 255;
376
377 let output = BatchRenderOutput {
378 request,
379 rgba,
380 depth: vec![1.0; 4],
381 width: 2,
382 height: 2,
383 intrinsics: RenderConfig::tbp_default().intrinsics(),
384 status: RenderStatus::Success,
385 error_message: None,
386 };
387
388 let rgb = output.to_rgb_image();
389 assert_eq!(rgb.len(), 2); assert_eq!(rgb[0].len(), 2); assert_eq!(rgb[0][0], [255, 0, 0]); }
393}