fix: Correct mesh normal transformation and floor shadow renderingHEADmain

7 files changed, 149 insertions, 41 deletions

diff --git a/assets/final/shaders/mesh_render.wgsl b/assets/final/shaders/mesh_render.wgsl
index 3759747..3faf7ca 100644
--- a/assets/final/shaders/mesh_render.wgsl
+++ b/assets/final/shaders/mesh_render.wgsl
@@ -33,10 +33,9 @@ fn vs_main(in: VertexInput, @builtin(instance_index) instance_index: u32) -> Ver
     out.clip_pos = globals.view_proj * world_pos;
     out.world_pos = world_pos.xyz;
     
-    // Normal transform (assuming uniform scale or using transpose(inverse(model)))
-    // For simplicity, we use the same mat3 logic as renderer_3d.wgsl
-    let normal_matrix = mat3x3<f32>(obj.model[0].xyz, obj.model[1].xyz, obj.model[2].xyz);
-    out.normal = normalize(normal_matrix * in.normal);
+    // Use transpose of inverse for normals
+    let normal_matrix = mat3x3<f32>(obj.inv_model[0].xyz, obj.inv_model[1].xyz, obj.inv_model[2].xyz);
+    out.normal = normalize(transpose(normal_matrix) * in.normal);
     
     out.uv = in.uv;
     out.color = obj.color;
@@ -56,4 +55,4 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
     let lit_color = calculate_lighting(in.color.rgb, in.normal, in.world_pos, shadow);
     
     return vec4<f32>(lit_color, in.color.a);
