[API] Add type index query mode

include/tvm/ffi/c_api.h

@@ -1358,15 +1358,17 @@ TVM_FFI_DLL const TVMFFIByteArray* TVMFFIBacktrace(const char* filename, int lin
  * If the static_tindex is non-negative, the function will
  * allocate a runtime type index.
  * Otherwise, we will populate the type table and return the static index.
+ * If parent_type_index is -2, the function queries the existing type index:
+ * it returns the registered index for type_key, or -2 if type_key is not registered.
  *
  * \param type_key The type key.
  * \param type_depth The type depth.
  * \param static_type_index Static type index if any, can be -1, which means this is a dynamic index
  * \param num_child_slots Number of slots reserved for its children.
  * \param child_slots_can_overflow Whether to allow child to overflow the slots.
- * \param parent_type_index Parent type index, pass in -1 if it is root.
+ * \param parent_type_index Parent type index, pass in -1 if it is root, or -2 to query only.
  *
- * \return The allocated type index.
+ * \return The existing or allocated type index; -2 when query-only mode misses.
  */
 TVM_FFI_DLL int32_t TVMFFITypeGetOrAllocIndex(const TVMFFIByteArray* type_key,
                                               int32_t static_type_index, int32_t type_depth,

src/ffi/object.cc

@@ -126,6 +126,9 @@ class TypeTable {
     if (it != type_key2index_.end()) {
       return type_table_[(*it).second]->type_index;
     }
+    if (parent_type_index == -2) {
+      return -2;
+    }
 
     // get parent's entry
     Entry* parent = [&]() -> Entry* {

tests/cpp/test_object.cc

@@ -86,6 +86,21 @@ TEST(Object, CRTPObjectInfo) {
   EXPECT_GE(info->type_index, TypeIndex::kTVMFFIDynObjectBegin);
 }
 
+TEST(Object, TypeGetOrAllocIndexQueryRegistered) {
+  TVMFFIByteArray type_key{TIntObj::_type_key, std::char_traits<char>::length(TIntObj::_type_key)};
+  EXPECT_EQ(TVMFFITypeGetOrAllocIndex(&type_key, -1, 0, 0, 0, -2), TIntObj::RuntimeTypeIndex());
+}
+
+TEST(Object, TypeGetOrAllocIndexQueryMissDoesNotRegister) {
+  const char* type_key_data = "test.TypeGetOrAllocIndexQueryMiss";
+  TVMFFIByteArray type_key{type_key_data, std::char_traits<char>::length(type_key_data)};
+  EXPECT_EQ(TVMFFITypeGetOrAllocIndex(&type_key, -1, 0, 0, 0, -2), -2);
+
+  int32_t type_index = -1;
+  EXPECT_NE(TVMFFITypeKeyToIndex(&type_key, &type_index), 0);
+  EXPECT_EQ(type_index, -1);
+}
+
 TEST(Object, InstanceCheck) {
   ObjectPtr<Object> a = make_object<TIntObj>(11);
   ObjectPtr<Object> b = make_object<TFloatObj>(11);