[mjspec] mjs_findAll() in C user_api

[Tadinu]

Currently, MjSpec python api supports handy properties, eg geoms, bodies, produced from FindAllImpl() in pybind11 specs.cc. In C api, in case of such a need for a recursive property set to contype, conaffinity of all geoms of a body OR eg gravcomp of all bodies under a mjSpec, mjs_findAll() would be convenient to be used for looping over child specs.

An implementation example for parent type as either of mjSpec, mjsBody can be, as largely based on FindAllImpl()

template <typename TSpec, typename TChildSpec,
                 typename = std::enable_if<std::is_same_v<TSpec, mjSpec> | std::is_same_v<TSpec, mjsBody>>>
inline std::vector<TChildSpec*> mjs_findAll(TSpec* base_spec, bool recursive = true) {
  constexpr mjtObj objtype = std::is_same_v<TChildSpec, mjsGeom>
                               ? mjOBJ_GEOM
                               : std::is_same_v<TChildSpec, mjsBody>
                               ? mjOBJ_BODY
                               : std::is_same_v<TChildSpec, mjsCamera>
                               ? mjOBJ_CAMERA
                               : std::is_same_v<TChildSpec, mjsFrame>
                               ? mjOBJ_FRAME
                               : std::is_same_v<TChildSpec, mjsGeom>
                               ? mjOBJ_GEOM
                               : std::is_same_v<TChildSpec, mjsJoint>
                               ? mjOBJ_JOINT
                               : std::is_same_v<TChildSpec, mjsLight>
                               ? mjOBJ_LIGHT
                               : std::is_same_v<TChildSpec, mjsSite>
                               ? mjOBJ_SITE
                               : mjOBJ_UNKNOWN;

  if (objtype == mjOBJ_UNKNOWN) {
    // this should never happen
    throw mjCError(
        "[mjs_findAll] supports the types: body, frame, geom, site, "
        "joint, light, camera.");
  }

  std::vector<TChildSpec*> child_specs;
  mjsElement* el = nullptr;
  if constexpr (std::is_same_v<TSpec, mjSpec>) {
    el = mjs_firstElement(base_spec, objtype);
  } else {
    el = mjs_firstChild(base_spec, objtype, recursive);
  }

  const std::string error = mjs_getError(mjs_getSpec(base_spec->element));
  if (!el && !error.empty()) {
    throw mjCError(error);
  }
  while (el) {
    if constexpr (std::is_same_v<TChildSpec, mjsElement>) {
      child_specs.push_back(el);
    }
    else {
      TChildSpec* child_spec = nullptr;
      if constexpr (type == mjOBJ_BODY) {
        child_spec = mjs_asBody(el);
      } else if constexpr (type == mjOBJ_CAMERA) {
        child_spec = mjs_asCamera(el);
      } else if constexpr (type == mjOBJ_FRAME) {
        child_spec = mjs_asFrame(el);
      } else if constexpr (type == mjOBJ_GEOM) {
        child_spec = mjs_asGeom(el);
      } else if constexpr (type == mjOBJ_JOINT) {
        child_spec = mjs_asJoint(el);
      } else if constexpr (type == mjOBJ_LIGHT) {
        child_spec = mjs_asLight(el);
      } else if constexpr (type == mjOBJ_SITE) {
        child_spec = mjs_asSite(el);
      }

      if (child_spec) {
        child_specs.push_back(child_spec);
      }
    }

    if constexpr (std::is_same_v<TSpec, mjSpec>) {
      el = mjs_nextElement(base_spec, el);
    } else {
      el = mjs_nextChild(base_spec, el, recursive);
    }
  }
  return child_specs;
}

Example usecases:

void SetBodyTreeCollisionEnabled(mjsBody* base_body_spec, bool enabled) {
  for (const auto& geom_spec : mjs_findAll<mjsGeom, mjsBody>(base_body_spec)) {
    geom_spec->contype = enabled;
    geom_spec->conaffinity = enabled;
  }
  for (const auto& child_body_spec : mjs_findAll<mjsBody, mjsBody>(base_body_spec)) {
    SetBodyTreeCollisionEnabled(child_body_spec, enabled);
  }
}

void SetBodyTreeGravityCompensationEnabled(mjsBody* base_body_spec, bool enabled) {
  base_body_spec->gravcomp = enabled;
  for (const auto& child_body_spec : FindAllChildSpecs<mjsBody, mjsBody>(base_body_spec)) {
    SetBodyTreeGravityCompensationEnabled(child_body_spec, enabled);
  }
}

[yuvaltassa]

cc @quagla

[quagla]

Would you like to send a PR adding that function? Then we can probably use it in the Python API directly, instead of FindAllImpl.

[Tadinu]

@quagla I've created PR #2561. Thank you!