usd/ndr/registry.cpp

//
// Copyright 2018 Pixar
//
// Licensed under the Apache License, Version 2.0 (the "Apache License")
// with the following modification; you may not use this file except in
// compliance with the Apache License and the following modification to it:
// Section 6. Trademarks. is deleted and replaced with:
//
// 6. Trademarks. This License does not grant permission to use the trade
//    names, trademarks, service marks, or product names of the Licensor
//    and its affiliates, except as required to comply with Section 4(c) of
//    the License and to reproduce the content of the NOTICE file.
//
// You may obtain a copy of the Apache License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the Apache License with the above modification is
// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the Apache License for the specific
// language governing permissions and limitations under the Apache License.
//

#include "pxr/pxr.h"
#include "pxr/base/tf/pathUtils.h"
#include "pxr/base/tf/stringUtils.h"
#include "pxr/base/tf/type.h"
#include "pxr/base/trace/trace.h"
#include "pxr/base/work/loops.h"
#include "pxr/base/work/withScopedParallelism.h"
#include "pxr/usd/ar/resolver.h"
#include "pxr/usd/ndr/debugCodes.h"
#include "pxr/usd/ndr/discoveryPlugin.h"
#include "pxr/usd/ndr/node.h"
#include "pxr/usd/ndr/nodeDiscoveryResult.h"
#include "pxr/usd/ndr/property.h"
#include "pxr/usd/ndr/registry.h"
#include "pxr/usd/sdf/types.h"

#include "pxr/base/plug/registry.h"
#include "pxr/base/tf/envSetting.h"

#include <boost/functional/hash.hpp>

PXR_NAMESPACE_OPEN_SCOPE

TF_DEFINE_ENV_SETTING(
    PXR_NDR_SKIP_DISCOVERY_PLUGIN_DISCOVERY, 0,
    "The auto-discovery of discovery plugins in ndr can be skipped. "
    "This is used mostly for testing purposes.");

TF_DEFINE_ENV_SETTING(
    PXR_NDR_SKIP_PARSER_PLUGIN_DISCOVERY, 0,
    "The auto-discovery of parser plugins in ndr can be skipped. "
    "This is used mostly for testing purposes.");

// This function is used for property validation. It is written as a non-static
// freestanding function so that we can exercise it in a test without needing
// to expose it in the header file.  It is also written without using unique
// pointers for ease of python wrapping and testability.
NDR_API
bool
NdrRegistry_ValidateProperty(
    const NdrNodeConstPtr& node,
    const NdrPropertyConstPtr& property,
    std::string* errorMessage)
{
    const VtValue& defaultValue = property->GetDefaultValue();
    const NdrSdfTypeIndicator sdfTypeIndicator = property->GetTypeAsSdfType();
    const SdfValueTypeName sdfType = sdfTypeIndicator.first;

    // We allow default values to be unspecified, but if they aren't empty, then
    // we want to error if the value's type is different from the specified type
    // for the property.
    if (!defaultValue.IsEmpty()) {
        if (defaultValue.GetType() != sdfType.GetType()) {

            if (errorMessage) {
                *errorMessage = TfStringPrintf(
                    "Default value type does not match specified type for "
                    "property.\n"
                    "Node identifier: %s\n"
                    "Source type: %s\n"
                    "Property name: %s.\n"
                    "Type from SdfType: %s.\n"
                    "Type from default value: %s.\n",
                    node->GetIdentifier().GetString().c_str(),
                    node->GetSourceType().GetString().c_str(),
                    property->GetName().GetString().c_str(),
                    sdfType.GetType().GetTypeName().c_str(),
                    defaultValue.GetType().GetTypeName().c_str());
            }

            return false;
        }
    }
    return true;
}

