(tvm-ffi) branch main updated: [FEAT] Add a name membership query to EnumDef (#849)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git


The following commit(s) were added to refs/heads/main by this push:
     new db59370e [FEAT] Add a name membership query to EnumDef (#849)
db59370e is described below

commit db59370efb9a55a74e3549d70268f979a1ea1eeb
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Wed Oct 14 18:17:01 2026 -0400

    [FEAT] Add a name membership query to EnumDef (#849)
---
 include/tvm/ffi/reflection/enum_def.h | 20 ++++++++++++++++++++
 tests/cpp/test_reflection.cc          | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/include/tvm/ffi/reflection/enum_def.h 
b/include/tvm/ffi/reflection/enum_def.h
index f444014a..dd0317df 100644
--- a/include/tvm/ffi/reflection/enum_def.h
+++ b/include/tvm/ffi/reflection/enum_def.h
@@ -48,6 +48,26 @@ namespace reflection {
 template <typename EnumClsObj, typename = 
std::enable_if_t<std::is_base_of_v<EnumObj, EnumClsObj>>>
 class EnumDef : public ReflectionDefBase {
  public:
+  /*!
+   * \brief Check whether a canonical name is defined without registering enum 
state or entries.
+   * \param name The canonical string index.
+   * \return Whether the name is defined for this enum type.
+   */
+  static bool IsDefined(const String& name) {
+    constexpr TVMFFIByteArray column_name = 
AsByteArray(type_attr::kEnumState.name);
+    if (TVMFFIGetTypeAttrColumn(&column_name) == nullptr) return false;
+    AnyView value =
+        
TypeAttrColumn(type_attr::kEnumState)[EnumClsObj::_GetOrAllocRuntimeTypeIndex()];
+    if (value == nullptr) return false;
+    if (value.type_index() != EnumStateObj::_GetOrAllocRuntimeTypeIndex()) {
+      TVM_FFI_THROW(TypeError) << "Expected `" << EnumStateObj::_type_key
+                               << "` in enum state column, but got `" << 
value.GetTypeKey() << "`";
+    }
+    // Avoid cast<EnumState> because its inline type index may not be 
initialized yet.
+    const auto* state = static_cast<const 
EnumStateObj*>(value.cast<ObjectRef>().get());
+    return state->indexes.count(name) != 0;
+  }
+
   /*!
    * \brief Register an enum singleton with an automatically assigned integer 
index.
    *
diff --git a/tests/cpp/test_reflection.cc b/tests/cpp/test_reflection.cc
index 6fd600c9..f507b48a 100644
--- a/tests/cpp/test_reflection.cc
+++ b/tests/cpp/test_reflection.cc
@@ -71,6 +71,11 @@ struct PrefixLookupObj : public Object {
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.PrefixLookup", PrefixLookupObj, 
Object);
 };
 
+class QueryEnumObj : public EnumObj {
+ public:
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.QueryEnum", QueryEnumObj, EnumObj);
+};
+
 TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
 
@@ -117,6 +122,35 @@ TEST(Reflection, GetFieldByteOffset) {
   EXPECT_EQ(reflection::GetFieldByteOffsetToObject(&TIntObj::value), 
sizeof(TVMFFIObject));
 }
 
+TEST(Reflection, EnumIsDefined) {
+  using Definition = reflection::EnumDef<QueryEnumObj>;
+  reflection::TypeAttrColumn states(type_attr::kEnumState);
+  int32_t type_index = QueryEnumObj::_GetOrAllocRuntimeTypeIndex();
+  EXPECT_FALSE(Definition::IsDefined("Alpha"));
+  EXPECT_EQ(states[type_index], nullptr);
+
+  Definition("Alpha").set_attr("description", String("first"));
+  Enum original = EnumObj::_GetByStrIndex<QueryEnumObj>("Alpha");
+  EnumState state = states[type_index].cast<EnumState>();
+  EXPECT_TRUE(Definition::IsDefined("Alpha"));
+  EXPECT_FALSE(Definition::IsDefined("alpha"));
+  EXPECT_FALSE(Definition::IsDefined("Beta"));
+  EXPECT_FALSE(Definition::IsDefined("0"));
+  EXPECT_FALSE(Definition::IsDefined(String("Alpha\0suffix", 12)));
+  EXPECT_TRUE(states[type_index].cast<EnumState>().same_as(state));
+  EXPECT_EQ(state->entries.size(), 1);
+  EXPECT_EQ(state->indexes.size(), 2);
+  EXPECT_EQ(state->attrs.size(), 1);
+  EXPECT_EQ(state->attrs.at("description").size(), 1);
+  EXPECT_EQ(state->attrs.at("description").at(original).cast<String>(), 
"first");
+  
EXPECT_TRUE(EnumObj::_GetByStrIndex<QueryEnumObj>("Alpha").same_as(original));
+  EXPECT_TRUE(EnumObj::_GetByIntIndex<QueryEnumObj>(0).same_as(original));
+  Definition("Beta");  // NOLINT(bugprone-unused-raii)
+  EXPECT_TRUE(Definition::IsDefined("Beta"));
+  EXPECT_TRUE(EnumObj::_GetByStrIndex<QueryEnumObj>("Beta").same_as(
+      EnumObj::_GetByIntIndex<QueryEnumObj>(1)));
+}
+
 TEST(Reflection, FieldGetter) {
   ObjectRef a = TInt(10);
   reflection::FieldGetter getter("test.Int", "value");