1#define_import_path bevy_gaussian_splatting::packed
2
3#import bevy_gaussian_splatting::bindings::points
4#ifdef BINARY_GAUSSIAN_OP
5    #import bevy_gaussian_splatting::bindings::{rhs_points, out_points}
6#endif
7
8#import bevy_gaussian_splatting::spherical_harmonics::{
9    spherical_harmonics_lookup,
10    srgb_to_linear,
11}
12
13#ifdef PACKED_F32
14
15    fn gaussian_position(point: Gaussian) -> vec3<f32> {
16        return point.position_visibility.xyz;
17    }
18
19    fn gaussian_color(point: Gaussian, ray_direction: vec3<f32>) -> vec3<f32> {
20        let sh = gaussian_spherical_harmonics(point);
21        let color = spherical_harmonics_lookup(ray_direction, sh);
22        return srgb_to_linear(color);
23    }
24
25    fn gaussian_spherical_harmonics(point: Gaussian) -> array<f32, #{SH_COEFF_COUNT}> {
26        return point.sh;
27    }
28
29    fn gaussian_rotation(point: Gaussian) -> vec4<f32> {
30        return point.rotation;
31    }
32
33    fn gaussian_scale(point: Gaussian) -> vec3<f32> {
34        return point.scale_opacity.xyz;
35    }
36
37    fn gaussian_opacity(point: Gaussian) -> f32 {
38        return point.scale_opacity.w;
39    }
40
41    fn gaussian_visibility(point: Gaussian) -> f32 {
42        return point.position_visibility.w;
43    }
44
45    fn get_position(index: u32) -> vec3<f32> {
46        return gaussian_position(points[index]);
47    }
48
49    fn get_color(
50        index: u32,
51        ray_direction: vec3<f32>,
52    ) -> vec3<f32> {
53        return gaussian_color(points[index], ray_direction);
54    }
55
56    fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
57        return gaussian_spherical_harmonics(points[index]);
58    }
59
60    fn get_rotation(index: u32) -> vec4<f32> {
61        return gaussian_rotation(points[index]);
62    }
63
64    fn get_scale(index: u32) -> vec3<f32> {
65        return gaussian_scale(points[index]);
66    }
67
68    fn get_opacity(index: u32) -> f32 {
69        return gaussian_opacity(points[index]);
70    }
71
72    fn get_visibility(index: u32) -> f32 {
73        return gaussian_visibility(points[index]);
74    }
75
76    #ifdef BINARY_GAUSSIAN_OP
77
78        fn get_rhs_position(index: u32) -> vec3<f32> {
79            return gaussian_position(rhs_points[index]);
80        }
81
82        fn get_rhs_color(
83            index: u32,
84            ray_direction: vec3<f32>,
85        ) -> vec3<f32> {
86            return gaussian_color(rhs_points[index], ray_direction);
87        }
88
89        fn get_rhs_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
90            return gaussian_spherical_harmonics(rhs_points[index]);
91        }
92
93        fn get_rhs_rotation(index: u32) -> vec4<f32> {
94            return gaussian_rotation(rhs_points[index]);
95        }
96
97        fn get_rhs_scale(index: u32) -> vec3<f32> {
98            return gaussian_scale(rhs_points[index]);
99        }
100
101        fn get_rhs_opacity(index: u32) -> f32 {
102            return gaussian_opacity(rhs_points[index]);
103        }
104
105        fn get_rhs_visibility(index: u32) -> f32 {
106            return gaussian_visibility(rhs_points[index]);
107        }
108
109        fn set_output_position_visibility(
110            index: u32,
111            position: vec3<f32>,
112            visibility: f32,
113        ) {
114            out_points[index].position_visibility = vec4<f32>(position, visibility);
115        }
116
117        fn set_output_spherical_harmonics(
118            index: u32,
119            sh: array<f32, #{SH_COEFF_COUNT}>,
120        ) {
121            out_points[index].sh = sh;
122        }
123
124        fn set_output_transform(
125            index: u32,
126            rotation: vec4<f32>,
127            scale: vec3<f32>,
128            opacity: f32,
129        ) {
130            out_points[index].rotation = rotation;
131            out_points[index].scale_opacity = vec4<f32>(scale, opacity);
132        }
133
134    #endif
135
136
137#endif