namespace {

// Helpers to allow template functions to treat discovery results and
// nodes equally.
template <typename T> struct _NdrObjectAccess { };
template <> struct _NdrObjectAccess<NdrNodeDiscoveryResult> {
    typedef NdrNodeDiscoveryResult Type;
    static const std::string& GetName(const Type& x) { return x.name; }
    static const TfToken& GetFamily(const Type& x) { return x.family; }
    static NdrVersion GetVersion(const Type& x) { return x.version; }
};
template <> struct _NdrObjectAccess<NdrNodeUniquePtr> {
    typedef NdrNodeUniquePtr Type;
    static const std::string& GetName(const Type& x) { return x->GetName(); }
    static const TfToken& GetFamily(const Type& x) { return x->GetFamily(); }
    static NdrVersion GetVersion(const Type& x) { return x->GetVersion(); }
};

template <typename T>
static
bool
_MatchesFamilyAndFilter(
    const T& object,
    const TfToken& family,
    NdrVersionFilter filter)
{
    using Access = _NdrObjectAccess<T>;

    // Check the family.
    if (!family.IsEmpty() && family != Access::GetFamily(object)) {
        return false;
    }

    // Check the filter.
    switch (filter) {
    case NdrVersionFilterDefaultOnly:
        if (!Access::GetVersion(object).IsDefault()) {
            return false;
        }
        break;

    default:
        break;
    }

    return true;
}

static NdrIdentifier
_GetIdentifierForAsset(const SdfAssetPath &asset,
                       const NdrTokenMap &metadata,
                       const TfToken &subIdentifier,
                       const TfToken &sourceType)
{
    size_t h = 0;
    boost::hash_combine(h, asset);
    for (const auto &i : metadata) {
        boost::hash_combine(h, i.first.GetString());
        boost::hash_combine(h, i.second);
    }

    return NdrIdentifier(TfStringPrintf(
        "%s<%s><%s>",
        std::to_string(h).c_str(),
        subIdentifier.GetText(),
        sourceType.GetText()));
}

static NdrIdentifier
_GetIdentifierForSourceCode(const std::string &sourceCode,
                            const NdrTokenMap &metadata)
{
    size_t h = 0;
    boost::hash_combine(h, sourceCode);
    for (const auto &i : metadata) {
        boost::hash_combine(h, i.first.GetString());
        boost::hash_combine(h, i.second);
    }
    return NdrIdentifier(std::to_string(h));
}

static bool
_ValidateProperty(
    const NdrNodeConstPtr& node,
    const NdrPropertyConstPtr& property)
{
    std::string errorMessage;
    if (!NdrRegistry_ValidateProperty(node, property, &errorMessage)) {
        // This warning may eventually want to be a runtime error and return
        // false to indicate an invalid node, but we didn't want to introduce
        // unexpected behaviors by introducing this error.
        TF_WARN(errorMessage);
    }
    return true;
}

static
bool
_ValidateNode(const NdrNodeUniquePtr &newNode,
              const NdrNodeDiscoveryResult &dr)
{
    // Validate the node.
    if (!newNode) {
        TF_RUNTIME_ERROR("Parser for asset @%s@ of type %s returned null",
            dr.resolvedUri.c_str(), dr.discoveryType.GetText());
        return false;
    }

    // The node is invalid; continue without further error checking.
    //
    // XXX -- WBN if these were just automatically copied and parser plugins
    //        didn't have to deal with them.
    if (newNode->IsValid() &&
        !(newNode->GetIdentifier() == dr.identifier &&
          newNode->GetName() == dr.name &&
          newNode->GetVersion() == dr.version &&
          newNode->GetFamily() == dr.family &&
          newNode->GetSourceType() == dr.sourceType)) {
        TF_RUNTIME_ERROR(
               "Parsed node %s:%s:%s:%s:%s doesn't match discovery result "
               "created for asset @%s@ - "
               "%s:%s:%s:%s:%s (identifier:version:name:family:source type); "
               "discarding.",
               NdrGetIdentifierString(newNode->GetIdentifier()).c_str(),
               newNode->GetVersion().GetString().c_str(),
               newNode->GetName().c_str(),
               newNode->GetFamily().GetText(),
               newNode->GetSourceType().GetText(),
               dr.resolvedUri.c_str(),
               NdrGetIdentifierString(dr.identifier).c_str(),
               dr.version.GetString().c_str(),
               dr.name.c_str(),
               dr.family.GetText(),
               dr.sourceType.GetText());
        return false;
    }

    // It is safe to get the raw pointer from the unique pointer here since
    // this raw pointer will not be passed beyond the scope of this function.
    NdrNodeConstPtr node = newNode.get();

    // Validate the node's properties.  Always validate each property even if
    // we have already found an invalid property because we want to report
    // errors on all properties.
    bool valid = true;
    for (const TfToken& inputName : newNode->GetInputNames()) {
        const NdrPropertyConstPtr& input = newNode->GetInput(inputName);
        valid &= _ValidateProperty(node, input);
    }

    for (const TfToken& outputName : newNode->GetOutputNames()) {
        const NdrPropertyConstPtr& output = newNode->GetOutput(outputName);
        valid &= _ValidateProperty(node, output);
    }

    return valid;
}

} // anonymous namespace

class NdrRegistry::_DiscoveryContext : public NdrDiscoveryPluginContext {
public:
    _DiscoveryContext(const NdrRegistry& registry) : _registry(registry) { }
    ~_DiscoveryContext() override = default;

