Consolidate ray normal functions.

PiperOrigin-RevId: 842300829
Change-Id: Ic5e921084710cb044a4f49ea9b566d277838aaa1

src/engine/engine_ray.c

@@ -197,6 +197,7 @@ mjtNum ray_triangle(mjtNum v[][3], const mjtNum lpnt[3], const mjtNum lvec[3],
   return x;
 }
 
+
 //---------------------------- geom-specific intersection functions --------------------------------
 
 // plane
@@ -558,8 +559,8 @@ static mjtNum ray_box(const mjtNum pos[3], const mjtNum mat[9], const mjtNum siz
 
 
 // intersect ray with hfield, compute normal if given
-mjtNum mj_rayHfieldNormal(const mjModel* m, const mjData* d, int geomid,
+static mjtNum mj_rayHfieldNormal(const mjModel* m, const mjData* d, int geomid,
-                          const mjtNum pnt[3], const mjtNum vec[3], mjtNum normal[3]) {
+                                 const mjtNum pnt[3], const mjtNum vec[3], mjtNum normal[3]) {
   // clear normal if given
   if (normal) mju_zero3(normal);
 
@@ -736,6 +737,7 @@ mjtNum mj_rayHfieldNormal(const mjModel* m, const mjData* d, int geomid,
   return x;
 }
 
+
 // intersect ray with hfield
 mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int geomid,
                     const mjtNum pnt[3], const mjtNum vec[3]) {
@@ -885,10 +887,11 @@ mjtNum mju_rayTree(const mjModel* m, const mjData* d, int id, const mjtNum pnt[3
   return x;
 }
 
+
 // intersect ray with signed distance field, compute normal if given
-mjtNum mj_raySdfNormal(const mjModel* m, const mjData* d, int g,
+static mjtNum mj_raySdfNormal(const mjModel* m, const mjData* d, int g,
-                       const mjtNum pnt[3], const mjtNum vec[3],
+                              const mjtNum pnt[3], const mjtNum vec[3],
-                       mjtNum normal[3]) {
+                              mjtNum normal[3]) {
   if (normal) mju_zero3(normal);
 
   mjtNum distance_total = 0;
@@ -953,15 +956,17 @@ mjtNum mj_raySdfNormal(const mjModel* m, const mjData* d, int g,
   return -1;
 }
 
+
 // intersect ray with signed distance field
 static mjtNum ray_sdf(const mjModel* m, const mjData* d, int g,
                       const mjtNum pnt[3], const mjtNum vec[3]) {
   return mj_raySdfNormal(m, d, g, pnt, vec, NULL);
 }
 
+
 // intersect ray with mesh, compute normal if given
-mjtNum mj_rayMeshNormal(const mjModel* m, const mjData* d, int id, const mjtNum pnt[3],
+static mjtNum mj_rayMeshNormal(const mjModel* m, const mjData* d, int id, const mjtNum pnt[3],
-                        const mjtNum vec[3], mjtNum normal[3]) {
+                               const mjtNum vec[3], mjtNum normal[3]) {
   // clear normal if given
   if (normal) mju_zero3(normal);
 
@@ -978,15 +983,18 @@ mjtNum mj_rayMeshNormal(const mjModel* m, const mjData* d, int id, const mjtNum
   return mju_rayTree(m, d, id, pnt, vec, normal);
 }
 
+
 // intersect ray with mesh
 mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int id, const mjtNum pnt[3],
                   const mjtNum vec[3]) {
   return mj_rayMeshNormal(m, d, id, pnt, vec, NULL);
 }
 
+
 // intersect ray and find normal with primitive geom, no meshes or hfields, compute normal if given
 mjtNum mju_rayGeomNormal(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3],
-                         const mjtNum pnt[3], const mjtNum vec[3], int geomtype, mjtNum normal[3]) {
+                         const mjtNum pnt[3], const mjtNum vec[3], int geomtype,
+                         mjtNum normal[3]) {
   switch ((mjtGeom) geomtype) {
   case mjGEOM_PLANE:
     return ray_plane(pos, mat, size, pnt, vec, normal);
@@ -1019,6 +1027,7 @@ mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3
   return mju_rayGeomNormal(pos, mat, size, pnt, vec, geomtype, NULL);
 }
 
+
 // intersect ray with flex, return nearest vertex id
 mjtNum mju_rayFlex(const mjModel* m, const mjData* d, int flex_layer, mjtByte flg_vert,
                    mjtByte flg_edge, mjtByte flg_face, mjtByte flg_skin, int flexid,
@@ -1314,9 +1323,13 @@ static int point_in_box(const mjtNum aabb[6], const mjtNum xpos[3],
 // intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms on bodyexclude
 //  return geomid and distance (x) to nearest surface, or -1 if no intersection
 //  geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion
-mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
+mjtNum mj_rayNormal(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
-              const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude, int geomid[1]) {
+                    const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                    int geomid[1], mjtNum normal[3]) {
+  int ngeom = m->ngeom;
   mjtNum dist, newdist;
+  mjtNum normal_local[3];
+  mjtNum* p_normal = normal ? normal_local : NULL;
 
   // check vector length
   if (mju_norm3(vec) < mjMINVAL) {
@@ -1326,29 +1339,28 @@ mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtN
   // clear result
   dist = -1;
   if (geomid) *geomid = -1;
+  if (normal) mju_zero3(normal);
 
   // loop over geoms not eliminated by mask and bodyexclude
-  for (int i=0; i < m->ngeom; i++) {
+  for (int i=0; i < ngeom; i++) {
     if (!ray_eliminate(m, d, i, geomgroup, flg_static, bodyexclude)) {
-      // handle mesh and hfield separately
+      int type = m->geom_type[i];
-      if (m->geom_type[i] == mjGEOM_MESH) {
+      if (type == mjGEOM_MESH) {
-        newdist = mj_rayMesh(m, d, i, pnt, vec);
+        newdist = mj_rayMeshNormal(m, d, i, pnt, vec, p_normal);
-      } else if (m->geom_type[i] == mjGEOM_HFIELD) {
+      } else if (type == mjGEOM_HFIELD) {
-        newdist = mj_rayHfield(m, d, i, pnt, vec);
+        newdist = mj_rayHfieldNormal(m, d, i, pnt, vec, p_normal);
-      } else if (m->geom_type[i] == mjGEOM_SDF) {
+      } else if (type == mjGEOM_SDF) {
-        newdist = ray_sdf(m, d, i, pnt, vec);
+        newdist = mj_raySdfNormal(m, d, i, pnt, vec, p_normal);
-      }
+      } else {
-
+        newdist = mju_rayGeomNormal(d->geom_xpos+3*i, d->geom_xmat+9*i,
-      // otherwise general dispatch
+                                    m->geom_size+3*i, pnt, vec, type, p_normal);
-      else {
-        newdist = mju_rayGeom(d->geom_xpos+3*i, d->geom_xmat+9*i,
-                              m->geom_size+3*i, pnt, vec, m->geom_type[i]);
       }
 
       // update if closer intersection found
       if (newdist >= 0 && (newdist < dist || dist < 0)) {
         dist = newdist;
         if (geomid) *geomid = i;
+        if (normal) mju_copy3(normal, normal_local);
       }
     }
   }
@@ -1357,6 +1369,13 @@ mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtN
 }
 
 
+// intersect ray
+mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
+              const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude, int geomid[1]) {
+  return mj_rayNormal(m, d, pnt, vec, geomgroup, flg_static, bodyexclude, geomid, NULL);
+}
+
+
 // Initializes spherical bounding angles (geom_ba) and flag vector for a given source
 void mju_multiRayPrepare(const mjModel* m, const mjData* d, const mjtNum pnt[3],
                          const mjtNum ray_xmat[9], const mjtByte* geomgroup, mjtByte flg_static,

src/engine/engine_ray.h

@@ -46,10 +46,6 @@ MJAPI mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], cons
 MJAPI mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int geomid,
                           const mjtNum pnt[3], const mjtNum vec[3]);
 
-// intersect ray with hfield, compute normal if given
-MJAPI mjtNum mj_rayHfieldNormal(const mjModel* m, const mjData* d, int geomid,
-                                const mjtNum pnt[3], const mjtNum vec[3], mjtNum normal[3]);
-
 // intersect ray with triangle
 MJAPI mjtNum ray_triangle(mjtNum v[][3], const mjtNum lpnt[3], const mjtNum lvec[3],
                           const mjtNum b0[3], const mjtNum b1[3], mjtNum normal[3]);
@@ -58,22 +54,14 @@ MJAPI mjtNum ray_triangle(mjtNum v[][3], const mjtNum lpnt[3], const mjtNum lvec
 MJAPI mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int geomid,
                         const mjtNum pnt[3], const mjtNum vec[3]);
 
-// intersect ray with mesh, compute normal if given
-MJAPI mjtNum mj_rayMeshNormal(const mjModel* m, const mjData* d, int geomid,
-                              const mjtNum pnt[3], const mjtNum vec[3], mjtNum normal[3]);
-
-// intersect ray with signed distance field, compute normal if given
-MJAPI mjtNum mj_raySdfNormal(const mjModel* m, const mjData* d, int geomid,
-                             const mjtNum pnt[3], const mjtNum vec[3], mjtNum normal[3]);
-
 // intersect ray with primitive geom, no meshes or hfields
 MJAPI mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3],
                          const mjtNum pnt[3], const mjtNum vec[3], int geomtype);
 
-// intersect ray with primitive geom, no meshes or hfields, compute normal if given
+// intersect ray with geom, compute normal if given
-MJAPI mjtNum mju_rayGeomNormal(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3],
+MJAPI mjtNum mj_rayNormal(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
-                               const mjtNum pnt[3], const mjtNum vec[3], int geomtype,
+                          const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
-                               mjtNum normal[3]);
+                          int geomid[1], mjtNum normal[3]);
 
 // intersect ray with flex, return nearest vertex id
 MJAPI mjtNum mju_rayFlex(const mjModel* m, const mjData* d, int flex_layer, mjtByte flg_vert,

test/engine/engine_ray_test.cc

@@ -510,31 +510,14 @@ TEST_F(RayTest, GeomNormal) {
         continue;
       }
 
-      // geom info
-      const mjtNum* pos = d->geom_xpos;
-      const mjtNum* mat = d->geom_xmat;
-      const mjtNum* size = m->geom_size;
-      int type = m->geom_type[0];
-
       // site info
       const mjtNum* pnt = d->site_xpos;
       const mjtNum vec[3] = {d->site_xmat[2], d->site_xmat[5], d->site_xmat[8]};
 
       // compute ray length and normal
       mjtNum r, normal[3];
-      switch (type) {
+      int geomid;
-        case mjGEOM_MESH:
+      r = mj_rayNormal(m, d, pnt, vec, NULL, 1, -1, &geomid, normal);
-          r = mj_rayMeshNormal(m, d, 0, pnt, vec, normal);
-          break;
-        case mjGEOM_SDF:
-          r = mj_raySdfNormal(m, d, 0, pnt, vec, normal);
-          break;
-        case mjGEOM_HFIELD:
-          r = mj_rayHfieldNormal(m, d, 0, pnt, vec, normal);
-          break;
-        default:
-          r = mju_rayGeomNormal(pos, mat, size, pnt, vec, type, normal);
-      }
 
       // compare with sensor
       EXPECT_EQ(r, d->sensordata[0]) << path << ", time " << d->time;
@@ -559,19 +542,7 @@ TEST_F(RayTest, GeomNormal) {
                            d->site_xmat[6 + i]};
         mjtNum dr, dpnt[3];
         mju_addScl3(dpnt, pnt, nudge, eps);
-        switch (type) {
+        dr = mj_rayNormal(m, d, dpnt, vec, NULL, 1, -1, nullptr, nullptr);
-          case mjGEOM_MESH:
-            dr = mj_rayMeshNormal(m, d, 0, dpnt, vec, nullptr);
-            break;
-          case mjGEOM_SDF:
-            dr = mj_raySdfNormal(m, d, 0, dpnt, vec, nullptr);
-            break;
-          case mjGEOM_HFIELD:
-            dr = mj_rayHfieldNormal(m, d, 0, dpnt, vec, nullptr);
-            break;
-          default:
-            dr = mju_rayGeomNormal(pos, mat, size, dpnt, vec, type, nullptr);
-        }
         mju_addScl3(ds[i], dpnt, vec, dr);
       }