USDPrimIterator class now encapsulates a stage.

3 files changed, 48 insertions, 19 deletions

diff --git a/source/blender/io/usd/import/usd_prim_iterator.cc b/source/blender/io/usd/import/usd_prim_iterator.cc
index dd9edbe6d68..6f40538c787 100644
--- a/source/blender/io/usd/import/usd_prim_iterator.cc
+++ b/source/blender/io/usd/import/usd_prim_iterator.cc
@@ -26,7 +26,6 @@
 #include "usd_reader_xformable.h"
 
 #include <iostream>
-#include <pxr/base/plug/registry.h>
 #include <pxr/pxr.h>
 #include <pxr/usd/usd/prim.h>
 #include <pxr/usd/usd/primRange.h>
@@ -38,8 +37,28 @@
 
 namespace blender::io::usd {
 
-USDXformableReader *USDPrimIterator::get_reader(const pxr::UsdPrim &prim,
-                                                const USDImporterContext &context)
+USDPrimIterator::USDPrimIterator(pxr::UsdStageRefPtr stage) : stage_(stage)
+{
+}
+
+void USDPrimIterator::create_object_readers(const USDImporterContext &context,
+                                            std::vector<USDXformableReader *> &r_readers) const
+{
+  if (!stage_) {
+    return;
+  }
+  std::vector<USDXformableReader *> child_readers;
+
+  create_object_readers(stage_->GetPseudoRoot(), context, r_readers, child_readers);
+}
+
+void USDPrimIterator::debug_traverse_stage() const
+{
+  debug_traverse_stage(stage_);
+}
+
+USDXformableReader *USDPrimIterator::get_object_reader(const pxr::UsdPrim &prim,
+                                                       const USDImporterContext &context)
 {
   USDXformableReader *result = nullptr;
 
@@ -53,10 +72,10 @@ USDXformableReader *USDPrimIterator::get_reader(const pxr::UsdPrim &prim,
   return result;
 }
 
-void USDPrimIterator::create_readers(const pxr::UsdPrim &prim,
-                                     const USDImporterContext &context,
-                                     std::vector<USDXformableReader *> &r_readers,
-                                     std::vector<USDXformableReader *> &r_child_readers)
+void USDPrimIterator::create_object_readers(const pxr::UsdPrim &prim,
+                                            const USDImporterContext &context,
+                                            std::vector<USDXformableReader *> &r_readers,
+                                            std::vector<USDXformableReader *> &r_child_readers)
 {
   if (!prim) {
     return;
@@ -69,7 +88,7 @@ void USDPrimIterator::create_readers(const pxr::UsdPrim &prim,
       pxr::UsdTraverseInstanceProxies(pxr::UsdPrimDefaultPredicate));
 
   for (const pxr::UsdPrim &child_prim : child_prims) {
-    create_readers(child_prim, context, r_readers, child_readers);
+    create_object_readers(child_prim, context, r_readers, child_readers);
   }
 
   if (prim.IsPseudoRoot()) {
@@ -100,7 +119,7 @@ void USDPrimIterator::create_readers(const pxr::UsdPrim &prim,
     return;
   }
 
-  USDXformableReader *reader = get_reader(prim, context);
+  USDXformableReader *reader = get_object_reader(prim, context);
 
   if (reader) {
     for (USDXformableReader *child_reader : child_readers) {
diff --git a/source/blender/io/usd/import/usd_prim_iterator.h b/source/blender/io/usd/import/usd_prim_iterator.h
index a1d9a8ad97a..e7a9c5484a6 100644
--- a/source/blender/io/usd/import/usd_prim_iterator.h
+++ b/source/blender/io/usd/import/usd_prim_iterator.h
@@ -28,14 +28,24 @@ struct USDImporterContext;
 class USDXformableReader;
 
 class USDPrimIterator {
+ protected:
+  pxr::UsdStageRefPtr stage_;
+
  public:
-  static USDXformableReader *get_reader(const pxr::UsdPrim &prim,
-                                        const USDImporterContext &context);
+  USDPrimIterator(pxr::UsdStageRefPtr stage);
+
+  void create_object_readers(const USDImporterContext &context,
+                             std::vector<USDXformableReader *> &r_readers) const;
+
+  void debug_traverse_stage() const;
+
+  static USDXformableReader *get_object_reader(const pxr::UsdPrim &prim,
+                                               const USDImporterContext &context);
 
-  static void create_readers(const pxr::UsdPrim &root,
-                             const USDImporterContext &context,
-                             std::vector<USDXformableReader *> &r_readers,
-                             std::vector<USDXformableReader *> &r_child_readers);
+  static void create_object_readers(const pxr::UsdPrim &root,
+                                    const USDImporterContext &context,
+                                    std::vector<USDXformableReader *> &r_readers,
+                                    std::vector<USDXformableReader *> &r_child_readers);
 
   static void debug_traverse_stage(const pxr::UsdStageRefPtr &usd_stage);
 };
diff --git a/source/blender/io/usd/intern/usd_capi.cc b/source/blender/io/usd/intern/usd_capi.cc
index 926713eaecd..aabe02bce04 100644
--- a/source/blender/io/usd/intern/usd_capi.cc
+++ b/source/blender/io/usd/intern/usd_capi.cc
@@ -260,9 +260,11 @@ static void import_startjob(void *user_data, short *stop, short *do_update, floa
   pxr::TfToken up_axis = pxr::UsdGeomGetStageUpAxis(data->stage);
   USDImporterContext import_ctx{up_axis, data->params};
 
+  USDPrimIterator usd_prim_iter(data->stage);
+
   // Optionally print the stage contents for debugging.
   if (data->params.debug) {
-    USDPrimIterator::debug_traverse_stage(data->stage);
+    usd_prim_iter.debug_traverse_stage();
   }
 
   if (G.is_break) {
@@ -273,9 +275,7 @@ static void import_startjob(void *user_data, short *stop, short *do_update, floa
   *data->do_update = true;
   *data->progress = 0.1f;
 
-  std::vector<USDXformableReader *> child_readers;
-  USDPrimIterator::create_readers(
-      data->stage->GetPseudoRoot(), import_ctx, data->readers, child_readers);
+  usd_prim_iter.create_object_readers(import_ctx, data->readers);
 
   // Create objects