    TfToken GetSourceType(const TfToken& discoveryType) const override
    {
        auto parser = _registry._GetParserForDiscoveryType(discoveryType);
        return parser ? parser->GetSourceType() : TfToken();
    }

private:
    const NdrRegistry& _registry;
};

NdrRegistry::NdrRegistry()
{
    TRACE_FUNCTION();
    _FindAndInstantiateParserPlugins();
    _FindAndInstantiateDiscoveryPlugins();
    _RunDiscoveryPlugins(_discoveryPlugins);
}

NdrRegistry::~NdrRegistry()
{
    // nothing yet
}

void
NdrRegistry::SetExtraDiscoveryPlugins(DiscoveryPluginRefPtrVec plugins)
{
    {
        std::lock_guard<std::mutex> nmLock(_nodeMapMutex);

        // This policy was implemented in order to keep internal registry
        // operations simpler, and it "just makes sense" to have all plugins
        // run before asking for information from the registry.
        if (!_nodeMap.empty()) {
            TF_CODING_ERROR("SetExtraDiscoveryPlugins() cannot be called after"
                            " nodes have been parsed; ignoring.");
            return;
        }
    }

    _RunDiscoveryPlugins(plugins);

    _discoveryPlugins.insert(_discoveryPlugins.end(),
                             std::make_move_iterator(plugins.begin()),
                             std::make_move_iterator(plugins.end()));
}

void
NdrRegistry::SetExtraDiscoveryPlugins(const std::vector<TfType>& pluginTypes)
{
    // Validate the types and remove duplicates.
    std::set<TfType> discoveryPluginTypes;
    auto& discoveryPluginType = TfType::Find<NdrDiscoveryPlugin>();
    for (auto&& type: pluginTypes) {
        if (!TF_VERIFY(type.IsA(discoveryPluginType),
                       "Type %s is not a %s",
                       type.GetTypeName().c_str(),
                       discoveryPluginType.GetTypeName().c_str())) {
            return;
        }
        discoveryPluginTypes.insert(type);
    }

    // Instantiate any discovery plugins that were found
    DiscoveryPluginRefPtrVec discoveryPlugins;
    for (const TfType& discoveryPluginType : discoveryPluginTypes) {
        NdrDiscoveryPluginFactoryBase* pluginFactory =
            discoveryPluginType.GetFactory<NdrDiscoveryPluginFactoryBase>();

        if (TF_VERIFY(pluginFactory)) {
            discoveryPlugins.emplace_back(pluginFactory->New());
        }
    }

    // Add the discovery plugins.
    SetExtraDiscoveryPlugins(std::move(discoveryPlugins));
}

void
NdrRegistry::SetExtraParserPlugins(const std::vector<TfType>& pluginTypes)
{
    {
        std::lock_guard<std::mutex> nmLock(_nodeMapMutex);

        // This policy was implemented in order to keep internal registry
        // operations simpler, and it "just makes sense" to have all plugins
        // run before asking for information from the registry.
        if (!_nodeMap.empty()) {
            TF_CODING_ERROR("SetExtraParserPlugins() cannot be called after"
                            " nodes have been parsed; ignoring.");
            return;
        }
    }

    // Validate the types and remove duplicates.
    std::set<TfType> parserPluginTypes;
    auto& parserPluginType = TfType::Find<NdrParserPlugin>();
    for (auto&& type: pluginTypes) {
        if (!TF_VERIFY(type.IsA(parserPluginType),
                       "Type %s is not a %s",
                       type.GetTypeName().c_str(),
                       parserPluginType.GetTypeName().c_str())) {
            return;
        }
        parserPluginTypes.insert(type);
    }

    _InstantiateParserPlugins(parserPluginTypes);
}