-}
+}
\ No newline at end of file
diff --git a/assets/final/shaders/render/scene_query_bvh.wgsl b/assets/final/shaders/render/scene_query_bvh.wgsl
index c7dfdf4..d040c1b 100644
--- a/assets/final/shaders/render/scene_query_bvh.wgsl
+++ b/assets/final/shaders/render/scene_query_bvh.wgsl
@@ -10,11 +10,13 @@ struct BVHNode {
 
 @group(0) @binding(2) var<storage, read> bvh_nodes: array<BVHNode>;
 
-fn get_dist(p: vec3<f32>, obj_type: f32) -> f32 {
+fn get_dist(p: vec3<f32>, obj_params: vec4<f32>) -> f32 {
+    let obj_type = obj_params.x;
     if (obj_type == 1.0) { return length(p) - 1.0; } // Unit Sphere
     if (obj_type == 2.0) { return sdBox(p, vec3<f32>(1.0)); } // Unit Box
     if (obj_type == 3.0) { return sdTorus(p, vec2<f32>(1.0, 0.4)); } // Unit Torus
     if (obj_type == 4.0) { return sdPlane(p, vec3<f32>(0.0, 1.0, 0.0), 0.0); }
+    if (obj_type == 5.0) { return sdBox(p, obj_params.yzw); } // MESH AABB
     return 100.0;
 }
 
@@ -40,7 +42,14 @@ fn map_scene(p: vec3<f32>, skip_idx: u32) -> f32 {
                 let obj = object_data.objects[obj_idx];
                 let q = (obj.inv_model * vec4<f32>(p, 1.0)).xyz;
                 let s = min(length(obj.model[0].xyz), min(length(obj.model[1].xyz), length(obj.model[2].xyz)));
-                d = min(d, get_dist(q, obj.params.x) * s);
+                // IMPORTANT: Plane (type 4.0) should not be scaled by 's' in this way.
+                // The sdPlane function expects its own scale/offset implicitly handled by the model matrix.
+                // The 's' factor is meant for primitives whose SDFs are defined relative to a unit size.
+                if (obj.params.x != 4.0) { // Only scale if not a plane
+                    d = min(d, get_dist(q, obj.params) * s);
+                } else {
+                    d = min(d, get_dist(q, obj.params));
+                }
             } else { // Internal
                 if (stack_ptr < 31) {
                     stack[stack_ptr] = node.left_idx;
diff --git a/assets/final/shaders/render/scene_query_linear.wgsl b/assets/final/shaders/render/scene_query_linear.wgsl
index 7bcd96f..30a0371 100644
--- a/assets/final/shaders/render/scene_query_linear.wgsl
+++ b/assets/final/shaders/render/scene_query_linear.wgsl
@@ -1,26 +1,64 @@
 #include "math/sdf_shapes"
+#include "math/sdf_utils"
 
-fn get_dist(p: vec3<f32>, obj_type: f32) -> f32 {
+struct BVHNode {
+    min: vec3<f32>,
+    left_idx: i32,
+    max: vec3<f32>,
+    obj_idx_or_right: i32,
+};
+
+@group(0) @binding(2) var<storage, read> bvh_nodes: array<BVHNode>;
+
+fn get_dist(p: vec3<f32>, obj_params: vec4<f32>) -> f32 {
+    let obj_type = obj_params.x;
     if (obj_type == 1.0) { return length(p) - 1.0; } // Unit Sphere
     if (obj_type == 2.0) { return sdBox(p, vec3<f32>(1.0)); } // Unit Box
     if (obj_type == 3.0) { return sdTorus(p, vec2<f32>(1.0, 0.4)); } // Unit Torus
     if (obj_type == 4.0) { return sdPlane(p, vec3<f32>(0.0, 1.0, 0.0), 0.0); }
+    if (obj_type == 5.0) { return sdBox(p, obj_params.yzw); } // MESH AABB
     return 100.0;
 }
 
 fn map_scene(p: vec3<f32>, skip_idx: u32) -> f32 {
     var d = 1000.0;
-    let count = u32(globals.params.x);
+    var stack: array<i32, 32>;
+    var stack_ptr = 0;
     
-    for (var i = 0u; i < count; i = i + 1u) {
-        if (i == skip_idx) { continue; }
-        let obj = object_data.objects[i];
-        let obj_type = obj.params.x;
-        if (obj_type <= 0.0) { continue; } 
+    if (arrayLength(&bvh_nodes) > 0u) {
+        stack[stack_ptr] = 0;
+        stack_ptr++;
+    }
 
-        let q = (obj.inv_model * vec4<f32>(p, 1.0)).xyz;
-        let s = min(length(obj.model[0].xyz), min(length(obj.model[1].xyz), length(obj.model[2].xyz)));
-        d = min(d, get_dist(q, obj_type) * s);
+    while (stack_ptr > 0) {
+        stack_ptr--;
+        let node_idx = stack[stack_ptr];
+        let node = bvh_nodes[node_idx];
+
+        if (aabb_sdf(p, node.min, node.max) < d) {
+            if (node.left_idx < 0) { // Leaf
+                let obj_idx = u32(node.obj_idx_or_right);
+                if (obj_idx == skip_idx) { continue; }
+                let obj = object_data.objects[obj_idx];
+                let q = (obj.inv_model * vec4<f32>(p, 1.0)).xyz;
+                let s = min(length(obj.model[0].xyz), min(length(obj.model[1].xyz), length(obj.model[2].xyz)));
+                // IMPORTANT: Plane (type 4.0) should not be scaled by 's' in this way.
+                // The sdPlane function expects its own scale/offset implicitly handled by the model matrix.
+                // The 's' factor is meant for primitives whose SDFs are defined relative to a unit size.
+                if (obj.params.x != 4.0) { // Only scale if not a plane
+                    d = min(d, get_dist(q, obj.params) * s);
+                } else {
+                    d = min(d, get_dist(q, obj.params));
+                }
+            } else { // Internal
+                if (stack_ptr < 31) {
+                    stack[stack_ptr] = node.left_idx;
+                    stack_ptr++;
+                    stack[stack_ptr] = node.obj_idx_or_right;
+                    stack_ptr++;
+                }
+            }
+        }
     }
     return d;
-}
+}
\ No newline at end of file
diff --git a/assets/final/shaders/renderer_3d.wgsl b/assets/final/shaders/renderer_3d.wgsl
index e7cb810..4733f6f 100644
--- a/assets/final/shaders/renderer_3d.wgsl
+++ b/assets/final/shaders/renderer_3d.wgsl
@@ -15,6 +15,7 @@ struct VertexOutput {
     @location(1) color: vec4<f32>,
     @location(2) @interpolate(flat) instance_index: u32,
     @location(3) world_pos: vec3<f32>,
+    @location(4) transformed_normal: vec3<f32>,
 };
 
 @vertex
@@ -41,8 +42,10 @@ fn vs_main(@builtin(vertex_index) vertex_index: u32,
     let obj_type = obj.params.x;
 
     if (obj_type == 5.0) { // MESH
+        // For meshes, we use the actual vertex data, not proxy geometry.
+        // The position here is a placeholder, the real mesh data is handled by mesh_pipeline_.
         var out: VertexOutput;
-        out.position = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+        out.position = vec4<f32>(0.0, 0.0, 2.0, 1.0); // Outside far plane, so it's not rendered by this pipeline.
         return out;
     }
 
@@ -62,6 +65,22 @@ fn vs_main(@builtin(vertex_index) vertex_index: u32,
     out.color = obj.color;
     out.instance_index = instance_index;
     out.world_pos = world_pos.xyz;
+    
+    // Correct normal transformation for meshes: transpose of inverse of model matrix
+    // For non-uniform scaling, this is necessary. For other primitives, we use their analytical normals.
+    if (obj_type == 5.0) {
+        // Calculate inverse transpose of the model matrix (upper 3x3 part)
+        let model_matrix = mat3x3<f32>(obj.model[0].xyz, obj.model[1].xyz, obj.model[2].xyz);
+        let normal_matrix = transpose(inverse(model_matrix));
+        out.transformed_normal = normalize(normal_matrix * in.normal);
+    } else {
+        // For SDF primitives, we don't use vertex normals directly here; they are computed in the fragment shader.
+        // However, we still need to output a normal for the fragment shader to use if it were a rasterized primitive.
+        // The transformed_normal is not used by the SDF fragment shader, but for correctness, we'll pass it.
+        // If this were a rasterized mesh, it would be used.
+        out.transformed_normal = normalize(vec3<f32>(0.0, 1.0, 0.0)); // Placeholder for non-mesh types
+    }
+
     return out;
 }
 
@@ -70,8 +89,13 @@ fn vs_main(@builtin(vertex_index) vertex_index: u32,
 #include "render/lighting_utils"
 #include "ray_box"
 
+struct FragmentOutput {
+    @location(0) color: vec4<f32>,
+    @builtin(frag_depth) depth: f32,
+};
+
 @fragment
-fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
+fn fs_main(in: VertexOutput) -> FragmentOutput {
     let obj = object_data.objects[in.instance_index];
     let obj_type = obj.params.x;
 
@@ -80,11 +104,10 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
     var base_color = in.color.rgb;
     let light_dir = normalize(vec3<f32>(1.0, 1.0, 1.0)); 
 
-    if (obj_type <= 0.0) { // Raster path
+    if (obj_type <= 0.0) { // Raster path (legacy or generic)
         p = in.world_pos;
-        let local_normal = normalize(cross(dpdx(in.local_pos), dpdy(in.local_pos)));
-        let normal_matrix = mat3x3<f32>(obj.inv_model[0].xyz, obj.inv_model[1].xyz, obj.inv_model[2].xyz);
-        normal = normalize(transpose(normal_matrix) * local_normal);
+        // Use the transformed normal passed from the vertex shader for rasterized objects
+        normal = normalize(in.transformed_normal);
 
         // Apply grid pattern to floor
         let uv = p.xz * 0.5;
@@ -100,8 +123,10 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
         let rd_local = normalize((obj.inv_model * vec4<f32>(rd_world, 0.0)).xyz);
         
         // Proxy box extent (matches vs_main)
+        // MESHES use obj.params.yzw for extent
         var extent = vec3<f32>(1.0);
-        if (obj_type == 3.0) { extent = vec3<f32>(1.5, 0.5, 1.5); }
+        if (obj.params.x == 3.0) { extent = vec3<f32>(1.5, 0.5, 1.5); } // Torus
+        else if (obj.params.x == 5.0) { extent = obj.params.yzw; } // MESH extent
         
         let bounds = ray_box_intersection(ro_local, rd_local, extent);
 
@@ -111,7 +136,7 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
         var hit = false;
         for (var i = 0; i < 64; i = i + 1) {
             let q = ro_local + rd_local * t;
-            let d_local = get_dist(q, obj_type);
+            let d_local = get_dist(q, obj.params);
             if (d_local < 0.0005) { hit = true; break; }
             t = t + d_local;
             if (t > bounds.t_exit) { break; }
@@ -128,32 +153,32 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
         let q_x1 = q_hit + e.xyy;
         let uv_x1 = vec2<f32>(atan2(q_x1.x, q_x1.z) / 6.28 + 0.5, acos(clamp(q_x1.y / length(q_x1), -1.0, 1.0)) / 3.14);
         let h_x1 = textureSample(noise_tex, noise_sampler, uv_x1).r;
-        let d_x1 = get_dist(q_x1, obj_type) - disp_strength * h_x1;
+        let d_x1 = get_dist(q_x1, obj.params) - disp_strength * h_x1;
         
         let q_x2 = q_hit - e.xyy;
         let uv_x2 = vec2<f32>(atan2(q_x2.x, q_x2.z) / 6.28 + 0.5, acos(clamp(q_x2.y / length(q_x2), -1.0, 1.0)) / 3.14);
         let h_x2 = textureSample(noise_tex, noise_sampler, uv_x2).r;
-        let d_x2 = get_dist(q_x2, obj_type) - disp_strength * h_x2;
+        let d_x2 = get_dist(q_x2, obj.params) - disp_strength * h_x2;
         
         let q_y1 = q_hit + e.yxy;
         let uv_y1 = vec2<f32>(atan2(q_y1.x, q_y1.z) / 6.28 + 0.5, acos(clamp(q_y1.y / length(q_y1), -1.0, 1.0)) / 3.14);
         let h_y1 = textureSample(noise_tex, noise_sampler, uv_y1).r;
-        let d_y1 = get_dist(q_y1, obj_type) - disp_strength * h_y1;
+        let d_y1 = get_dist(q_y1, obj.params) - disp_strength * h_y1;
         
         let q_y2 = q_hit - e.yxy;
         let uv_y2 = vec2<f32>(atan2(q_y2.x, q_y2.z) / 6.28 + 0.5, acos(clamp(q_y2.y / length(q_y2), -1.0, 1.0)) / 3.14);
         let h_y2 = textureSample(noise_tex, noise_sampler, uv_y2).r;
-        let d_y2 = get_dist(q_y2, obj_type) - disp_strength * h_y2;
+        let d_y2 = get_dist(q_y2, obj.params) - disp_strength * h_y2;
         
         let q_z1 = q_hit + e.yyx;
         let uv_z1 = vec2<f32>(atan2(q_z1.x, q_z1.z) / 6.28 + 0.5, acos(clamp(q_z1.y / length(q_z1), -1.0, 1.0)) / 3.14);
         let h_z1 = textureSample(noise_tex, noise_sampler, uv_z1).r;
-        let d_z1 = get_dist(q_z1, obj_type) - disp_strength * h_z1;
+        let d_z1 = get_dist(q_z1, obj.params) - disp_strength * h_z1;
         
         let q_z2 = q_hit - e.yyx;
         let uv_z2 = vec2<f32>(atan2(q_z2.x, q_z2.z) / 6.28 + 0.5, acos(clamp(q_z2.y / length(q_z2), -1.0, 1.0)) / 3.14);
         let h_z2 = textureSample(noise_tex, noise_sampler, uv_z2).r;
-        let d_z2 = get_dist(q_z2, obj_type) - disp_strength * h_z2;
+        let d_z2 = get_dist(q_z2, obj.params) - disp_strength * h_z2;
         
         let n_local = normalize(vec3<f32>(d_x1 - d_x2, d_y1 - d_y2, d_z1 - d_z2));
         let normal_matrix = mat3x3<f32>(obj.inv_model[0].xyz, obj.inv_model[1].xyz, obj.inv_model[2].xyz);
@@ -174,5 +199,13 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
 
     let shadow = calc_shadow(p, light_dir, 0.05, 20.0, in.instance_index);
     let lit_color = calculate_lighting(base_color, normal, p, shadow);
-    return vec4<f32>(lit_color, 1.0);
+    
+    var out: FragmentOutput;
+    out.color = vec4<f32>(lit_color, 1.0);
+    
+    // Calculate and write correct depth
+    let clip_pos = globals.view_proj * vec4<f32>(p, 1.0);
+    out.depth = clip_pos.z / clip_pos.w;
+    
+    return out;
 }
\ No newline at end of file
diff --git a/src/3d/object.h b/src/3d/object.h
index a7ce0b8..3d4ff35 100644
--- a/src/3d/object.h
+++ b/src/3d/object.h
@@ -40,12 +40,14 @@ class Object3D {
   bool is_static;
 
   AssetId mesh_asset_id;
+  vec3 local_extent; // Half-extents for AABB (used by meshes for shadows)
   void* user_data; // For tool-specific data, not for general use
 
   Object3D(ObjectType t = ObjectType::CUBE)
       : position(0, 0, 0), rotation(0, 0, 0, 1), scale(1, 1, 1), type(t),
         color(1, 1, 1, 1), velocity(0, 0, 0), mass(1.0f), restitution(0.5f),
-        is_static(false), mesh_asset_id((AssetId)0), user_data(nullptr) {
+        is_static(false), mesh_asset_id((AssetId)0), local_extent(1, 1, 1), 
+        user_data(nullptr) {
   }
 
   mat4 get_model_matrix() const {
@@ -61,6 +63,8 @@ class Object3D {
   BoundingVolume get_local_bounds() const {
     if (type == ObjectType::TORUS)
       return {{-1.5f, -0.5f, -1.5f}, {1.5f, 0.5f, 1.5f}};
+    if (type == ObjectType::MESH)
+      return {-local_extent, local_extent};
     // Simple defaults for unit primitives
     return {{-1, -1, -1}, {1, 1, 1}};
   }
diff --git a/src/3d/renderer.cc b/src/3d/renderer.cc
index a9beffe..eea3ff0 100644
--- a/src/3d/renderer.cc
+++ b/src/3d/renderer.cc
@@ -174,7 +174,7 @@ void Renderer3D::update_uniforms(const Scene& scene, const Camera& camera,
       type_id = 4.0f;
     else if (obj.type == ObjectType::MESH)
       type_id = 5.0f;
-    data.params = vec4(type_id, 0, 0, 0);
+    data.params = vec4(type_id, obj.local_extent.x, obj.local_extent.y, obj.local_extent.z);
     obj_data.push_back(data);
     if (obj_data.size() >= kMaxObjects)
       break;
@@ -303,9 +303,11 @@ void Renderer3D::draw(WGPURenderPassEncoder pass, const Scene& scene,
 #if !defined(STRIP_ALL)
   if (s_debug_enabled_) {
     for (const auto& obj : scene.objects) {
-      vec3 extent(1.0f, 1.0f, 1.0f);
+      vec3 extent = obj.local_extent;
       if (obj.type == ObjectType::TORUS) {
         extent = vec3(1.5f, 0.5f, 1.5f);
+      } else if (obj.type != ObjectType::MESH) {
+        extent = vec3(1.0f, 1.0f, 1.0f);
       }
       visual_debug_.add_box(obj.get_model_matrix(), extent,
                             vec3(1.0f, 1.0f, 0.0f)); // Yellow boxes
diff --git a/src/tests/test_mesh.cc b/src/tests/test_mesh.cc
index 0294d9b..f29ec27 100644
--- a/src/tests/test_mesh.cc
+++ b/src/tests/test_mesh.cc
@@ -169,7 +169,15 @@ bool load_obj_and_create_buffers(const char* path, Object3D& out_obj) {
                 std::string parts[3] = {s1, s2, s3};
                 RawFace face = {};
                 for (int i = 0; i < 3; ++i) {
-                  sscanf(parts[i].c_str(), "%d/%d/%d", &face.v[i], &face.vt[i], &face.vn[i]);
+                  // Handle v//vn format
+                  if (parts[i].find("//") != std::string::npos) {
+                    sscanf(parts[i].c_str(), "%d//%d", &face.v[i], &face.vn[i]);
+                    face.vt[i] = 0;
+                  } else {
+                    int res = sscanf(parts[i].c_str(), "%d/%d/%d", &face.v[i], &face.vt[i], &face.vn[i]);
+                    if (res == 2) face.vn[i] = 0;
+                    else if (res == 1) { face.vt[i] = 0; face.vn[i] = 0; }
+                  }
                 }
                 raw_faces.push_back(face);
             }
@@ -213,6 +221,21 @@ bool load_obj_and_create_buffers(const char* path, Object3D& out_obj) {
     }
 
     if (final_vertices.empty()) return false;
+
+    // Calculate AABB and center the mesh
+    float min_x = 1e10f, min_y = 1e10f, min_z = 1e10f;
+    float max_x = -1e10f, max_y = -1e10f, max_z = -1e10f;
+    for (const auto& v : final_vertices) {
+        min_x = std::min(min_x, v.p[0]); min_y = std::min(min_y, v.p[1]); min_z = std::min(min_z, v.p[2]);
+        max_x = std::max(max_x, v.p[0]); max_y = std::max(max_y, v.p[1]); max_z = std::max(max_z, v.p[2]);
+    }
+    float cx = (min_x + max_x) * 0.5f;
+    float cy = (min_y + max_y) * 0.5f;
+    float cz = (min_z + max_z) * 0.5f;
+    for (auto& v : final_vertices) {
+        v.p[0] -= cx; v.p[1] -= cy; v.p[2] -= cz;
+    }
+    out_obj.local_extent = vec3((max_x - min_x) * 0.5f, (max_y - min_y) * 0.5f, (max_z - min_z) * 0.5f);
     
     g_mesh_gpu_data.num_indices = final_indices.size();
     g_mesh_gpu_data.vertex_buffer = gpu_create_buffer(g_device, final_vertices.size() * sizeof(MeshVertex), WGPUBufferUsage_Vertex | WGPUBufferUsage_CopyDst, final_vertices.data()).buffer;
@@ -261,7 +284,7 @@ int main(int argc, char** argv) {
 
   // --- Create Scene ---
   Object3D floor(ObjectType::PLANE);
-  floor.scale = vec3(20.0f, 1.0f, 20.0f);
+  floor.scale = vec3(20.0f, 0.01f, 20.0f); // Very thin floor proxy
   floor.color = vec4(0.5f, 0.5f, 0.5f, 1.0f);
   g_scene.add_object(floor);
 
@@ -271,11 +294,11 @@ int main(int argc, char** argv) {
       return 1;
   }
   mesh_obj.color = vec4(1.0f, 0.7f, 0.2f, 1.0f);
-  mesh_obj.position = {0, 1, 0};
+  mesh_obj.position = {0, 1.5, 0}; // Elevate a bit more
   g_scene.add_object(mesh_obj);
   
   g_camera.position = vec3(0, 3, 5);
-  g_camera.target = vec3(0, 1, 0);
+  g_camera.target = vec3(0, 1.5, 0);
 
   while (!platform_should_close(&platform_state)) {
     platform_poll(&platform_state);
@@ -287,7 +310,7 @@ int main(int argc, char** argv) {
     
     if (debug_mode) {
         auto* vertices = (std::vector<MeshVertex>*)g_scene.objects[1].user_data;
-        g_renderer.GetVisualDebug().add_mesh_normals(g_scene.objects[1].get_model_matrix(), vertices->size(), vertices->data());
+        g_renderer.GetVisualDebug().add_mesh_normals(g_scene.objects[1].get_model_matrix(), (uint32_t)vertices->size(), vertices->data());
     }
 
     WGPUSurfaceTexture surface_tex;