Add feature to find Prim by prim_id

Add unit test for handle-allocator.

examples/api_tutorial/api-tutorial-main.cc

@@ -9,6 +9,7 @@
 
 // Tydra is a collection of APIs to access/convert USD Prim data
 // (e.g. Can get Attribute by name)
+// See <tinyusdz>/examples/tydra_api for more Tydra API examples.
 #include "tydra/scene-access.hh"
 
 //
@@ -225,7 +226,7 @@ void CreateScene(tinyusdz::Stage *stage) {
         uvIndices.push_back(2);
 
         uvPrimvar.set_indices(uvIndices);
-      
+
         // primvar name is extracted from Primvar::name
         std::string err;
         if (!mesh.set_primvar(uvPrimvar, &err)) {
@@ -235,46 +236,6 @@ void CreateScene(tinyusdz::Stage *stage) {
     }
 
 
-    // Access GeomPrimvar
-    {
-      std::cout << "uv is primvar? " << mesh.has_primvar("uv") << "\n";
-      tinyusdz::GeomPrimvar primvar;
-      std::string err;
-      if (mesh.get_primvar("uv", &primvar, &err)) {
-        std::cout << "uv primvar is Indexed Primvar? " << primvar.has_indices() << "\n";
-      } else {
-        std::cerr << "get_primvar(\"uv\") failed. err = " << err << "\n";
-      }
-
-      // Equivalent to pxr::UsdGeomPrimvar::ComputeFlattened().
-      // elems[i] = values[indices[i]]
-      tinyusdz::value::Value value;
-      if (primvar.flatten_with_indices(&value, &err)) {
-        // value;:Value can contain any types, but value.array_size() should work well only for primvar types(e.g. `float[]`, `color3f[]`)
-        // It would report 0 for non-primvar types(e.g.`std::vector<Xform>`)
-        std::cout << "uv primvars. array size = " << value.array_size() << "\n";
-        std::cout << "uv primvars. expand_by_indices result = " << tinyusdz::value::pprint_value(value) << "\n";
-      } else {
-        std::cerr << "expand_by_indices failed. err = " << err << "\n";
-      }
-
-      // Typed version
-      std::vector<tinyusdz::value::texcoord2f> uvs;
-      if (primvar.flatten_with_indices(&uvs, &err)) {
-        // value;:Value can contain any types, but value.array_size() should work well only for primvar types(e.g. `float[]`, `color3f[]`)
-        // It would report 0 for non-primvar types(e.g.`std::vector<Xform>`)
-        std::cout << "uv primvars. array size = " << uvs.size() << "\n";
-        std::cout << "uv primvars. expand_by_indices result = " << tinyusdz::value::pprint_value(uvs) << "\n";
-      } else {
-        std::cerr << "expand_by_indices failed. err = " << err << "\n";
-      }
-
-      std::vector<tinyusdz::GeomPrimvar> gpvars = mesh.get_primvars();
-      std::cout << "# of primvars = " << gpvars.size();
-      for (const auto &item : gpvars) {
-        std::cout << "  primvar = " << item.name() << "\n";
-      }
-    }
   }
 
   tinyusdz::GeomSphere sphere;
@@ -294,7 +255,7 @@ void CreateScene(tinyusdz::Stage *stage) {
   //  +- [Mesh]
   //  +- [Sphere]
   //
-  
+
 
 
   // Prim's elementName is read from concrete Prim class(GeomMesh::name,
@@ -342,7 +303,7 @@ void CreateScene(tinyusdz::Stage *stage) {
 
 
     // You can also use SetCustomDataByKey to set custom value with key having namespaces(':')
-    
+
     tinyusdz::MetaVariable intval = int(5);
     tinyusdz::SetCustomDataByKey("mydict:myval", intval, /* inout */customData);
 
@@ -392,11 +353,23 @@ int main(int argc, char **argv) {
 
   }
 
-  // GeomPrimvar
+  // find Prim by prim_id
   {
+    uint64_t prim_id = 2;
+
+    const tinyusdz::Prim *prim{nullptr};
+    std::string err;
+    bool ret = stage.find_prim_by_prim_id(prim_id, prim, &err);
+    if (ret && prim) {
+      std::cout << "Found Prim by ID: " << prim_id << "\n";
+      std::cout << "Prim's absolute_path: " << tinyusdz::to_string(prim->absolute_path()) << "\n";
+    } else {
+      std::cerr << err << "\n";
+    }
 
   }
 
+  // GetAttribute and GeomPrimvar
   {
     tinyusdz::Path path(/* absolute prim path */ "/root/quad",
                         /* property path */ "");
@@ -437,6 +410,55 @@ int main(int argc, char **argv) {
     } else {
       std::cerr << err << "\n";
     }
+
+
+    const tinyusdz::GeomMesh *mesh = prim->as<tinyusdz::GeomMesh>();
+    if (!mesh) {
+      std::cerr << "Expected GeomMesh.\n";
+      return -1;
+    }
+
+    // GeomPrimvar
+    // Access GeomPrimvar
+    {
+      std::cout << "uv is primvar? " << mesh->has_primvar("uv") << "\n";
+      tinyusdz::GeomPrimvar primvar;
+      std::string err;
+      if (mesh->get_primvar("uv", &primvar, &err)) {
+        std::cout << "uv primvar is Indexed Primvar? " << primvar.has_indices() << "\n";
+      } else {
+        std::cerr << "get_primvar(\"uv\") failed. err = " << err << "\n";
+      }
+
+      // Equivalent to pxr::UsdGeomPrimvar::ComputeFlattened().
+      // elems[i] = values[indices[i]]
+      tinyusdz::value::Value value;
+      if (primvar.flatten_with_indices(&value, &err)) {
+        // value;:Value can contain any types, but value.array_size() should work well only for primvar types(e.g. `float[]`, `color3f[]`)
+        // It would report 0 for non-primvar types(e.g.`std::vector<Xform>`)
+        std::cout << "uv primvars. array size = " << value.array_size() << "\n";
+        std::cout << "uv primvars. expand_by_indices result = " << tinyusdz::value::pprint_value(value) << "\n";
+      } else {
+        std::cerr << "expand_by_indices failed. err = " << err << "\n";
+      }
+
+      // Typed version
+      std::vector<tinyusdz::value::texcoord2f> uvs;
+      if (primvar.flatten_with_indices(&uvs, &err)) {
+        // value;:Value can contain any types, but value.array_size() should work well only for primvar types(e.g. `float[]`, `color3f[]`)
+        // It would report 0 for non-primvar types(e.g.`std::vector<Xform>`)
+        std::cout << "uv primvars. array size = " << uvs.size() << "\n";
+        std::cout << "uv primvars. expand_by_indices result = " << tinyusdz::value::pprint_value(uvs) << "\n";
+      } else {
+        std::cerr << "expand_by_indices failed. err = " << err << "\n";
+      }
+
+      std::vector<tinyusdz::GeomPrimvar> gpvars = mesh->get_primvars();
+      std::cout << "# of primvars = " << gpvars.size();
+      for (const auto &item : gpvars) {
+        std::cout << "  primvar = " << item.name() << "\n";
+      }
+    }
   }
 
 

src/handle-allocator.hh

@@ -3,6 +3,8 @@
 #pragma once
 
 #include <cstdint>
+#include <algorithm>
+//#include <iostream>
 //#include <cassert>
 #include <limits>
 #include <vector>
@@ -31,9 +33,11 @@ public:
     T handle = 0;
 
     if (!freeList_.empty()) {
-      // Reuse previously issued handle.
+      // Reuse last element.
       handle = freeList_.back();
       freeList_.pop_back();
+      // Delay sort until required
+      dirty_ = true;
       (*dst) = handle;
       return true;
     }
@@ -41,6 +45,7 @@ public:
     handle = counter_;
     if ((handle >= static_cast<T>(1)) && (handle < std::numeric_limits<T>::max())) {
       counter_++;
+      //std::cout << "conter = " << counter_ << "\n";
       (*dst) = handle;
       return true;
     }
@@ -53,10 +58,14 @@ public:
     if (handle == counter_ - static_cast<T>(1)) {
       if (counter_ > static_cast<T>(1)) {
         counter_--;
+      } else {
+        return false;
       }
     } else {
       if (handle >= static_cast<T>(1)) {
         freeList_.push_back(handle);
+        // Delay sort until required
+        dirty_ = true;
       } else {
         // invalid handle
         return false;
@@ -66,9 +75,38 @@ public:
     return true;
   }
 
+  bool Has(const T handle) const {
+    if (dirty_) {
+      std::sort(freeList_.begin(), freeList_.end());
+      dirty_ = false;
+    }
+
+    if (handle < 1) {
+      return false;
+    }
+
+    // Do binary search.
+    if (std::binary_search(freeList_.begin(), freeList_.end(), handle)) {
+      return false;
+    }
+
+    if (handle >= counter_) {
+      return false;
+    }
+
+    return true;
+  }
+
+  int64_t Size() const {
+    return counter_ - freeList_.size() - 1;
+  }
+
 private:
-  std::vector<T> freeList_;
+  // TODO: Use unorderd_set or unorderd_map for efficiency?
+  // worst case complexity is still c.size() though.
+  mutable std::vector<T> freeList_; // will be sorted in `Has` call.
   T counter_{};
+  mutable bool dirty_{true};
 };
 
 } // namespace tinyusdz

src/stage.cc

@@ -242,6 +242,73 @@ bool Stage::find_prim_at_path(const Path &path, const Prim *&prim,
   }
 }
 
+namespace {
+
+bool FindPrimByPrimIdRec(uint64_t prim_id, const Prim *root, const Prim **primFound, int level, std::string *err) {
+
+  if (level > 1024*1024*128) {
+    // too deep node.
+    return false;
+  }
+
+  if (!primFound) {
+    return false;
+  }
+
+  if (root->prim_id() == int64_t(prim_id)) {
+    (*primFound) = root;
+    return true;
+  }
+
+  // Brute-force search.
+  for (const auto &child : root->children()) {
+    if (FindPrimByPrimIdRec(prim_id, &child, primFound, level+1, err)) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+} // namespace local
+
+bool Stage::find_prim_by_prim_id(const uint64_t prim_id, const Prim *&prim,
+                              std::string *err) const {
+  if (prim_id < 1) {
+    if (err) {
+      (*err) = "Input prim_id must be 1 or greater.";
+    }
+    return false;
+  }
+
+  if (_prim_id_dirty) {
+    DCOUT("clear prim_id cache.");
+    // Clear cache.
+    _prim_id_cache.clear();
+
+    _prim_id_dirty = false;
+  } else {
+    // First find from a cache.
+    auto ret = _prim_id_cache.find(prim_id);
+    if (ret != _prim_id_cache.end()) {
+      DCOUT("Found cache.");
+      return ret->second;
+    }
+  }
+
+  const Prim *p{nullptr};
+  for (const auto &root : root_prims()) {
+    if (FindPrimByPrimIdRec(prim_id, &root, &p, 0, err)) {
+      _prim_id_cache[prim_id] = p;
+      prim = p;
+      return true;
+    }
+  }
+
+  return false;
+
+}
+
 nonstd::expected<const Prim *, std::string> Stage::GetPrimFromRelativePath(
     const Prim &root, const Path &path) const {
   // TODO: Resolve "../"
@@ -476,7 +543,7 @@ bool Stage::allocate_prim_id(uint64_t *prim_id) const {
 
   uint64_t val;
   if (_prim_id_allocator.Allocate(&val)) {
-    (*prim_id) = val; 
+    (*prim_id) = val;
     return true;
   }
 
@@ -487,9 +554,13 @@ bool Stage::release_prim_id(const uint64_t prim_id) const {
   return _prim_id_allocator.Release(prim_id);
 }
 
+bool Stage::has_prim_id(const uint64_t prim_id) const {
+  return _prim_id_allocator.Has(prim_id);
+}
+
 namespace {
 
-bool ComputeAbsPathAndAssignPrimIdRec(const Stage &stage, Prim &prim, const Path &parentPath, uint32_t depth) {
+bool ComputeAbsPathAndAssignPrimIdRec(const Stage &stage, Prim &prim, const Path &parentPath, uint32_t depth, bool assign_prim_id, bool force_assign_prim_id = true) {
   if (depth > 1024*1024*128) {
     // too deep node.
     return false;
@@ -500,16 +571,18 @@ bool ComputeAbsPathAndAssignPrimIdRec(const Stage &stage, Prim &prim, const Path
   Path abs_path = parentPath.AppendPrim(prim.element_name());
 
   prim.absolute_path() = abs_path;
-  if (prim.prim_id() < 1) {
-    uint64_t prim_id{0};
-    if (!stage.allocate_prim_id(&prim_id)) {
-      return false;
+  if (assign_prim_id) {
+    if (force_assign_prim_id || (prim.prim_id() < 1)) {
+      uint64_t prim_id{0};
+      if (!stage.allocate_prim_id(&prim_id)) {
+        return false;
+      }
+      prim.prim_id() = int64_t(prim_id);
     }
-    prim.prim_id() = int64_t(prim_id);
   }
 
   for (Prim &child : prim.children()) {
-    if (!ComputeAbsPathAndAssignPrimIdRec(stage, child, abs_path, depth+1)) {
+    if (!ComputeAbsPathAndAssignPrimIdRec(stage, child, abs_path, depth+1, assign_prim_id, force_assign_prim_id)) {
       return false;
     }
   }
@@ -519,16 +592,32 @@ bool ComputeAbsPathAndAssignPrimIdRec(const Stage &stage, Prim &prim, const Path
 
 } // namespace local
 
-bool Stage::compute_absolute_prim_path_and_assign_prim_id() {
+bool Stage::compute_absolute_prim_path_and_assign_prim_id(bool force_assign_prim_id) {
 
   Path rootPath("/", "");
   for (Prim &root : root_prims()) {
-    if (!ComputeAbsPathAndAssignPrimIdRec(*this, root, rootPath, 1)) {
+    if (!ComputeAbsPathAndAssignPrimIdRec(*this, root, rootPath, 1, /* assign_prim_id */true, force_assign_prim_id)) {
+      return false;
+    }
+  }
+
+  // TODO: Only set dirty when prim_id changed.
+  _prim_id_dirty = true;
+
+  return true;
+}
+
+bool Stage::compute_absolute_prim_path() {
+
+  Path rootPath("/", "");
+  for (Prim &root : root_prims()) {
+    if (!ComputeAbsPathAndAssignPrimIdRec(*this, root, rootPath, 1, /* assign prim_id */false)) {
       return false;
     }
   }
 
   return true;
+
 }
 
 }  // namespace tinyusdz

src/stage.hh

@@ -102,12 +102,22 @@ class Stage {
   ///
   /// @param[in] root Find from this Prim
   /// @param[in] relative_path relative path(e.g. `dora/muda`)
-  /// @param[out] prim const reference to Prim(if found)
+  /// @param[out] prim const reference to the pointer to Prim(if found)
   /// @param[out] err Error message(filled when false is returned)
   ///
   /// @returns true if found a Prim.
   bool find_prim_from_relative_path(const Prim &root, const Path &relative_path, const Prim *&prim, std::string *err) const;
 
+  ///
+  /// Find(Get) Prim from Prim ID. Prim with no Prim ID assigned(-1 or 0) are ignored.
+  ///
+  /// @param[in] prim_id Prim ID(1 or greater)
+  /// @param[out] prim const reference to the pointer to Prim(if found)
+  /// @param[out] err Error message(filled when false is returned)
+  ///
+  /// @returns true if found a Prim.
+  bool find_prim_by_prim_id(const uint64_t prim_id, const Prim *&prim, std::string *err = nullptr) const;
+
   const std::vector<Prim> &root_prims() const {
     return root_nodes;
   }
@@ -128,18 +138,26 @@ class Stage {
   /// Assign unique Prim id inside this Stage.
   ///
   bool allocate_prim_id(uint64_t *prim_id) const;
+
   bool release_prim_id(const uint64_t prim_id) const;
+  bool has_prim_id(const uint64_t prim_id) const;
   
   ///
   /// Call this function after you finished adding Prims manually to Stage.
   /// (No need to call this if you just use ether USDA/USDC/USDZ reader).
   ///
   /// - Compute absolute path and set it to Prim::abs_path for each Prim currently added to this Stage.
-  /// - Assign unique ID to Prim(if Prim does not have prim_id)
+  /// - Assign unique ID to Prim
   ///
+  /// @param[in] force_assign_prim_id true Overwrite `prim_id` of each Prim. false only assign Prim id when `prim_id` is -1(preserve user-assgiend prim_id). Setting `false` is not recommended since prim_id may not be unique over Prims in Stage.
   /// @return false when the Stage contains any invalid Prim 
   ///
-  bool compute_absolute_prim_path_and_assign_prim_id();
+  bool compute_absolute_prim_path_and_assign_prim_id(bool force_assign_prim_id=true);
+
+  ///
+  /// Compute absolute Prim path only.
+  ///
+  bool compute_absolute_prim_path();
 
   ///
   /// Compose scene(Not implemented yet).
@@ -182,7 +200,13 @@ class Stage {
   // key : prim_part string (e.g. "/path/bora")
   mutable std::map<std::string, const Prim *> _prim_path_cache;
 
-  mutable bool _dirty{false}; // True when Stage content changes(addition, deletion, composition/flatten, etc.)
+  // Cached prim_id -> Prim lookup
+  // key : prim_id
+  mutable std::map<uint64_t, const Prim *> _prim_id_cache;
+
+  mutable bool _dirty{true}; // True when Stage content changes(addition, deletion, composition/flatten, etc.)
+
+  mutable bool _prim_id_dirty{true}; // True when Prim Id assignent changed(TODO: Unify with `_dirty` flag)
 
   mutable HandleAllocator<uint64_t> _prim_id_allocator;
 

tests/unit/CMakeLists.txt

@@ -4,6 +4,7 @@ set(TEST_SOURCES
 	unit-main.cc
 
 	unit-customdata.cc
+	unit-handle-allocator.cc
 	unit-prim-types.cc
 	unit-primvar.cc
 	unit-value-types.cc

tests/unit/unit-common.hh

@@ -4,6 +4,7 @@
 
 #include <iostream>
 #include <limits>
+#include <cmath>
 
 namespace tinyusdz_test {
 

tests/unit/unit-handle-allocator.cc

@@ -0,0 +1,102 @@
+#ifdef _MSC_VER
+#define NOMINMAX
+#endif
+
+#define TEST_NO_MAIN
+#include "acutest.h"
+
+#include "handle-allocator.hh"
+#include "unit-common.hh"
+
+#include <random>
+#include <functional>
+
+
+using namespace tinyusdz;
+using namespace tinyusdz_test;
+
+void handle_allocator_test(void) {
+
+  std::random_device seed_gen;
+  std::mt19937 engine(seed_gen());
+
+  uint64_t n = 1024*16;
+
+  std::vector<uint64_t> perm_handles;
+
+  HandleAllocator<uint64_t> allocator;
+
+  bool ok = true;
+  for (size_t i = 0; i < n; i++) {
+    uint64_t handle;
+    if (!allocator.Allocate(&handle)) {
+      ok = false;
+      break; 
+    }
+
+    perm_handles.push_back(handle);
+  }
+  TEST_CHECK(ok);
+  TEST_CHECK(allocator.Size() == n);
+
+  std::shuffle(perm_handles.begin(), perm_handles.end(), engine);
+
+  // handle should exist
+  ok = true;
+  for (size_t i = 0; i < perm_handles.size(); i++) {
+    //std::cout << "handle = " << perm_handles[i] << ", Size = " << allocator.Size() << "\n";
+    if (!allocator.Has(perm_handles[i])) {
+      ok = false;
+      break;
+    }
+  }
+  TEST_CHECK(ok);
+
+  //std::cout << "del \n";
+  // delete
+  ok = true;
+  for (size_t i = 0; i < perm_handles.size(); i++) {
+    //std::cout << "del handle " << perm_handles[i] << ", Size = " << allocator.Size() << "\n";
+    if (!allocator.Release(perm_handles[i])) {
+      ok = false;
+      break;
+    }
+  }
+  TEST_CHECK(ok);
+  //std::cout << "del done\n";
+  TEST_CHECK(allocator.Size() == 0);
+  
+  // no handle exists
+  ok = true;
+  for (size_t i = 0; i < perm_handles.size(); i++) {
+    //std::cout << "has handle " << perm_handles[i] << " = " << allocator.Has(perm_handles[i]) << "\n";
+    if (allocator.Has(perm_handles[i])) {
+      ok = false;
+      break;
+    }
+  }
+  TEST_CHECK(ok);
+
+  // reallocate test
+  ok = true;
+  std::vector<uint64_t> handles;
+  for (size_t i = 0; i < perm_handles.size(); i++) {
+    uint64_t handle;
+    if (!allocator.Allocate(&handle)) {
+      ok = false;
+      break;
+    }
+    handles.push_back(handle);
+  }
+  TEST_CHECK(ok);
+  // uniqueness check
+  std::sort(handles.begin(), handles.end());
+  // array items = [1, ..., n]
+  TEST_CHECK(handles.front() == 1); 
+  uint64_t &last = handles.back();
+  TEST_CHECK(last == n);
+  handles.erase(std::unique(handles.begin(), handles.end()), handles.end());
+  TEST_CHECK(handles.size() == perm_handles.size());
+
+
+}

tests/unit/unit-handle-allocator.h

@@ -0,0 +1,3 @@
+#pragma once
+
+void handle_allocator_test(void);

tests/unit/unit-main.cc

@@ -9,6 +9,7 @@
 #include "unit-value-types.h"
 #include "unit-xform.h"
 #include "unit-customdata.h"
+#include "unit-handle-allocator.h"
 
 #if defined(TINYUSDZ_WITH_PXR_COMPAT_API)
 #include "unit-pxr-compat-api.h"
@@ -22,6 +23,7 @@ TEST_LIST = {
   { "value_types_test", value_types_test },
   { "xformOp_test", xformOp_test },
   { "customdata_test", customdata_test },
+  { "handle_allocator_test", handle_allocator_test },
 #if defined(TINYUSDZ_WITH_PXR_COMPAT_API)
   { "pxr_compat_api_test", pxr_compat_api_test },
 #endif