NdrNodeConstPtr
NdrRegistry::GetNodeFromAsset(const SdfAssetPath &asset,
                              const NdrTokenMap &metadata,
                              const TfToken &subIdentifier,
                              const TfToken &sourceType)
{
    // Ensure there is a parser plugin that can handle this asset.
    TfToken discoveryType(ArGetResolver().GetExtension(asset.GetAssetPath()));
    auto parserIt = _parserPluginMap.find(discoveryType);

    // Ensure that there is a parser registered corresponding to the
    // discoveryType of the asset.
    if (parserIt == _parserPluginMap.end()) {
        TF_DEBUG(NDR_PARSING).Msg("Encountered a asset @%s@ of type [%s], but "
                                  "a parser for the type could not be found; "
                                  "ignoring.\n", asset.GetAssetPath().c_str(),
                                  discoveryType.GetText());
        return nullptr;
    }

    NdrIdentifier identifier =
        _GetIdentifierForAsset(asset, metadata, subIdentifier, sourceType);

    // Use given sourceType if there is one, else use sourceType from the parser
    // plugin.
    const TfToken &thisSourceType = (!sourceType.IsEmpty()) ? sourceType :
        parserIt->second->GetSourceType();
    NodeMapKey key{identifier, thisSourceType};

    // Return the existing node in the map if an entry for the constructed node
    // key already exists.
    std::unique_lock<std::mutex> nmLock(_nodeMapMutex);
    auto it = _nodeMap.find(key);
    if (it != _nodeMap.end()) {
        // Get the raw ptr from the unique_ptr
        return it->second.get();
    }

    // Ensure the map is not locked at this point. The parse is the bulk of the
    // operation, and concurrency is the most valuable here.
    nmLock.unlock();

    // Construct a NdrNodeDiscoveryResult object to pass into the parser
    // plugin's Parse() method.
    // XXX: Should we try resolving the assetPath if the resolved path is empty.
    std::string resolvedUri = asset.GetResolvedPath().empty() ?
        asset.GetAssetPath() : asset.GetResolvedPath();

    NdrNodeDiscoveryResult dr(identifier,
                              NdrVersion(), /* use an invalid version */
                              /* name */ TfGetBaseName(resolvedUri),
                              /*family*/ TfToken(),
                              discoveryType,
                              /* sourceType */ thisSourceType,
                              /* uri */ asset.GetAssetPath(),
                              resolvedUri,
                              /* sourceCode */ "",
                              metadata,
                              /* blindData */ "",
                              /* subIdentifier */ subIdentifier);

    NdrNodeUniquePtr newNode = parserIt->second->Parse(dr);

    if (!_ValidateNode(newNode, dr)) {
        return nullptr;
    }

    // Move the discovery result into _discoveryResults so the node can be found
    // in the Get*() methods
    {
        std::lock_guard<std::mutex> drLock(_discoveryResultMutex);
        _discoveryResultIndicesBySourceType[dr.sourceType].push_back(
            _discoveryResults.size());
        _discoveryResults.emplace_back(std::move(dr));
    }

    nmLock.lock();

    NodeMap::const_iterator result =
        _nodeMap.emplace(std::move(key), std::move(newNode));

    // Get the unique_ptr from the iterator, then get its raw ptr
    return result->second.get();
}

NdrNodeConstPtr
NdrRegistry::GetNodeFromSourceCode(const std::string &sourceCode,
                                   const TfToken &sourceType,
                                   const NdrTokenMap &metadata)
{
    // Ensure that there is a parser registered corresponding to the
    // given sourceType.
    NdrParserPlugin *parserForSourceType = nullptr;
    for (const auto &parserIt : _parserPlugins) {
        if (parserIt->GetSourceType() == sourceType) {
            parserForSourceType = parserIt.get();
        }
    }

    if (!parserForSourceType) {
        // XXX: Should we try looking for sourceType in _parserPluginMap,
        // in case it corresponds to a discovery type in Ndr?

        TF_DEBUG(NDR_PARSING).Msg("Encountered source code of type [%s], but "
                                  "a parser for the type could not be found; "
                                  "ignoring.\n", sourceType.GetText());
        return nullptr;
    }

    NdrIdentifier identifier = _GetIdentifierForSourceCode(sourceCode,
            metadata);
    NodeMapKey key{identifier, sourceType};

    // Return the existing node in the map if an entry for the constructed node
    // key already exists.
    std::unique_lock<std::mutex> nmLock(_nodeMapMutex);
    auto it = _nodeMap.find(key);
    if (it != _nodeMap.end()) {
        // Get the raw ptr from the unique_ptr
        return it->second.get();
    }

    // Ensure the map is not locked at this point. The parse is the bulk of the
    // operation, and concurrency is the most valuable here.
    nmLock.unlock();

    NdrNodeDiscoveryResult dr(identifier,
                              NdrVersion(), /* use an invalid version */
                              /* name */ identifier,
                              /*family*/ TfToken(),
                              // XXX: Setting discoveryType also to sourceType.
                              // Do ParserPlugins rely on it? If yes, should they?
                              /* discoveryType */ sourceType,
                              sourceType,
                              /* uri */ "",
                              /* resolvedUri */ "",
                               sourceCode,
                               metadata);

    NdrNodeUniquePtr newNode = parserForSourceType->Parse(dr);
    if (!newNode) {
        TF_RUNTIME_ERROR("Could not create node for the given source code of "
            "source type '%s'.", sourceType.GetText());
        return nullptr;
    }

    // Move the discovery result into _discoveryResults so the node can be found
    // in the Get*() methods
    {
        std::lock_guard<std::mutex> drLock(_discoveryResultMutex);
        _discoveryResultIndicesBySourceType[dr.sourceType].push_back(
            _discoveryResults.size());
        _discoveryResults.emplace_back(std::move(dr));
    }

    nmLock.lock();

    NodeMap::const_iterator result =
        _nodeMap.emplace(std::move(key), std::move(newNode));

    // Get the unique_ptr from the iterator, then get its raw ptr
    return result->second.get();
}

