Skip to main content

loam_render/raymarch/scene/
scene3.rs

1use std::boxed::Box;
2
3use glam::Vec3;
4
5use super::literal::wgsl_f32;
6use super::primitive::Primitive;
7use loam_math::{Space, WgslSpace};
8pub use loam_shape::Shape as PrimitiveKind;
9
10#[derive(Debug, Clone)]
11pub enum SceneNode {
12    Leaf(PrimitiveKind),
13    Union(Box<SceneNode>, Box<SceneNode>),
14    Intersection(Box<SceneNode>, Box<SceneNode>),
15    /// `max(left, -right)`: the right subtree is carved from the left.
16    Difference(Box<SceneNode>, Box<SceneNode>),
17    /// `k` is the blend radius in Space units.
18    SmoothUnion {
19        k: f32,
20        left: Box<SceneNode>,
21        right: Box<SceneNode>,
22    },
23}
24
25impl SceneNode {
26    pub fn sphere(center: Vec3, radius: f32) -> Self {
27        SceneNode::Leaf(PrimitiveKind::Sphere { center, radius })
28    }
29
30    /// Emits chart-coord `dot(p, n) - d` in flat charts and [`super::SENTINEL_DISTANCE`] in curved ones.
31    pub fn plane(normal: Vec3, offset: f32) -> Self {
32        SceneNode::Leaf(PrimitiveKind::HalfSpace { normal, offset })
33    }
34
35    pub fn box_(half_extents: Vec3) -> Self {
36        SceneNode::Leaf(PrimitiveKind::Box3 { half_extents })
37    }
38
39    pub fn cube(half_side: f32) -> Self {
40        SceneNode::Leaf(PrimitiveKind::Box3 {
41            half_extents: Vec3::splat(half_side),
42        })
43    }
44
45    pub fn union(self, other: SceneNode) -> Self {
46        SceneNode::Union(Box::new(self), Box::new(other))
47    }
48
49    pub fn intersect(self, other: SceneNode) -> Self {
50        SceneNode::Intersection(Box::new(self), Box::new(other))
51    }
52
53    pub fn subtract(self, other: SceneNode) -> Self {
54        SceneNode::Difference(Box::new(self), Box::new(other))
55    }
56
57    pub fn smooth_union(self, other: SceneNode, k: f32) -> Self {
58        SceneNode::SmoothUnion {
59            k,
60            left: Box::new(self),
61            right: Box::new(other),
62        }
63    }
64}
65
66#[derive(Debug, Clone)]
67pub struct Scene {
68    pub root: SceneNode,
69}
70
71impl Scene {
72    pub fn new(root: SceneNode) -> Self {
73        Self { root }
74    }
75
76    /// Emits `loam_scene_sdf` and helpers; prepend the space prelude.
77    pub fn to_wgsl<S: WgslSpace>(&self, space: &S) -> String {
78        let mut helpers = String::new();
79        let mut body = String::new();
80        let mut counter = 0u32;
81
82        let result_var = emit_node(&self.root, space, &mut counter, &mut helpers, &mut body);
83
84        format!(
85            "// ---- loam-scene scene (typed) ----\n\
86             {helpers}\
87             fn loam_scene_sdf(p: vec3<f32>) -> f32 {{\n\
88             {body}\
89             \treturn {result_var};\n\
90             }}\n"
91        )
92    }
93
94    /// Evaluates the scene without heap allocation.
95    pub fn eval<S: Space<Point = Vec3, Vector = Vec3>>(&self, space: &S, p: Vec3) -> f32 {
96        eval_node(&self.root, space, p)
97    }
98}
99
100fn emit_node<S: WgslSpace>(
101    node: &SceneNode,
102    space: &S,
103    counter: &mut u32,
104    helpers: &mut String,
105    body: &mut String,
106) -> String {
107    let idx = *counter;
108    *counter += 1;
109
110    match node {
111        SceneNode::Leaf(prim) => {
112            let fn_name = format!("sdf_p{idx}");
113            helpers.push_str(&prim.to_wgsl(space, &fn_name));
114            let var = format!("d{idx}");
115            body.push_str(&format!("\tlet {var} = {fn_name}(p);\n"));
116            var
117        }
118
119        SceneNode::Union(left, right) => {
120            let lv = emit_node(left, space, counter, helpers, body);
121            let rv = emit_node(right, space, counter, helpers, body);
122            let var = format!("d{idx}");
123            body.push_str(&format!("\tlet {var} = min({lv}, {rv});\n"));
124            var
125        }
126
127        SceneNode::Intersection(left, right) => {
128            let lv = emit_node(left, space, counter, helpers, body);
129            let rv = emit_node(right, space, counter, helpers, body);
130            let var = format!("d{idx}");
131            body.push_str(&format!("\tlet {var} = max({lv}, {rv});\n"));
132            var
133        }
134
135        SceneNode::Difference(left, right) => {
136            let lv = emit_node(left, space, counter, helpers, body);
137            let rv = emit_node(right, space, counter, helpers, body);
138            let var = format!("d{idx}");
139            body.push_str(&format!("\tlet {var} = max({lv}, -({rv}));\n"));
140            var
141        }
142
143        SceneNode::SmoothUnion { k, left, right } => {
144            let fn_name = format!("sdf_smin{idx}");
145            helpers.push_str(&smooth_min_fn(&fn_name, *k));
146            let lv = emit_node(left, space, counter, helpers, body);
147            let rv = emit_node(right, space, counter, helpers, body);
148            let var = format!("d{idx}");
149            body.push_str(&format!("\tlet {var} = {fn_name}({lv}, {rv});\n"));
150            var
151        }
152    }
153}
154
155fn smooth_min_fn(name: &str, k: f32) -> String {
156    let k = wgsl_f32(k);
157    format!(
158        "fn {name}(a: f32, b: f32) -> f32 {{\n\
159         \tlet h = clamp(0.5 + 0.5 * (b - a) / ({k}), 0.0, 1.0);\n\
160         \treturn mix(b, a, h) - ({k}) * h * (1.0 - h);\n\
161         }}\n",
162    )
163}
164
165fn eval_node<S: Space<Point = Vec3, Vector = Vec3>>(node: &SceneNode, space: &S, p: Vec3) -> f32 {
166    match node {
167        SceneNode::Leaf(prim) => prim.eval(space, p),
168
169        SceneNode::Union(left, right) => eval_node(left, space, p).min(eval_node(right, space, p)),
170
171        SceneNode::Intersection(left, right) => {
172            eval_node(left, space, p).max(eval_node(right, space, p))
173        }
174
175        SceneNode::Difference(left, right) => {
176            eval_node(left, space, p).max(-eval_node(right, space, p))
177        }
178
179        SceneNode::SmoothUnion { k, left, right } => {
180            let a = eval_node(left, space, p);
181            let b = eval_node(right, space, p);
182            let h = (0.5 + 0.5 * (b - a) / k).clamp(0.0, 1.0);
183            (b * (1.0 - h) + a * h) - k * h * (1.0 - h)
184        }
185    }
186}