TransformIterator.h

#pragma once
#include "AL/usdmaya/Common.h"

#include "maya/MObject.h"
#include "maya/MDagPath.h"

#include "pxr/usd/sdf/layer.h"
#include "pxr/usd/usd/prim.h"

#include <vector>

namespace AL {
namespace usdmaya {
namespace fileio {

//----------------------------------------------------------------------------------------------------------------------
//----------------------------------------------------------------------------------------------------------------------
class TransformIterator
{
public:

  TransformIterator(UsdStageRefPtr stage, const MDagPath& parentPath = MDagPath());

  TransformIterator(const UsdPrim& startPrim, const MDagPath& startMayaPath);

  inline bool done() const
    { return m_primStack.empty(); }

  inline size_t depth() const
    { return m_primStack.size(); }

  inline const UsdPrim& prim() const
    { return (m_primStack.end() - 1)->m_prim; }

  void prune();

  bool next();

  inline void append(const MObject newNode)
    {
      m_primStack[m_primStack.size() - 1].m_object = newNode;
    }

  inline MObject parent() const
    {
      if(m_primStack.size() > 1)
      {
        if( (m_primStack.end() - 2)->m_object != MObject::kNullObj)
        {
          return (m_primStack.end() - 2)->m_object;
        }
        if( (m_primStack.end() - 1)->m_object != MObject::kNullObj)
        {
          return (m_primStack.end() - 1)->m_object;
        }
      }
      if(m_primStack.size() == 1)
      {
        return (m_primStack.end() - 1)->m_object;
      }
      return m_parentPath.node();
    }

  MDagPath currentPath() const;

private:
  struct StackRef
  {
    StackRef(const UsdPrim& prim);
    StackRef(const StackRef& prim);
    StackRef();
    UsdPrim m_prim;
    MObject m_object;
    UsdPrim::SiblingIterator m_begin;
    UsdPrim::SiblingIterator m_end;
    int32_t m_currentChild;
    int32_t m_numChildren;
  };

  std::vector<StackRef> m_primStack;
  UsdStageRefPtr m_stage;
  size_t m_currentItem;
  MDagPath m_parentPath;

  TfHashSet<SdfPath, SdfPath::Hash> m_visitedMasterPrimPaths;
};

//----------------------------------------------------------------------------------------------------------------------
} // fileio
} // usdmaya
} // AL
//----------------------------------------------------------------------------------------------------------------------