NdrStringVec
NdrRegistry::GetSearchURIs() const
{
    NdrStringVec searchURIs;

    for (const NdrDiscoveryPluginRefPtr& dp : _discoveryPlugins) {
        NdrStringVec uris = dp->GetSearchURIs();

        searchURIs.insert(searchURIs.end(),
                          std::make_move_iterator(uris.begin()),
                          std::make_move_iterator(uris.end()));
    }

    return searchURIs;
}

NdrIdentifierVec
NdrRegistry::GetNodeIdentifiers(
    const TfToken& family, NdrVersionFilter filter) const
{
    //
    // This should not trigger a parse because node names come directly from
    // the discovery process.
    //

    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);

    NdrIdentifierVec result;
    result.reserve(_discoveryResults.size());

    NdrIdentifierSet visited;
    for (const NdrNodeDiscoveryResult& dr : _discoveryResults) {
        if (_MatchesFamilyAndFilter(dr, family, filter)) {
            // Avoid duplicates.
            if (visited.insert(dr.identifier).second) {
                result.push_back(dr.identifier);
            }
        }
    }

    return result;
}

NdrStringVec
NdrRegistry::GetNodeNames(const TfToken& family) const
{
    //
    // This should not trigger a parse because node names come directly from
    // the discovery process.
    //

    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);

    NdrStringVec nodeNames;
    nodeNames.reserve(_discoveryResults.size());

    NdrStringSet visited;
    for (const NdrNodeDiscoveryResult& dr : _discoveryResults) {
        if (family.IsEmpty() || dr.family == family) {
            // Avoid duplicates.
            if (visited.insert(dr.name).second) {
                nodeNames.push_back(dr.name);
            }
        }
    }

    return nodeNames;
}

NdrNodeConstPtr
NdrRegistry::GetNodeByIdentifier(
    const NdrIdentifier& identifier, const NdrTokenVec& sourceTypePriority)
{
    TRACE_FUNCTION();
    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);

    // If the type priority specifier is empty, pick the first node that matches
    // the identifier regardless of source type.
    if (sourceTypePriority.empty()) {
        // We check for any node that matches the identifier first. If no
        // matching node is found we'll try to find a node with an alias
        // matching the identifier.
        for (const NdrNodeDiscoveryResult& dr : _discoveryResults) {
            if (NdrNodeConstPtr node =
                    _ParseNodeMatchingIdentifier(dr, identifier)) {
                return node;
            }
        }

        for (const NdrNodeDiscoveryResult& dr : _discoveryResults) {
            if (NdrNodeConstPtr node =
                    _ParseNodeMatchingAlias(dr, identifier)) {
                return node;
            }
        }
    }

    // Otherwise we attempt to get a node for matching the identifier (possibly
    // through an alias) for each source type in priority order.
    for (const TfToken& sourceType : sourceTypePriority) {
        if (NdrNodeConstPtr node =
                _GetNodeByIdentifierAndTypeImpl(identifier, sourceType)) {
            return node;
        }
    }

    return nullptr;
}

NdrNodeConstPtr
NdrRegistry::GetNodeByIdentifierAndType(
    const NdrIdentifier& identifier, const TfToken& sourceType)
{
    TRACE_FUNCTION();
    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);

    return _GetNodeByIdentifierAndTypeImpl(identifier, sourceType);
}

NdrNodeConstPtr
NdrRegistry::_GetNodeByIdentifierAndTypeImpl(
    const NdrIdentifier& identifier, const TfToken& sourceType)
{
    // Note that this impl function doesn't lock the discovery result mutex as
    // it expects the caller to have appropriately locked it.

    const auto it = _discoveryResultIndicesBySourceType.find(sourceType);
    if (it == _discoveryResultIndicesBySourceType.end()) {
        return nullptr;
    }

    const std::vector<size_t> &indices = it->second;

    // We check for any node that matches the identifier first. If no
    // matching node is found we'll try to find a node with an alias
    // matching the identifier.
    for (size_t index : indices) {
        const NdrNodeDiscoveryResult &dr = _discoveryResults[index];
        if (NdrNodeConstPtr node = _ParseNodeMatchingIdentifier(dr, identifier)) {
            return node;
        }
    }
    for (size_t index : indices) {
        const NdrNodeDiscoveryResult &dr = _discoveryResults[index];
        if (NdrNodeConstPtr node = _ParseNodeMatchingAlias(dr, identifier)) {
            return node;
        }
    }

    return nullptr;
}

