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

use std::sync::Mutex;

use crate::tonemapping::{TonemappingLuts, TonemappingPipeline, ViewTonemappingPipeline};

use bevy_ecs::{prelude::*, query::QueryItem};
use bevy_render::{
    render_asset::RenderAssets,
    render_graph::{NodeRunError, RenderGraphContext, ViewNode},
    render_resource::{
        BindGroup, BindGroupEntries, BufferId, LoadOp, Operations, PipelineCache,
        RenderPassColorAttachment, RenderPassDescriptor, StoreOp, TextureViewId,
    },
    renderer::RenderContext,
    texture::Image,
    view::{ViewTarget, ViewUniformOffset, ViewUniforms},
};

use super::{get_lut_bindings, Tonemapping};

#[derive(Default)]
pub struct TonemappingNode {
    cached_bind_group: Mutex<Option<(BufferId, TextureViewId, BindGroup)>>,
    last_tonemapping: Mutex<Option<Tonemapping>>,
}

impl ViewNode for TonemappingNode {
    type ViewQuery = (
        &'static ViewUniformOffset,
        &'static ViewTarget,
        &'static ViewTonemappingPipeline,
        &'static Tonemapping,
    );

    fn run(
        &self,
        _graph: &mut RenderGraphContext,
        render_context: &mut RenderContext,
        (view_uniform_offset, target, view_tonemapping_pipeline, tonemapping): QueryItem<
            Self::ViewQuery,
        >,
        world: &World,
    ) -> Result<(), NodeRunError> {
        let pipeline_cache = world.resource::<PipelineCache>();
        let tonemapping_pipeline = world.resource::<TonemappingPipeline>();
        let gpu_images = world.get_resource::<RenderAssets<Image>>().unwrap();
        let view_uniforms_resource = world.resource::<ViewUniforms>();
        let view_uniforms = &view_uniforms_resource.uniforms;
        let view_uniforms_id = view_uniforms.buffer().unwrap().id();

        if !target.is_hdr() {
            return Ok(());
        }

        let Some(pipeline) = pipeline_cache.get_render_pipeline(view_tonemapping_pipeline.0) else {
            return Ok(());
        };

        let post_process = target.post_process_write();
        let source = post_process.source;
        let destination = post_process.destination;

        let mut last_tonemapping = self.last_tonemapping.lock().unwrap();

        let tonemapping_changed = if let Some(last_tonemapping) = &*last_tonemapping {
            tonemapping != last_tonemapping
        } else {
            true
        };
        if tonemapping_changed {
            *last_tonemapping = Some(*tonemapping);
        }

        let mut cached_bind_group = self.cached_bind_group.lock().unwrap();
        let bind_group = match &mut *cached_bind_group {
            Some((buffer_id, texture_id, bind_group))
                if view_uniforms_id == *buffer_id
                    && source.id() == *texture_id
                    && !tonemapping_changed =>
            {
                bind_group
            }
            cached_bind_group => {
                let tonemapping_luts = world.resource::<TonemappingLuts>();

                let lut_bindings = get_lut_bindings(gpu_images, tonemapping_luts, tonemapping);

                let bind_group = render_context.render_device().create_bind_group(
                    None,
                    &tonemapping_pipeline.texture_bind_group,
                    &BindGroupEntries::sequential((
                        view_uniforms,
                        source,
                        &tonemapping_pipeline.sampler,
                        lut_bindings.0,
                        lut_bindings.1,
                    )),
                );

                let (_, _, bind_group) =
                    cached_bind_group.insert((view_uniforms_id, source.id(), bind_group));
                bind_group
            }
        };

        let pass_descriptor = RenderPassDescriptor {
            label: Some("tonemapping_pass"),
            color_attachments: &[Some(RenderPassColorAttachment {
                view: destination,
                resolve_target: None,
                ops: Operations {
                    load: LoadOp::Clear(Default::default()), // TODO shouldn't need to be cleared
                    store: StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
        };

        let mut render_pass = render_context
            .command_encoder()
            .begin_render_pass(&pass_descriptor);

        render_pass.set_pipeline(pipeline);
        render_pass.set_bind_group(0, bind_group, &[view_uniform_offset.offset]);
        render_pass.draw(0..3, 0..1);

        Ok(())
    }
}