1#define_import_path bevy_gaussian_splatting::planar
2
3#ifdef PRECOMPUTE_COVARIANCE_3D
4#import bevy_gaussian_splatting::bindings::{
5    position_visibility,
6    spherical_harmonics,
7    covariance_3d_opacity,
8}
9#else
10#import bevy_gaussian_splatting::bindings::{
11    position_visibility,
12    spherical_harmonics,
13    rotation,
14    rotation_scale_opacity,
15    scale_opacity,
16}
17#endif
18
19#import bevy_gaussian_splatting::spherical_harmonics::{
20    spherical_harmonics_lookup,
21    srgb_to_linear,
22}
23
24
25#ifdef PLANAR_F16
26
27fn get_color(
28    index: u32,
29    ray_direction: vec3<f32>,
30) -> vec3<f32> {
31    let sh = get_spherical_harmonics(index);
32    let color = spherical_harmonics_lookup(ray_direction, sh);
33    return srgb_to_linear(color);
34}
35
36fn get_position(index: u32) -> vec3<f32> {
37    return position_visibility[index].xyz;
38}
39
40fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
41    var coefficients: array<f32, #{SH_COEFF_COUNT}>;
42
43    for (var i = 0u; i < #{HALF_SH_COEFF_COUNT}u; i = i + 1u) {
44        let values = unpack2x16float(spherical_harmonics[index][i]);
45
46        coefficients[i * 2u] = values[0];
47        coefficients[i * 2u + 1u] = values[1];
48    }
49
50    return coefficients;
51}
52
53#ifdef PRECOMPUTE_COVARIANCE_3D
54    fn get_cov3d(index: u32) -> array<f32, 6> {
55        let c0 = unpack2x16float(covariance_3d_opacity[index].x);
56        let c1 = unpack2x16float(covariance_3d_opacity[index].y);
57        let c2 = unpack2x16float(covariance_3d_opacity[index].z);
58
59        var cov3d: array<f32, 6>;
60
61        cov3d[0] = c0.y;
62        cov3d[1] = c0.x;
63        cov3d[2] = c1.y;
64        cov3d[3] = c1.x;
65        cov3d[4] = c2.y;
66        cov3d[5] = c2.x;
67
68        return cov3d;
69    }
70#else
71    fn get_rotation(index: u32) -> vec4<f32> {
72        let q0 = unpack2x16float(rotation_scale_opacity[index].x);
73        let q1 = unpack2x16float(rotation_scale_opacity[index].y);
74
75        return vec4<f32>(
76            q0.yx,
77            q1.yx,
78        );
79    }
80
81    fn get_scale(index: u32) -> vec3<f32> {
82        let s0 = unpack2x16float(rotation_scale_opacity[index].z);
83        let s1 = unpack2x16float(rotation_scale_opacity[index].w);
84
85        return vec3<f32>(
86            s0.yx,
87            s1.y,
88        );
89    }
90#endif
91
92fn get_opacity(index: u32) -> f32 {
93#ifdef PRECOMPUTE_COVARIANCE_3D
94    return unpack2x16float(covariance_3d_opacity[index].w).y;
95#else
96    return unpack2x16float(rotation_scale_opacity[index].w).x;
97#endif
98}
99
100fn get_visibility(index: u32) -> f32 {
101    return position_visibility[index].w;
102}
103#endif
104
105
106#ifdef PLANAR_F32
107fn get_color(
108    index: u32,
109    ray_direction: vec3<f32>,
110) -> vec3<f32> {
111    let sh = get_spherical_harmonics(index);
112    let color = spherical_harmonics_lookup(ray_direction, sh);
113    return srgb_to_linear(color);
114}
115
116fn get_position(index: u32) -> vec3<f32> {
117    return position_visibility[index].xyz;
118}
119
120fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
121    return spherical_harmonics[index];
122}
123
124fn get_rotation(index: u32) -> vec4<f32> {
125    return rotation[index];
126}
127
128fn get_scale(index: u32) -> vec3<f32> {
129    return scale_opacity[index].xyz;
130}
131
132fn get_opacity(index: u32) -> f32 {
133    return scale_opacity[index].w;
134}
135
136fn get_visibility(index: u32) -> f32 {
137    return position_visibility[index].w;
138}
139#endif