NdrNodeConstPtr
NdrRegistry::GetNodeByName(
    const std::string& name,
    const NdrTokenVec& sourceTypePriority,
    NdrVersionFilter filter)
{
    TRACE_FUNCTION();
    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);

    // If the type priority specifier is empty, pick the first node that matches
    // the name
    if (sourceTypePriority.empty()) {
        for (const NdrNodeDiscoveryResult& dr : _discoveryResults) {
            if (NdrNodeConstPtr node =
                    _ParseNodeMatchingNameAndFilter(dr, name, filter)) {
                return node;
            }
        }
    }

    // Although this is a doubly-nested loop, the number of types in the
    // priority list should be small as should the number of nodes.
    for (const TfToken& sourceType : sourceTypePriority) {
        if (NdrNodeConstPtr node =
                _GetNodeByNameAndTypeImpl(name, sourceType, filter)) {
            return node;
        }
    }

    return nullptr;
}

NdrNodeConstPtr
NdrRegistry::GetNodeByNameAndType(
    const std::string& name, const TfToken& sourceType, NdrVersionFilter filter)
{
    TRACE_FUNCTION();
    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);

    return _GetNodeByNameAndTypeImpl(name, sourceType, filter);
}

NdrNodeConstPtr
NdrRegistry::_GetNodeByNameAndTypeImpl(
    const std::string& name, const TfToken& sourceType,
    NdrVersionFilter filter)
{
    // Note that this impl function doesn't lock the discovery result mutex as
    // it expects the caller to have appropriately locked it.

    const auto it = _discoveryResultIndicesBySourceType.find(sourceType);
    if (it == _discoveryResultIndicesBySourceType.end()) {
        return nullptr;
    }

    const std::vector<size_t> &indices = it->second;

    for (size_t index : indices) {
        const NdrNodeDiscoveryResult &dr = _discoveryResults[index];
        if (NdrNodeConstPtr node =
                _ParseNodeMatchingNameAndFilter(dr, name, filter)) {
            return node;
        }
    }
    return nullptr;

}

NdrNodeConstPtrVec
NdrRegistry::GetNodesByIdentifier(const NdrIdentifier& identifier)
{
    TRACE_FUNCTION();
    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);
    NdrNodeConstPtrVec parsedNodes;

    for (const NdrNodeDiscoveryResult& dr : _discoveryResults) {
        if (NdrNodeConstPtr node =
                _ParseNodeMatchingIdentifier(dr, identifier)) {
            parsedNodes.push_back(node);
        }
    }

    for (const NdrNodeDiscoveryResult& dr : _discoveryResults) {
        if (NdrNodeConstPtr node = _ParseNodeMatchingAlias(dr, identifier)) {
            // In the extremely unlikely case that a node is supplied with an
            // alias that is the same as its identifier, we just make sure it
            // doesn't show up in the list twice as it will have already been
            // added by the above loop.
            if (ARCH_LIKELY(dr.identifier != identifier)) {
                parsedNodes.push_back(node);
            }
        }
    }

    return parsedNodes;
}

NdrNodeConstPtrVec
NdrRegistry::GetNodesByName(const std::string& name, NdrVersionFilter filter)
{
    TRACE_FUNCTION();
    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);
    NdrNodeConstPtrVec parsedNodes;

    for (const NdrNodeDiscoveryResult& dr : _discoveryResults) {
        if (NdrNodeConstPtr node = _ParseNodeMatchingNameAndFilter(
                dr, name, filter)) {
            parsedNodes.push_back(node);
        }
    }

    return parsedNodes;
}

