1// Copyright (c) 2023 Tomasz Stachowiak2//3// This contribution is dual licensed under EITHER OF4//5// Apache License, Version 2.0, (http://www.apache.org/licenses/LICENSE-2.0)6// MIT license (http://opensource.org/licenses/MIT)7//8// at your option.9//10// This is a port of the original [`raymarch.hlsl`] to WGSL. It's deliberately11// kept as close as possible so that patches to the original `raymarch.hlsl`12// have the greatest chances of applying to this version.13//14// [`raymarch.hlsl`]:15// https://gist.github.com/h3r2tic/9c8356bdaefbe80b1a22ae0aaee192db17#define_import_path bevy_pbr::raymarch19#import bevy_pbr::mesh_view_bindings::depth_prepass_texture20#import bevy_pbr::view_transformations::{21 direction_world_to_clip,22 ndc_to_uv,23 perspective_camera_near,24 position_world_to_ndc,25}27// Allows us to sample from the depth buffer with bilinear filtering.28@group(2) @binding(2) var depth_linear_sampler: sampler;30// Allows us to sample from the depth buffer with nearest-neighbor filtering.31@group(2) @binding(3) var depth_nearest_sampler: sampler;33// Main code35struct HybridRootFinder {36 linear_steps: u32,37 bisection_steps: u32,38 use_secant: bool,39 linear_march_exponent: f32,41 jitter: f32,42 min_t: f32,43 max_t: f32,44}46fn hybrid_root_finder_new_with_linear_steps(v: u32) -> HybridRootFinder {47 var res: HybridRootFinder;48 res.linear_steps = v;49 res.bisection_steps = 0u;50 res.use_secant = false;51 res.linear_march_exponent = 1.0;52 res.jitter = 1.0;53 res.min_t = 0.0;54 res.max_t = 1.0;55 return res;56}58fn hybrid_root_finder_find_root(59 root_finder: ptr<function, HybridRootFinder>,60 start: vec3<f32>,61 end: vec3<f32>,62 distance_fn: ptr<function, DepthRaymarchDistanceFn>,63 hit_t: ptr<function, f32>,64 miss_t: ptr<function, f32>,65 hit_d: ptr<function, DistanceWithPenetration>,66) -> bool {67 let dir = end - start;69 var min_t = (*root_finder).min_t;70 var max_t = (*root_finder).max_t;72 var min_d = DistanceWithPenetration(0.0, false, 0.0);73 var max_d = DistanceWithPenetration(0.0, false, 0.0);75 let step_size = (max_t - min_t) / f32((*root_finder).linear_steps);77 var intersected = false;79 //80 // Ray march using linear steps82 if ((*root_finder).linear_steps > 0u) {83 let candidate_t = mix(84 min_t,85 max_t,86 pow(87 (*root_finder).jitter / f32((*root_finder).linear_steps),88 (*root_finder).linear_march_exponent89 )90 );92 let candidate = start + dir * candidate_t;93 let candidate_d = depth_raymarch_distance_fn_evaluate(distance_fn, candidate);94 intersected = candidate_d.distance < 0.0 && candidate_d.valid;96 if (intersected) {97 max_t = candidate_t;98 max_d = candidate_d;99 // The `[min_t .. max_t]` interval contains an intersection. End the linear search.100 } else {101 // No intersection yet. Carry on.102 min_t = candidate_t;103 min_d = candidate_d;105 for (var step = 1u; step < (*root_finder).linear_steps; step += 1u) {106 let candidate_t = mix(107 (*root_finder).min_t,108 (*root_finder).max_t,109 pow(110 (f32(step) + (*root_finder).jitter) / f32((*root_finder).linear_steps),111 (*root_finder).linear_march_exponent112 )113 );115 let candidate = start + dir * candidate_t;116 let candidate_d = depth_raymarch_distance_fn_evaluate(distance_fn, candidate);117 intersected = candidate_d.distance < 0.0 && candidate_d.valid;119 if (intersected) {120 max_t = candidate_t;121 max_d = candidate_d;122 // The `[min_t .. max_t]` interval contains an intersection.123 // End the linear search.124 break;125 } else {126 // No intersection yet. Carry on.127 min_t = candidate_t;128 min_d = candidate_d;129 }130 }131 }132 }134 *miss_t = min_t;135 *hit_t = min_t;137 //138 // Refine the hit using bisection140 if (intersected) {141 for (var step = 0u; step < (*root_finder).bisection_steps; step += 1u) {142 let mid_t = (min_t + max_t) * 0.5;143 let candidate = start + dir * mid_t;144 let candidate_d = depth_raymarch_distance_fn_evaluate(distance_fn, candidate);146 if (candidate_d.distance < 0.0 && candidate_d.valid) {147 // Intersection at the mid point. Refine the first half.148 max_t = mid_t;149 max_d = candidate_d;150 } else {151 // No intersection yet at the mid point. Refine the second half.152 min_t = mid_t;153 min_d = candidate_d;154 }155 }157 if ((*root_finder).use_secant) {158 // Finish with one application of the secant method159 let total_d = min_d.distance + -max_d.distance;161 let mid_t = mix(min_t, max_t, min_d.distance / total_d);162 let candidate = start + dir * mid_t;163 let candidate_d = depth_raymarch_distance_fn_evaluate(distance_fn, candidate);165 // Only accept the result of the secant method if it improves upon166 // the previous result.167 //168 // Technically root_finder should be `abs(candidate_d.distance) <169 // min(min_d.distance, -max_d.distance) * frac`, but root_finder seems170 // sufficient.171 if (abs(candidate_d.distance) < min_d.distance * 0.9 && candidate_d.valid) {172 *hit_t = mid_t;173 *hit_d = candidate_d;174 } else {175 *hit_t = max_t;176 *hit_d = max_d;177 }179 return true;180 } else {181 *hit_t = max_t;182 *hit_d = max_d;183 return true;184 }185 } else {186 // Mark the conservative miss distance.187 *hit_t = min_t;188 return false;189 }190}192struct DistanceWithPenetration {193 /// Distance to the surface of which a root we're trying to find194 distance: f32,196 /// Whether to consider this sample valid for intersection.197 /// Mostly relevant for allowing the ray marcher to travel behind surfaces,198 /// as it will mark surfaces it travels under as invalid.199 valid: bool,201 /// Conservative estimate of depth to which the ray penetrates the marched surface.202 penetration: f32,203}205struct DepthRaymarchDistanceFn {206 depth_tex_size: vec2<f32>,208 march_behind_surfaces: bool,209 depth_thickness: f32,211 use_sloppy_march: bool,212}214fn depth_raymarch_distance_fn_evaluate(215 distance_fn: ptr<function, DepthRaymarchDistanceFn>,216 ray_point_cs: vec3<f32>,217) -> DistanceWithPenetration {218 let interp_uv = ndc_to_uv(ray_point_cs.xy);220 let ray_depth = 1.0 / ray_point_cs.z;222 // We're using both point-sampled and bilinear-filtered values from the depth buffer.223 //224 // That's really stupid but works like magic. For samples taken near the ray origin,225 // the discrete nature of the depth buffer becomes a problem. It's not a land of continuous surfaces,226 // but a bunch of stacked duplo bricks.227 //228 // Technically we should be taking discrete steps in distance_fn duplo land, but then we're at the mercy229 // of arbitrary quantization of our directions -- and sometimes we'll take a step which would230 // claim that the ray is occluded -- even though the underlying smooth surface wouldn't occlude it.231 //232 // If we instead take linear taps from the depth buffer, we reconstruct the linear surface.233 // That fixes acne, but introduces false shadowing near object boundaries, as we now pretend234 // that everything is shrink-wrapped by distance_fn continuous 2.5D surface, and our depth thickness235 // heuristic ends up falling apart.236 //237 // The fix is to consider both the smooth and the discrete surfaces, and only claim occlusion238 // when the ray descends below both.239 //240 // The two approaches end up fixing each other's artifacts:241 // * The false occlusions due to duplo land are rejected because the ray stays above the smooth surface.242 // * The shrink-wrap surface is no longer continuous, so it's possible for rays to miss it.244 let linear_depth =245 1.0 / textureSampleLevel(depth_prepass_texture, depth_linear_sampler, interp_uv, 0u);246 let unfiltered_depth =247 1.0 / textureSampleLevel(depth_prepass_texture, depth_nearest_sampler, interp_uv, 0u);249 var max_depth: f32;250 var min_depth: f32;252 if ((*distance_fn).use_sloppy_march) {253 max_depth = unfiltered_depth;254 min_depth = unfiltered_depth;255 } else {256 max_depth = max(linear_depth, unfiltered_depth);257 min_depth = min(linear_depth, unfiltered_depth);258 }260 let bias = 0.000002;262 var res: DistanceWithPenetration;263 res.distance = max_depth * (1.0 + bias) - ray_depth;265 // distance_fn will be used at the end of the ray march to potentially discard the hit.266 res.penetration = ray_depth - min_depth;268 if ((*distance_fn).march_behind_surfaces) {269 res.valid = res.penetration < (*distance_fn).depth_thickness;270 } else {271 res.valid = true;272 }274 return res;275}277struct DepthRayMarchResult {278 /// True if the raymarch hit something.279 hit: bool,281 /// In case of a hit, the normalized distance to it.282 ///283 /// In case of a miss, the furthest the ray managed to travel, which could either be284 /// exceeding the max range, or getting behind a surface further than the depth thickness.285 ///286 /// Range: `0..=1` as a lerp factor over `ray_start_cs..=ray_end_cs`.287 hit_t: f32,289 /// UV corresponding to `hit_t`.290 hit_uv: vec2<f32>,292 /// The distance that the hit point penetrates into the hit surface.293 /// Will normally be non-zero due to limited precision of the ray march.294 ///295 /// In case of a miss: undefined.296 hit_penetration: f32,298 /// Ditto, within the range `0..DepthRayMarch::depth_thickness_linear_z`299 ///300 /// In case of a miss: undefined.301 hit_penetration_frac: f32,302}304struct DepthRayMarch {305 /// Number of steps to be taken at regular intervals to find an initial intersection.306 /// Must not be zero.307 linear_steps: u32,309 /// Exponent to be applied in the linear part of the march.310 ///311 /// A value of 1.0 will result in equidistant steps, and higher values will compress312 /// the earlier steps, and expand the later ones. This might be desirable in order313 /// to get more detail close to objects in SSR or SSGI.314 ///315 /// For optimal performance, this should be a small compile-time unsigned integer,316 /// such as 1 or 2.317 linear_march_exponent: f32,319 /// Number of steps in a bisection (binary search) to perform once the linear search320 /// has found an intersection. Helps narrow down the hit, increasing the chance of321 /// the secant method finding an accurate hit point.322 ///323 /// Useful when sampling color, e.g. SSR or SSGI, but pointless for contact shadows.324 bisection_steps: u32,326 /// Approximate the root position using the secant method -- by solving for line-line327 /// intersection between the ray approach rate and the surface gradient.328 ///329 /// Useful when sampling color, e.g. SSR or SSGI, but pointless for contact shadows.330 use_secant: bool,332 /// Jitter to apply to the first step of the linear search; 0..=1 range, mapping333 /// to the extent of a single linear step in the first phase of the search.334 /// Use 1.0 if you don't want jitter.335 jitter: f32,337 /// Clip space coordinates (w=1) of the ray.338 ray_start_cs: vec3<f32>,339 ray_end_cs: vec3<f32>,341 /// Should be used for contact shadows, but not for any color bounce, e.g. SSR.342 ///343 /// For SSR etc. this can easily create leaks, but with contact shadows it allows the rays344 /// to pass over invalid occlusions (due to thickness), and find potentially valid ones ahead.345 ///346 /// Note that this will cause the linear search to potentially miss surfaces,347 /// because when the ray overshoots and ends up penetrating a surface further than348 /// `depth_thickness_linear_z`, the ray marcher will just carry on.349 ///350 /// For this reason, this may require a lot of samples, or high depth thickness,351 /// so that `depth_thickness_linear_z >= world space ray length / linear_steps`.352 march_behind_surfaces: bool,354 /// If `true`, the ray marcher only performs nearest lookups of the depth buffer,355 /// resulting in aliasing and false occlusion when marching tiny detail.356 /// It should work fine for longer traces with fewer rays though.357 use_sloppy_march: bool,359 /// When marching the depth buffer, we only have 2.5D information, and don't know how360 /// thick surfaces are. We shall assume that the depth buffer fragments are little squares361 /// with a constant thickness defined by this parameter.362 depth_thickness_linear_z: f32,364 /// Size of the depth buffer we're marching in, in pixels.365 depth_tex_size: vec2<f32>,366}368fn depth_ray_march_new_from_depth(depth_tex_size: vec2<f32>) -> DepthRayMarch {369 var res: DepthRayMarch;370 res.jitter = 1.0;371 res.linear_steps = 4u;372 res.bisection_steps = 0u;373 res.linear_march_exponent = 1.0;374 res.depth_tex_size = depth_tex_size;375 res.depth_thickness_linear_z = 1.0;376 res.march_behind_surfaces = false;377 res.use_sloppy_march = false;378 return res;379}381fn depth_ray_march_to_cs_dir_impl(382 raymarch: ptr<function, DepthRayMarch>,383 dir_cs: vec4<f32>,384 infinite: bool,385) {386 var end_cs = vec4((*raymarch).ray_start_cs, 1.0) + dir_cs;388 // Perform perspective division, but avoid dividing by zero for rays389 // heading directly towards the eye.390 end_cs /= select(-1.0, 1.0, end_cs.w >= 0.0) * max(1e-10, abs(end_cs.w));392 // Clip ray start to the view frustum393 var delta_cs = end_cs.xyz - (*raymarch).ray_start_cs;394 let near_edge = select(vec3(-1.0, -1.0, 0.0), vec3(1.0, 1.0, 1.0), delta_cs < vec3(0.0));395 let dist_to_near_edge = (near_edge - (*raymarch).ray_start_cs) / delta_cs;396 let max_dist_to_near_edge = max(dist_to_near_edge.x, dist_to_near_edge.y);397 (*raymarch).ray_start_cs += delta_cs * max(0.0, max_dist_to_near_edge);399 // Clip ray end to the view frustum401 delta_cs = end_cs.xyz - (*raymarch).ray_start_cs;402 let far_edge = select(vec3(-1.0, -1.0, 0.0), vec3(1.0, 1.0, 1.0), delta_cs >= vec3(0.0));403 let dist_to_far_edge = (far_edge - (*raymarch).ray_start_cs) / delta_cs;404 let min_dist_to_far_edge = min(405 min(dist_to_far_edge.x, dist_to_far_edge.y),406 dist_to_far_edge.z407 );409 if (infinite) {410 delta_cs *= min_dist_to_far_edge;411 } else {412 // If unbounded, would make the ray reach the end of the frustum413 delta_cs *= min(1.0, min_dist_to_far_edge);414 }416 (*raymarch).ray_end_cs = (*raymarch).ray_start_cs + delta_cs;417}419/// March from a clip-space position (w = 1)420fn depth_ray_march_from_cs(raymarch: ptr<function, DepthRayMarch>, v: vec3<f32>) {421 (*raymarch).ray_start_cs = v;422}424/// March to a clip-space position (w = 1)425///426/// Must be called after `from_cs`, as it will clip the world-space ray to the view frustum.427fn depth_ray_march_to_cs(raymarch: ptr<function, DepthRayMarch>, end_cs: vec3<f32>) {428 let dir = vec4(end_cs - (*raymarch).ray_start_cs, 0.0) * sign(end_cs.z);429 depth_ray_march_to_cs_dir_impl(raymarch, dir, false);430}432/// March towards a clip-space direction. Infinite (ray is extended to cover the whole view frustum).433///434/// Must be called after `from_cs`, as it will clip the world-space ray to the view frustum.435fn depth_ray_march_to_cs_dir(raymarch: ptr<function, DepthRayMarch>, dir: vec4<f32>) {436 depth_ray_march_to_cs_dir_impl(raymarch, dir, true);437}439/// March to a world-space position.440///441/// Must be called after `from_cs`, as it will clip the world-space ray to the view frustum.442fn depth_ray_march_to_ws(raymarch: ptr<function, DepthRayMarch>, end: vec3<f32>) {443 depth_ray_march_to_cs(raymarch, position_world_to_ndc(end));444}446/// March towards a world-space direction. Infinite (ray is extended to cover the whole view frustum).447///448/// Must be called after `from_cs`, as it will clip the world-space ray to the view frustum.449fn depth_ray_march_to_ws_dir(raymarch: ptr<function, DepthRayMarch>, dir: vec3<f32>) {450 depth_ray_march_to_cs_dir_impl(raymarch, direction_world_to_clip(dir), true);451}453/// Perform the ray march.454fn depth_ray_march_march(raymarch: ptr<function, DepthRayMarch>) -> DepthRayMarchResult {455 var res = DepthRayMarchResult(false, 0.0, vec2(0.0), 0.0, 0.0);457 let ray_start_uv = ndc_to_uv((*raymarch).ray_start_cs.xy);458 let ray_end_uv = ndc_to_uv((*raymarch).ray_end_cs.xy);460 let ray_uv_delta = ray_end_uv - ray_start_uv;461 let ray_len_px = ray_uv_delta * (*raymarch).depth_tex_size;463 let min_px_per_step = 1u;464 let step_count = max(465 2,466 min(i32((*raymarch).linear_steps), i32(floor(length(ray_len_px) / f32(min_px_per_step))))467 );469 let linear_z_to_scaled_linear_z = 1.0 / perspective_camera_near();470 let depth_thickness = (*raymarch).depth_thickness_linear_z * linear_z_to_scaled_linear_z;472 var distance_fn: DepthRaymarchDistanceFn;473 distance_fn.depth_tex_size = (*raymarch).depth_tex_size;474 distance_fn.march_behind_surfaces = (*raymarch).march_behind_surfaces;475 distance_fn.depth_thickness = depth_thickness;476 distance_fn.use_sloppy_march = (*raymarch).use_sloppy_march;478 var hit: DistanceWithPenetration;480 var hit_t = 0.0;481 var miss_t = 0.0;482 var root_finder = hybrid_root_finder_new_with_linear_steps(u32(step_count));483 root_finder.bisection_steps = (*raymarch).bisection_steps;484 root_finder.use_secant = (*raymarch).use_secant;485 root_finder.linear_march_exponent = (*raymarch).linear_march_exponent;486 root_finder.jitter = (*raymarch).jitter;487 let intersected = hybrid_root_finder_find_root(488 &root_finder,489 (*raymarch).ray_start_cs,490 (*raymarch).ray_end_cs,491 &distance_fn,492 &hit_t,493 &miss_t,494 &hit495 );497 res.hit_t = hit_t;499 if (intersected && hit.penetration < depth_thickness && hit.distance < depth_thickness) {500 res.hit = true;501 res.hit_uv = mix(ray_start_uv, ray_end_uv, res.hit_t);502 res.hit_penetration = hit.penetration / linear_z_to_scaled_linear_z;503 res.hit_penetration_frac = hit.penetration / depth_thickness;504 return res;505 }507 res.hit_t = miss_t;508 res.hit_uv = mix(ray_start_uv, ray_end_uv, res.hit_t);510 return res;511}