1use std::collections::HashMap;
2
3use crate::{
4 composition::Composition,
5 error::RenderError,
6 expr::ExpressionContext,
7 gpu::{
8 BoundFrame, CompiledComposition, CompiledOutput, FrameBindContext, GpuCompiledNode,
9 RasterHandle, RasterMetadata,
10 },
11 media::MediaStore,
12 node::{Deferred, NodeId, NodeKind, PortRef},
13};
14
15pub(crate) use super::params::*;
16
17#[derive(Debug, Clone)]
18struct BackgroundClearBinding {
19 node_id: NodeId,
20 color: [u8; 4],
21 buffer: lumen_gpu::BufferId,
22}
23
24impl GpuCompiledNode for BackgroundClearBinding {
25 fn node_id(&self) -> NodeId {
26 self.node_id
27 }
28
29 fn bind(&self, _ctx: &FrameBindContext<'_>, bound: &mut BoundFrame) -> crate::Result<()> {
30 bound.write_buffer(
31 self.buffer,
32 0,
33 bytemuck::bytes_of(&ColorParams::from_rgba8(self.color)),
34 );
35 Ok(())
36 }
37}
38
39pub trait GpuCompileNode {
40 fn compile_gpu(
41 &self,
42 ctx: &mut CompileContext<'_>,
43 port: &PortRef,
44 ) -> crate::Result<CompiledOutput>;
45}
46
47#[derive(Debug)]
48struct CompiledPortKey {
49 port: PortRef,
50 frame: u32,
51}
52
53impl PartialEq for CompiledPortKey {
54 fn eq(&self, other: &Self) -> bool {
55 self.frame == other.frame && self.port == other.port
56 }
57}
58
59impl Eq for CompiledPortKey {}
60
61impl std::hash::Hash for CompiledPortKey {
62 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
63 self.port.hash(state);
64 self.frame.hash(state);
65 }
66}
67
68#[derive(Debug)]
69pub struct CompileContext<'a> {
70 composition: &'a Composition,
71 frame: u32,
72 media: Option<&'a dyn MediaStore>,
73 builder: lumen_gpu::RenderPlanBuilder,
74 outputs: HashMap<CompiledPortKey, CompiledOutput>,
75 public_outputs: HashMap<PortRef, CompiledOutput>,
76 compiled_nodes: HashMap<NodeId, Box<dyn GpuCompiledNode>>,
77 output_format: lumen_gpu::wgpu::TextureFormat,
78}
79
80impl<'a> CompileContext<'a> {
81 pub fn new(composition: &'a Composition) -> Self {
82 Self::with_output_format(composition, lumen_gpu::wgpu::TextureFormat::Rgba8Unorm)
83 }
84
85 pub fn with_output_format(
86 composition: &'a Composition,
87 output_format: lumen_gpu::wgpu::TextureFormat,
88 ) -> Self {
89 Self::with_options(composition, 0, None, output_format)
90 }
91
92 pub fn with_frame(
93 composition: &'a Composition,
94 frame: u32,
95 output_format: lumen_gpu::wgpu::TextureFormat,
96 ) -> Self {
97 Self::with_options(composition, frame, None, output_format)
98 }
99
100 pub fn with_media<M: MediaStore>(
101 composition: &'a Composition,
102 media: &'a M,
103 output_format: lumen_gpu::wgpu::TextureFormat,
104 ) -> Self {
105 Self::with_media_for_frame(composition, 0, media, output_format)
106 }
107
108 pub fn with_media_for_frame<M: MediaStore>(
109 composition: &'a Composition,
110 frame: u32,
111 media: &'a M,
112 output_format: lumen_gpu::wgpu::TextureFormat,
113 ) -> Self {
114 Self::with_options(composition, frame, Some(media), output_format)
115 }
116
117 fn with_options(
118 composition: &'a Composition,
119 frame: u32,
120 media: Option<&'a dyn MediaStore>,
121 output_format: lumen_gpu::wgpu::TextureFormat,
122 ) -> Self {
123 Self {
124 composition,
125 frame,
126 media,
127 builder: lumen_gpu::RenderPlan::builder(),
128 outputs: HashMap::new(),
129 public_outputs: HashMap::new(),
130 compiled_nodes: HashMap::new(),
131 output_format,
132 }
133 }
134
135 pub fn compile(mut self) -> crate::Result<CompiledComposition> {
136 let output_node = self.media_output_node()?;
137 let output_ref = PortRef::new(output_node, "output".to_string());
138 let output = self
139 .compile_port(&output_ref)?
140 .into_raster(output_node, "output")?;
141 Ok(CompiledComposition {
142 plan: self.builder.build(),
143 output,
144 node_outputs: self.public_outputs,
145 compiled_nodes: self.compiled_nodes,
146 })
147 }
148
149 pub(crate) fn composition(&self) -> &Composition {
150 self.composition
151 }
152
153 pub(crate) fn media(&self) -> Option<&dyn MediaStore> {
154 self.media
155 }
156
157 pub(crate) fn output_format(&self) -> lumen_gpu::wgpu::TextureFormat {
158 self.output_format
159 }
160
161 pub(crate) fn builder_mut(&mut self) -> &mut lumen_gpu::RenderPlanBuilder {
162 &mut self.builder
163 }
164
165 pub(crate) fn register_compiled_node<N>(&mut self, node: N)
166 where
167 N: GpuCompiledNode + 'static,
168 {
169 self.compiled_nodes.insert(node.node_id(), Box::new(node));
170 }
171
172 pub(crate) fn compile_port(&mut self, port: &PortRef) -> crate::Result<CompiledOutput> {
173 let key = CompiledPortKey {
174 port: port.clone(),
175 frame: self.frame,
176 };
177 if let Some(output) = self.outputs.get(&key) {
178 return Ok(output.clone());
179 }
180
181 let node = self
182 .composition
183 .graph
184 .nodes
185 .get(&port.id)
186 .ok_or(RenderError::MissingNode {
187 frame: 0,
188 node_id: port.id,
189 })?;
190 let output = match node {
191 NodeKind::MediaIn(node) => node.compile_gpu(self, port)?,
192 NodeKind::Background(node) => node.compile_gpu(self, port)?,
193 NodeKind::Text(node) => node.compile_gpu(self, port)?,
194 NodeKind::Path(node) => node.compile_gpu(self, port)?,
195 NodeKind::Shape(node) => node.compile_gpu(self, port)?,
196 NodeKind::Boolean(node) => node.compile_gpu(self, port)?,
197 NodeKind::Merge(node) => node.compile_gpu(self, port)?,
198 NodeKind::RasterMultiMerge(node) => node.compile_gpu(self, port)?,
199 NodeKind::AlphaPremultiply(node) => node.compile_gpu(self, port)?,
200 NodeKind::Blur(node) => node.compile_gpu(self, port)?,
201 NodeKind::ChannelShuffle(node) => node.compile_gpu(self, port)?,
202 NodeKind::ColorGrade(node) => node.compile_gpu(self, port)?,
203 NodeKind::Curves(node) => node.compile_gpu(self, port)?,
204 NodeKind::Exposure(node) => node.compile_gpu(self, port)?,
205 NodeKind::HueSaturation(node) => node.compile_gpu(self, port)?,
206 NodeKind::Levels(node) => node.compile_gpu(self, port)?,
207 NodeKind::Memo(node) => node.compile_gpu(self, port)?,
208 NodeKind::TimeRemap(node) => node.compile_gpu(self, port)?,
209 NodeKind::Transform(node) => node.compile_gpu(self, port)?,
210 NodeKind::Crop(node) => node.compile_gpu(self, port)?,
211 NodeKind::Resize(node) => node.compile_gpu(self, port)?,
212 NodeKind::Shadow(node) => node.compile_gpu(self, port)?,
213 NodeKind::WgslShader(node) => node.compile_gpu(self, port)?,
214 NodeKind::Switch(node) => node.compile_gpu(self, port)?,
215 NodeKind::MediaOutput(node) => node.compile_gpu(self, port)?,
216 };
217
218 self.outputs.insert(key, output.clone());
219 self.public_outputs
220 .entry(port.clone())
221 .or_insert_with(|| output.clone());
222 Ok(output)
223 }
224
225 pub(crate) fn with_frame_context<T>(
226 &mut self,
227 frame: u32,
228 f: impl FnOnce(&mut Self) -> crate::Result<T>,
229 ) -> crate::Result<T> {
230 let original_frame = self.frame;
231 self.frame = frame;
232 let result = f(self);
233 self.frame = original_frame;
234 result
235 }
236
237 pub(crate) fn compile_unary_filter(
238 &mut self,
239 node_id: NodeId,
240 source_ref: &PortRef,
241 port: &PortRef,
242 label: &str,
243 shader: &str,
244 params_size: u64,
245 ) -> crate::Result<(RasterHandle, lumen_gpu::TextureId, lumen_gpu::BufferId)> {
246 if port.port != "output" {
247 return Err(self.missing_output(node_id, &port.port));
248 }
249
250 let source = self
251 .compile_port(source_ref)?
252 .into_raster(source_ref.id, &source_ref.port)?;
253 let size = source.domain.storage_size;
254 let texture = self.builder.texture_for(
255 lumen_gpu::NodeKey(node_id.0),
256 Some(format!("{label}:{}:output", node_id.0)),
257 lumen_gpu::TextureDesc::storage(size, lumen_gpu::wgpu::TextureFormat::Rgba8Unorm),
258 );
259 let params = self.builder.buffer_for(
260 lumen_gpu::NodeKey(node_id.0),
261 Some(format!("{label}:{}:params", node_id.0)),
262 lumen_gpu::BufferDesc::uniform(params_size),
263 );
264 let program = self.builder.program_for(
265 lumen_gpu::NodeKey(node_id.0),
266 lumen_gpu::ProgramDesc::Compute(lumen_gpu::ComputeProgramDesc {
267 label: Some(label.to_string()),
268 shader: shader.to_string(),
269 entry: "cs_main".to_string(),
270 bind_groups: lumen_gpu::BindGroupLayoutSpec::single(vec![
271 lumen_gpu::BindingLayoutEntry::texture(
272 0,
273 lumen_gpu::wgpu::ShaderStages::COMPUTE,
274 ),
275 lumen_gpu::BindingLayoutEntry::uniform(
276 1,
277 lumen_gpu::wgpu::ShaderStages::COMPUTE,
278 ),
279 lumen_gpu::BindingLayoutEntry::storage_texture(
280 2,
281 lumen_gpu::wgpu::ShaderStages::COMPUTE,
282 lumen_gpu::wgpu::TextureFormat::Rgba8Unorm,
283 lumen_gpu::wgpu::StorageTextureAccess::WriteOnly,
284 ),
285 ]),
286 }),
287 );
288 self.builder.compute_pass(lumen_gpu::ComputePassDesc {
289 label: Some(format!("{label}:{}:apply", node_id.0)),
290 owner: Some(lumen_gpu::NodeKey(node_id.0)),
291 program,
292 bindings: vec![
293 lumen_gpu::Binding::sampled_texture(0, 0, source.texture),
294 lumen_gpu::Binding::uniform(0, 1, params),
295 lumen_gpu::Binding::storage_texture(0, 2, texture),
296 ],
297 dispatch: dispatch_for(size).into(),
298 });
299 self.builder.param(
300 lumen_gpu::ParamKey {
301 owner: lumen_gpu::NodeKey(node_id.0),
302 slot: 0,
303 },
304 lumen_gpu::ParamTarget::Buffer(params),
305 );
306 Ok((source, texture, params))
307 }
308
309 pub(crate) fn compile_transparent(&mut self, node_id: NodeId) -> CompiledOutput {
310 let size = lumen_gpu::Size::new(
311 self.composition.render_settings.width.max(1),
312 self.composition.render_settings.height.max(1),
313 );
314 let texture = self.builder.texture_for(
315 lumen_gpu::NodeKey(node_id.0),
316 Some(format!("transparent:{}:output", node_id.0)),
317 lumen_gpu::TextureDesc::storage(size, lumen_gpu::wgpu::TextureFormat::Rgba8Unorm),
318 );
319 let params = self.builder.buffer_for(
320 lumen_gpu::NodeKey(node_id.0),
321 Some(format!("transparent:{}:params", node_id.0)),
322 lumen_gpu::BufferDesc::uniform(std::mem::size_of::<ColorParams>() as u64),
323 );
324 let program = self.builder.program_for(
325 lumen_gpu::NodeKey(node_id.0),
326 lumen_gpu::ProgramDesc::Compute(lumen_gpu::ComputeProgramDesc {
327 label: Some("transparent".to_string()),
328 shader: crate::node::source::background::SHADER.to_string(),
329 entry: "cs_main".to_string(),
330 bind_groups: lumen_gpu::BindGroupLayoutSpec::single(vec![
331 lumen_gpu::BindingLayoutEntry::uniform(
332 0,
333 lumen_gpu::wgpu::ShaderStages::COMPUTE,
334 ),
335 lumen_gpu::BindingLayoutEntry::storage_texture(
336 1,
337 lumen_gpu::wgpu::ShaderStages::COMPUTE,
338 lumen_gpu::wgpu::TextureFormat::Rgba8Unorm,
339 lumen_gpu::wgpu::StorageTextureAccess::WriteOnly,
340 ),
341 ]),
342 }),
343 );
344 self.builder.compute_pass(lumen_gpu::ComputePassDesc {
345 label: Some(format!("transparent:{}:fill", node_id.0)),
346 owner: Some(lumen_gpu::NodeKey(node_id.0)),
347 program,
348 bindings: vec![
349 lumen_gpu::Binding::uniform(0, 0, params),
350 lumen_gpu::Binding::storage_texture(0, 1, texture),
351 ],
352 dispatch: dispatch_for(size).into(),
353 });
354 self.builder.param(
355 lumen_gpu::ParamKey {
356 owner: lumen_gpu::NodeKey(node_id.0),
357 slot: 0,
358 },
359 lumen_gpu::ParamTarget::Buffer(params),
360 );
361 self.register_compiled_node(BackgroundClearBinding {
362 node_id,
363 color: [0, 0, 0, 0],
364 buffer: params,
365 });
366
367 CompiledOutput::Raster(RasterHandle {
368 texture,
369 domain: lumen_gpu::TextureDomain::full_frame(size),
370 metadata: RasterMetadata::default(),
371 })
372 }
373
374 pub(crate) fn static_dimension(
375 &self,
376 property: &Deferred<i64>,
377 node_id: NodeId,
378 property_path: &str,
379 ) -> crate::Result<u32> {
380 let value = property.resolve_int(
381 node_id,
382 property_path,
383 &self.expr_context(node_id, property_path),
384 )?;
385 let value = if value <= 0 {
386 match property_path {
387 "width" => i64::from(self.composition.render_settings.width),
388 "height" => i64::from(self.composition.render_settings.height),
389 _ => value,
390 }
391 } else {
392 value
393 };
394 Ok(value.clamp(1, i64::from(u32::MAX)) as u32)
395 }
396
397 pub(crate) fn static_dimension_value(&self, value: i64, property_path: &str) -> u32 {
398 let value = if value <= 0 {
399 match property_path {
400 "width" => i64::from(self.composition.render_settings.width),
401 "height" => i64::from(self.composition.render_settings.height),
402 _ => value,
403 }
404 } else {
405 value
406 };
407 value.clamp(1, i64::from(u32::MAX)) as u32
408 }
409
410 pub(crate) fn spatial_program(
411 &mut self,
412 node_id: NodeId,
413 label: &str,
414 shader: &str,
415 ) -> lumen_gpu::ProgramId {
416 self.builder.program_for(
417 lumen_gpu::NodeKey(node_id.0),
418 lumen_gpu::ProgramDesc::Compute(lumen_gpu::ComputeProgramDesc {
419 label: Some(label.to_string()),
420 shader: shader.to_string(),
421 entry: "cs_main".to_string(),
422 bind_groups: lumen_gpu::BindGroupLayoutSpec::single(vec![
423 lumen_gpu::BindingLayoutEntry::texture(
424 0,
425 lumen_gpu::wgpu::ShaderStages::COMPUTE,
426 ),
427 lumen_gpu::BindingLayoutEntry::uniform(
428 1,
429 lumen_gpu::wgpu::ShaderStages::COMPUTE,
430 ),
431 lumen_gpu::BindingLayoutEntry::storage_texture(
432 2,
433 lumen_gpu::wgpu::ShaderStages::COMPUTE,
434 lumen_gpu::wgpu::TextureFormat::Rgba8Unorm,
435 lumen_gpu::wgpu::StorageTextureAccess::WriteOnly,
436 ),
437 ]),
438 }),
439 )
440 }
441
442 #[allow(dead_code)]
443 fn compile_unary_compute(
444 &mut self,
445 node_id: NodeId,
446 port: &PortRef,
447 source_port: &PortRef,
448 label: &str,
449 shader: &str,
450 param_size: u64,
451 ) -> crate::Result<(RasterHandle, lumen_gpu::TextureId, lumen_gpu::BufferId)> {
452 if port.port != "output" {
453 return Err(self.missing_output(node_id, &port.port));
454 }
455
456 let source = self
457 .compile_port(source_port)?
458 .into_raster(source_port.id, &source_port.port)?;
459 let size = source.domain.storage_size;
460 let texture = self.builder.texture_for(
461 lumen_gpu::NodeKey(node_id.0),
462 Some(format!("{label}:{}:output", node_id.0)),
463 lumen_gpu::TextureDesc::storage(size, lumen_gpu::wgpu::TextureFormat::Rgba8Unorm),
464 );
465 let params = self.builder.buffer_for(
466 lumen_gpu::NodeKey(node_id.0),
467 Some(format!("{label}:{}:params", node_id.0)),
468 lumen_gpu::BufferDesc::uniform(param_size),
469 );
470 let program = self.spatial_program(node_id, label, shader);
471 self.builder.compute_pass(lumen_gpu::ComputePassDesc {
472 label: Some(format!("{label}:{}:apply", node_id.0)),
473 owner: Some(lumen_gpu::NodeKey(node_id.0)),
474 program,
475 bindings: spatial_bindings(source.texture, params, texture),
476 dispatch: dispatch_for(size).into(),
477 });
478 self.builder.param(
479 lumen_gpu::ParamKey {
480 owner: lumen_gpu::NodeKey(node_id.0),
481 slot: 0,
482 },
483 lumen_gpu::ParamTarget::Buffer(params),
484 );
485 Ok((source, texture, params))
486 }
487
488 fn media_output_node(&self) -> crate::Result<NodeId> {
489 let mut outputs = self
490 .composition
491 .graph
492 .nodes
493 .iter()
494 .filter_map(|(node_id, node)| {
495 matches!(node, NodeKind::MediaOutput(_)).then_some(*node_id)
496 });
497 let Some(output) = outputs.next() else {
498 return Err(crate::error::GraphValidationError::MissingMediaOutput.into());
499 };
500 if outputs.next().is_some() {
501 return Err(
502 crate::error::GraphValidationError::MultipleMediaOutputs { count: 2 }.into(),
503 );
504 }
505 Ok(output)
506 }
507
508 pub(crate) fn expr_context(
509 &self,
510 node_id: NodeId,
511 property_path: &str,
512 ) -> ExpressionContext<'_> {
513 ExpressionContext {
514 frame: self.frame,
515 fps: self.composition.timeline.fps,
516 width: self.composition.render_settings.width,
517 height: self.composition.render_settings.height,
518 duration_frames: self.composition.timeline.duration_frames,
519 path: Some(format!("{node_id}.{property_path}")),
520 graph: Some(&self.composition.graph),
521 }
522 }
523
524 pub(crate) fn missing_output(&self, node_id: NodeId, port: &str) -> crate::error::LumenError {
525 crate::error::PropertyError::MissingProperty {
526 node_id,
527 property_path: format!("output port `{port}`"),
528 }
529 .into()
530 }
531}