NdrNodeConstPtrVec
NdrRegistry::GetNodesByFamily(const TfToken& family, NdrVersionFilter filter)
{
    // Locking the discovery results for the entire duration of the parse is a
    // bit heavy-handed, but it needs to be 100% guaranteed that the results are
    // not modified while they are being iterated over.
    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);

    // This method does a multi-threaded "bulk parse" of all discovered nodes
    // (or a partial parse if a family is specified). It's possible that another
    // node access method (potentially triggering a parse) could be called in
    // another thread during bulk parse. In that scenario, the worst that should
    // happen is that one of the parses (either from the other method, or this
    // bulk parse) is discarded in favor of the other parse result
    // (_InsertNodeIntoCache() will guard against nodes of the same name and
    // type from being cached).
    {
        std::lock_guard<std::mutex> nmLock(_nodeMapMutex);

        // Skip parsing if a parse was already completed for all nodes
        if (_nodeMap.size() == _discoveryResults.size()) {
            return _GetNodeMapAsNodePtrVec(family, filter);
        }
    }

    // Do the parsing. We need to release the Python GIL here to avoid
    // deadlocks since the code running in the worker threads may call into
    // Python and try to take the GIL when loading plugins. We also need
    // to use scoped parallelism to ensure we don't pick up other tasks
    // during the call to WorkParallelForN that may reenter this function
    // and also deadlock.
    {
        TF_PY_ALLOW_THREADS_IN_SCOPE();

        WorkWithScopedParallelism([&]() {
            WorkParallelForN(_discoveryResults.size(),
                [&](size_t begin, size_t end) {
                    for (size_t i = begin; i < end; ++i) {
                        const NdrNodeDiscoveryResult& dr = _discoveryResults.at(i);
                        if (_MatchesFamilyAndFilter(dr, family, filter)) {
                            _InsertNodeIntoCache(dr);
                        }
                    }
                });
            }
        );
    }
    // Expose the concurrent map as a normal vector to the outside world
    return _GetNodeMapAsNodePtrVec(family, filter);
}

NdrTokenVec
NdrRegistry::GetAllNodeSourceTypes() const
{
    // We're using the _discoveryResultMutex because we populate the
    // _availableSourceTypes while creating the _discoveryResults.
    //
    // We also have to return the source types by value instead of by const
    // reference because we don't want a client holding onto the reference
    // to read from it when _RunDiscoveryPlugins could potentially be running
    // and modifying _availableSourceTypes
    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);
    NdrTokenVec availableSourceTypes;
    availableSourceTypes.reserve(_discoveryResultIndicesBySourceType.size());
    for (const auto &valuePair : _discoveryResultIndicesBySourceType) {
        availableSourceTypes.push_back(valuePair.first);
    }
    return availableSourceTypes;
}

NdrNodeConstPtr
NdrRegistry::_ParseNodeMatchingIdentifier(
    const NdrNodeDiscoveryResult& dr, const NdrIdentifier& identifier)
{
    if (dr.identifier == identifier) {
        return _InsertNodeIntoCache(dr);
    }
    return nullptr;
}

NdrNodeConstPtr
NdrRegistry::_ParseNodeMatchingAlias(
    const NdrNodeDiscoveryResult& dr, const NdrIdentifier& identifier)
{
    for (const TfToken &alias : dr.aliases) {
        if (alias == identifier) {
            return _InsertNodeIntoCache(dr);
        }
    }
    return nullptr;
}

NdrNodeConstPtr
NdrRegistry::_ParseNodeMatchingNameAndFilter(
    const NdrNodeDiscoveryResult& dr, const std::string& name,
    NdrVersionFilter filter)
{
    // Check the filter.
    if (filter == NdrVersionFilterDefaultOnly && !dr.version.IsDefault()) {
        return nullptr;
    }

    if (dr.name == name) {
        return _InsertNodeIntoCache(dr);
    }
    return nullptr;
}

void
NdrRegistry::_FindAndInstantiateDiscoveryPlugins()
{
    // The auto-discovery of discovery plugins can be skipped. This is mostly
    // for testing purposes.
    if (TfGetEnvSetting(PXR_NDR_SKIP_DISCOVERY_PLUGIN_DISCOVERY)) {
        return;
    }

    // Find all of the available discovery plugins
    std::set<TfType> discoveryPluginTypes;
    PlugRegistry::GetInstance().GetAllDerivedTypes<NdrDiscoveryPlugin>(
        &discoveryPluginTypes);

    // Instantiate any discovery plugins that were found
    for (const TfType& discoveryPluginType : discoveryPluginTypes) {
        TF_DEBUG(NDR_DISCOVERY).Msg(
            "Found NdrDiscoveryPlugin '%s'\n",
            discoveryPluginType.GetTypeName().c_str());

        NdrDiscoveryPluginFactoryBase* pluginFactory =
            discoveryPluginType.GetFactory<NdrDiscoveryPluginFactoryBase>();

        if (TF_VERIFY(pluginFactory)) {
            _discoveryPlugins.emplace_back(pluginFactory->New());
        }
    }
}

void
NdrRegistry::_FindAndInstantiateParserPlugins()
{
    // The auto-discovery of parser plugins can be skipped. This is mostly
    // for testing purposes.
    if (TfGetEnvSetting(PXR_NDR_SKIP_PARSER_PLUGIN_DISCOVERY)) {
        return;
    }

    // Find all of the available parser plugins
    std::set<TfType> parserPluginTypes;
    PlugRegistry::GetInstance().GetAllDerivedTypes<NdrParserPlugin>(
        &parserPluginTypes);

    _InstantiateParserPlugins(parserPluginTypes);
}

