mujoco

// Copyright 2026 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef MUJOCO_SRC_EXPERIMENTAL_PLATFORM_ENUM_UTILS_H_
#define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_ENUM_UTILS_H_

#include <array>
#include <bit>
#include <cstdio>
#include <string_view>
#include <utility>

// Provides compile-time utilities for C++ enums, inspired by magic_enum.
// https://github.com/Neargye/magic_enum.
//
// It provides two useful functions:
// - enum_to_string: converts an enum to a string.
// - entries_v: returns an array of pairs of enums and their names.
//
// It relies on non-standard language extensions, so use should be limited to
// non-critical debugging code only. Specifically, it makes assumptions about
// the format of __PRETTY_FUNCTION__ in order to extract the name of an enum
// from a template function.

namespace mujoco::platform::enum_utils {

// Range of values to check for valid enums. This range should be large enough
// to accommodate MuJoCo enums, but not flags. Any enum value that falls outside
// of this range will be ignored by this library.
constexpr int ENUM_MIN_VALUE = -1;
constexpr int ENUM_MAX_VALUE = 1024;

// Fixed-size string for use in compile-time expressions.
template <int N>
struct fixed_string {
  constexpr fixed_string(std::string_view sv) noexcept {
    for (int n = 0; n < N; ++n) str[n] = sv[n];
  }
  constexpr operator std::string_view() const noexcept {
    return {str.data(), N};
  }
  std::array<char, N + 1> str{};
};

// Characters that are valid parts of an enum name.
constexpr auto is_valid_char(char ch) noexcept {
  return (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch == '_') ||
         (ch >= '0' && ch <= '9');
}

// Extracts the name of an enum from a string generated by __PRETTY_FUNCTION__.
constexpr auto pretty_name(std::string_view sv) noexcept {
  for (int n = sv.size() - 1; n > 0; --n) {
    if (!is_valid_char(sv[n])) {
      sv.remove_prefix(n + 1);
      break;
    }
  }
  if (sv[0] >= '0' && sv[0] <= '9') {
    return std::string_view();
  }
  return sv;
}

// Returns the name of the V which is an enum value of type E.
template <typename E, E V>
constexpr auto n() noexcept {
#if defined(__GNUC__) || defined(__clang__)
  return pretty_name({__PRETTY_FUNCTION__, sizeof(__PRETTY_FUNCTION__) - 2});
#elif defined(_MSC_VER)
  return pretty_name({__FUNCSIG__, sizeof(__FUNCSIG__) - 17});
#endif
}

// Returns true if the value V is a valid enum value of type E.
template <typename E, int V>
constexpr auto is_valid() {
  return !n<E, std::bit_cast<E>(V)>().empty();
}

// Returns the value of the Nth element of enum E.
template <typename E>
constexpr auto nth(int v) {
  return std::bit_cast<E>(ENUM_MIN_VALUE + v);
}

// Returns the number of bits set to true in the given array.
template <int N>
constexpr auto count_values(const bool (&valid)[N]) {
  int count = 0;
  for (int n = 0; n < N; ++n)
    if (valid[n]) ++count;
  return count;
}

// Returns an array of the valid values of enum E. This is a subset of the
// values in the range defined by the index sequence. If an enum has a value
// outside this range, it will be ignored.
template <typename E, std::size_t... I>
constexpr auto values(std::index_sequence<I...>) noexcept {
  constexpr bool valid[sizeof...(I)] = {is_valid<E, nth<E>(I)>()...};
  constexpr auto num_valid = count_values(valid);
  static_assert(num_valid > 0, "no support for empty enums");

  std::array<E, num_valid> values = {};
  for (int offset = 0, n = 0; n < num_valid; ++offset) {
    if (valid[offset]) {
      values[n] = nth<E>(offset);
      ++n;
    }
  }
  return values;
}

// Returns an array of valid values of enum E. The possible range of values is
// defined by the range [ENUM_MIN_VALUE, ENUM_MAX_VALUE]. If the enum has a
// value outside this range, it will be ignored.
template <typename E>
constexpr auto values() noexcept {
  constexpr int enum_size = ENUM_MAX_VALUE - ENUM_MIN_VALUE + 1;
  return values<E>(std::make_index_sequence<enum_size>({}));
}

template <typename E>
inline constexpr auto values_v = values<E>();

// Returns the name of V which is an enum value of type E.
template <typename E, E V>
constexpr auto enum_name() {
  constexpr auto name = n<E, V>();
  return fixed_string<name.size()>(name);
}

template <typename E, E V>
inline constexpr auto enum_name_v = enum_name<E, V>();

// Returns an array of pairs of enum values and their names.
template <typename E, std::size_t... I>
constexpr auto entries(std::index_sequence<I...>) noexcept {
  return std::array<std::pair<E, std::string_view>, sizeof...(I)>{
      {{values_v<E>[I], enum_name_v<E, values_v<E>[I]>}...}};
}

template <typename E>
inline constexpr auto entries_v =
    entries<E>(std::make_index_sequence<values_v<E>.size()>());

// Returns the name of the given enum value.
template <typename E>
constexpr std::string_view enum_to_string(E value) {
  for (const auto& [key, name] : entries_v<E>) {
    if (value == key) return name;
  }
  return {};
}

}  // namespace mujoco::platform::enum_utils

#endif  // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_ENUM_UTILS_H_