void
NdrRegistry::_InstantiateParserPlugins(
    const std::set<TfType>& parserPluginTypes)
{
    // Instantiate any parser plugins that were found
    for (const TfType& parserPluginType : parserPluginTypes) {
        TF_DEBUG(NDR_DISCOVERY).Msg(
            "Found NdrParserPlugin '%s' for discovery types:\n",
            parserPluginType.GetTypeName().c_str());

        NdrParserPluginFactoryBase* pluginFactory =
            parserPluginType.GetFactory<NdrParserPluginFactoryBase>();

        if (!TF_VERIFY(pluginFactory)) {
            continue;
        }

        NdrParserPlugin* parserPlugin = pluginFactory->New();
        _parserPlugins.emplace_back(parserPlugin);

        for (const TfToken& discoveryType : parserPlugin->GetDiscoveryTypes()) {
            TF_DEBUG(NDR_DISCOVERY).Msg("  - %s\n", discoveryType.GetText());

            auto i = _parserPluginMap.insert({discoveryType, parserPlugin});
            if (!i.second){
                const TfType otherType = TfType::Find(*i.first->second);
                TF_CODING_ERROR("Plugin type %s claims discovery type '%s' "
                                "but that's already claimed by type %s",
                                parserPluginType.GetTypeName().c_str(),
                                discoveryType.GetText(),
                                otherType.GetTypeName().c_str());
            }
        }
    }
}

void
NdrRegistry::_RunDiscoveryPlugins(const DiscoveryPluginRefPtrVec& discoveryPlugins)
{
    std::lock_guard<std::mutex> drLock(_discoveryResultMutex);

    for (const NdrDiscoveryPluginRefPtr& dp : discoveryPlugins) {
        NdrNodeDiscoveryResultVec results =
            dp->DiscoverNodes(_DiscoveryContext(*this));
        _discoveryResults.reserve(_discoveryResults.size() + results.size());

        for (NdrNodeDiscoveryResult &dr : results) {
            _discoveryResultIndicesBySourceType[dr.sourceType].push_back(
                _discoveryResults.size());
            _discoveryResults.emplace_back(std::move(dr));
        }
    }
}

NdrNodeConstPtr
NdrRegistry::_InsertNodeIntoCache(const NdrNodeDiscoveryResult& dr)
{
    // Return an existing node in the map if the new node matches the
    // identifier AND source type of a node in the map.
    std::unique_lock<std::mutex> nmLock(_nodeMapMutex);
    NodeMapKey key{dr.identifier, dr.sourceType};
    auto it = _nodeMap.find(key);
    if (it != _nodeMap.end()) {
        // Get the raw ptr from the unique_ptr
        return it->second.get();
    }

    // Ensure the map is not locked at this point. The parse is the bulk of the
    // operation, and concurrency is the most valuable here.
    nmLock.unlock();

    // Ensure there is a parser plugin that can handle this node
    auto i = _parserPluginMap.find(dr.discoveryType);
    if (i == _parserPluginMap.end()) {
        TF_DEBUG(NDR_PARSING).Msg("Encountered a node of type [%s], "
                                  "with name [%s], but a parser for that type "
                                  "could not be found; ignoring.\n",
                                  dr.discoveryType.GetText(),  dr.name.c_str());
        return nullptr;
    }

    NdrNodeUniquePtr newNode = i->second->Parse(dr);

    // Validate the node.
    if (!_ValidateNode(newNode, dr)) {
        return nullptr;
    }

    nmLock.lock();

    NodeMap::const_iterator result =
        _nodeMap.emplace(std::move(key), std::move(newNode));

    // Get the unique_ptr from the iterator, then get its raw ptr
    return result->second.get();
}

NdrNodeConstPtrVec
NdrRegistry::_GetNodeMapAsNodePtrVec(
    const TfToken& family, NdrVersionFilter filter) const
{
    NdrNodeConstPtrVec _nodeVec;

    for (const auto& nodePair : _nodeMap) {
        if (_MatchesFamilyAndFilter(nodePair.second, family, filter)) {
            _nodeVec.emplace_back(nodePair.second.get());
        }
    }

    return _nodeVec;
}

NdrParserPlugin*
NdrRegistry::_GetParserForDiscoveryType(const TfToken& discoveryType) const
{
    auto i = _parserPluginMap.find(discoveryType);
    return i == _parserPluginMap.end() ? nullptr : i->second;
}

PXR_NAMESPACE_CLOSE_SCOPE