weightedIterator.h

Go to the documentation of this file.

//
// Copyright 2025 Pixar
//
// Licensed under the terms set forth in the LICENSE.txt file available at
// https://openusd.org/license.
//
#ifndef PXR_EXEC_VDF_WEIGHTED_ITERATOR_H
#define PXR_EXEC_VDF_WEIGHTED_ITERATOR_H
 
 
#include "pxr/pxr.h"
 
#include "pxr/exec/vdf/connection.h"
#include "pxr/exec/vdf/indexedWeights.h"
#include "pxr/exec/vdf/iterator.h"
#include "pxr/exec/vdf/node.h"
#include "pxr/exec/vdf/vector.h"
 
#include "pxr/base/tf/diagnostic.h"
 
#include <cstddef>
#include <initializer_list>
 
PXR_NAMESPACE_OPEN_SCOPE
 
struct Vdf_WeightSlot
{
    const VdfIndexedWeights *weights;
 
    size_t currentIndex;
};
 
class Vdf_WeightSlotArray
{
    Vdf_WeightSlotArray(const Vdf_WeightSlotArray &) = delete;
    Vdf_WeightSlotArray& operator=(const Vdf_WeightSlotArray &) = delete;
 
public:
    using value_type = Vdf_WeightSlot;
    using reference = Vdf_WeightSlot &;
    using const_reference = const Vdf_WeightSlot &;
    using iterator = Vdf_WeightSlot *;
    using const_iterator = const Vdf_WeightSlot *;
 
    Vdf_WeightSlotArray()
        : _begin(nullptr)
        , _end(nullptr)
    {}
 
    ~Vdf_WeightSlotArray() {
        if (_begin != _local) {
            delete[] _begin;
        }
    }
 
    void Allocate(uint32_t numInputs) {
        if (!TF_VERIFY(!_begin && !_end)) {
            return;
        }
 
        _begin = numInputs > NUM_LOCAL_STORAGE
            ? new value_type[numInputs]
            : _local;
        _end = _begin + numInputs;
    }
 
    size_t size() const {
        return _end - _begin;
    }
 
    iterator begin() {
        return _begin;
    }
 
    iterator end() {
        return _end;
    }
 
    const_iterator begin() const {
        return _begin;
    }
 
    const_iterator end() const {
        return _end;
    }
 
    //
    reference operator[](size_t slot) {
        return _begin[slot];
    }
 
    //
    const_reference operator[](size_t slot) const {
        return _begin[slot];
    }
 
    int AdvanceToNextExplicitIndex(int index);
 
private:
    value_type *_begin;
    value_type *_end;
 
    // The amount of local storage reserved for slots.
    // A single slot is stored locally, multiple slots are stored remotely.
    static const size_t NUM_LOCAL_STORAGE = 1;
 
    value_type _local[NUM_LOCAL_STORAGE];
};
 
 
template<class IteratorType>
class VdfWeightedIterator final : public VdfIterator 
{
public:
 
    typedef typename IteratorType::value_type value_type;
 
    typedef typename IteratorType::reference reference;
 
    template<typename... Args>
    VdfWeightedIterator(
        const VdfContext &context,
        const TfToken &weightName,
        Args&&... args) :
        _iterator(context, std::forward<Args>(args)...) {
        const TfToken *weightNamePtr = &weightName;
        _Init(context, weightNamePtr, weightNamePtr + 1);
    }
 
    template<typename... Args>
    VdfWeightedIterator(
        const VdfContext &context,
        std::initializer_list<TfToken> weightNames,
        Args&&... args) :
        _iterator(context, std::forward<Args>(args)...) {
        _Init(context, weightNames.begin(), weightNames.end());    
    }
 
    template<typename... Args>
    VdfWeightedIterator(
        const VdfContext &context,
        const std::vector<TfToken> &weightNames,
        Args&&... args) :
        _iterator(context, std::forward<Args>(args)...) {
        const TfToken *weightNamesBegin = weightNames.data();
        const TfToken *weightNamesEnd = weightNamesBegin + weightNames.size();
        _Init(context, weightNamesBegin, weightNamesEnd);
    }
 
    ~VdfWeightedIterator() = default;
 
    VdfWeightedIterator &operator++();
 
    reference operator*() const {
        return _iterator.operator*();
    }
 
    bool IsAtEnd() const {
        return _iterator.IsAtEnd();
    }
 
    int GetCurrentIndex() const { 
        return Vdf_GetIteratorIndex(_iterator);
    }
 
    void AdvanceToEnd() {
        _iterator.AdvanceToEnd();
    }
 
    double GetWeight(size_t slot = 0, double defWeight = 0.0) const {
        _GetExplicitWeight(slot, &defWeight);
        return defWeight;
    }
 
    bool HasExplicitWeight(size_t slot) const {
        double ret;
        return _GetExplicitWeight(slot, &ret);
    }
 
    std::pair<bool, double> GetExplicitWeight(
        size_t slot = 0, double defWeight = 0.0) const;
 
    size_t GetNumSlots() const {
        return _slots.size();
    }
 
    size_t GetNumExplicitWeights(size_t slot = 0) const {
        return (slot < _slots.size()) ?
            _slots[slot].weights->GetSize() : 0;
    }
 
private:
 
    // Noncopyable
    VdfWeightedIterator(const VdfWeightedIterator &) = delete;
    VdfWeightedIterator &operator=(const VdfWeightedIterator &) = delete;
 
    // Returns the current index into the data source.
    friend int Vdf_GetIteratorIndex(const VdfWeightedIterator &it) {
        return Vdf_GetIteratorIndex(it._iterator);
    }
 
    // Two constructors above call this shared code upon construction
    void _Init(
        const VdfContext &context,
        const TfToken *begin,
        const TfToken *end);
 
    // Extracts weighted data from weightInput and adds it to our vector
    // of weight vectors.
    void _AddWeightsInput(
        Vdf_WeightSlot *slot,
        const VdfInput *weightInput,
        const VdfContext &context);
 
    // Returns true and the weigth in *ret if the weight is explicity set for
    // this slot, otherwise, it returns false.
    bool _GetExplicitWeight(size_t slot, double *ret) const;
 
    // Implements the template parameter specified behavior
    // 
    // The affects mask iterator is advanced to the next explicit weight
    // only if the template parameter is SkipElementsWithNoWeights.
    // The currentIndex is always updated.
    // 
    // Advances the current _iterator to the first index where we have
    // both a weight explicity set and an element set in the mask.
    void _AdvanceIterator();
 
private:
 
    // The underlying iterator.
    IteratorType _iterator;
    Vdf_WeightSlotArray _slots;
};
 
 
template<class IteratorType>
void
VdfWeightedIterator<IteratorType>::_Init(
    const VdfContext &context,
    const TfToken *begin,
    const TfToken *end)
{
    // If there's nothing set in the mask, there's no need to go on.
    if (_iterator.IsAtEnd())  {
        return;
    }
 
    // Get the number of slots.
    const std::ptrdiff_t numInputs = std::distance(begin, end);
 
    // Reserve storage for the slots, if required.
    if (numInputs > 0) {
        _slots.Allocate(numInputs);
    }
 
    // Add all the weights inputs.
    Vdf_WeightSlot *slot = _slots.begin();
    const VdfNode &node = _GetNode(context);
    for (const TfToken *it = begin; it != end; ++it, ++slot) {
        const VdfInput *input = node.GetInput(*it);
        if (!input) {
            TF_CODING_ERROR(
                "Can't find input '%s' on node %s",
                it->GetText(),
                node.GetDebugName().c_str());
        }
        _AddWeightsInput(slot, input, context);
    }
 
    // No inputs.
    if (numInputs == 0) {
        TF_CODING_ERROR("Weighted Iterator instantiated with no weights.");
        return;
    }   
    
    // By here, *_iterator has been initialized with the first set mask.
    // What we'd like to do is advance it to the next index such that
    // we have both an explicit weight and a set element in the mask.
    _AdvanceIterator();
}
 
template<class IteratorType>
void
VdfWeightedIterator<IteratorType>::_AddWeightsInput(
    Vdf_WeightSlot *slot,
    const VdfInput *weightInput,
    const VdfContext &context)
{
    // We always expect exactly one input connection.
    if (weightInput && weightInput->GetNumConnections() == 1) {
        const VdfConnection &connection  = (*weightInput)[0];
        const VdfVector &out = _GetRequiredInputValue(
            context, connection, connection.GetMask());
 
        // We always expect exactly one element.
        if (ARCH_LIKELY(out.GetSize() == 1)) {
            // Note that here we hold on to a pointer from the executor because
            // we don't want to copy the vector of weights.  It is okay to hold
            // on to the pointer only because the lifetime of this iterator is
            // limited.
            slot->weights = &out.GetReadAccessor<VdfIndexedWeights>()[0];
            slot->currentIndex = 0;
            return;
        }
    
        TF_CODING_ERROR("Weight input must have exactly one element (got %zu)",
                        out.GetSize());
    }
 
    else if (weightInput && weightInput->GetNumConnections() > 1) {
        // This is an error, all weight connectors must have exactly one input.
        TF_CODING_ERROR("Weight connector must have at most one input (got %zu)",
                        weightInput->GetNumConnections());
    }
 
    // Not exactly one input connection and data element.
    slot->weights = nullptr;
    slot->currentIndex = 0;
}
 
template<class IteratorType>
VdfWeightedIterator<IteratorType> &
VdfWeightedIterator<IteratorType>::operator++()
{
    // We need to differentiate between two cases here: 
    //
    // a) There can be holes in the mask _iterator works on, we want to skip
    //    those fast.
    // b) There can be holes in the explicit weights, which we also want to 
    //    skip fast.
 
    // Advance iterator to the next element as indicated by mask, this may
    // or may not skip holes in the mask.
 
    _iterator.operator++();
    _AdvanceIterator();
 
    return *this;
}
 
template<class IteratorType>
void
VdfWeightedIterator<IteratorType>::_AdvanceIterator()
{
    while (!_iterator.IsAtEnd()) {
 
        // Find the next index that has an explicit weight after currentIndex
        const int currentIndex = Vdf_GetIteratorIndex(_iterator);
        const int nextExplicitIndex =
            _slots.AdvanceToNextExplicitIndex(currentIndex);
 
        // If there are no more explicit weights, we're done.
        if (nextExplicitIndex == std::numeric_limits<int>::max()) {
            // We're done iterating, set _iterator to end().
            _iterator.AdvanceToEnd();
            break;
 
        } else if (nextExplicitIndex != currentIndex) {
 
            // In this case the next explicit weight was further along than
            // our iterator.  So let's try to advance the iterator.
            while (!_iterator.IsAtEnd() &&
                Vdf_GetIteratorIndex(_iterator) < nextExplicitIndex) {
                ++_iterator;
            }
 
            if (_iterator.IsAtEnd() ||
                Vdf_GetIteratorIndex(_iterator) == nextExplicitIndex) {
                // Great!  Both the iterator and the next explicit weights
                // have a value, or we have reached the end of iterator.  We're
                // done.
                break;
            }
 
            // If we get here, the iterator did not visit an element at the
            // explicit index, and we have now advanced it beyond the explicit
            // index. Let's retry at the current index.
 
        } else {
            // There is a next explicit weight at our current iterator index. 
            // We're done. 
            break;
        }
    }
}
 
inline int 
Vdf_WeightSlotArray::AdvanceToNextExplicitIndex(
    int index)
{
    int nextExplicitIndex = std::numeric_limits<int>::max();
 
    // Iterate over all the weight slots.
    for (Vdf_WeightSlot &p : *this) {
        const VdfIndexedWeights *w = p.weights;
 
        // If we're done with that slot, don't bother trying to find one.
        if (ARCH_UNLIKELY(!w || p.currentIndex >= w->GetSize())) {
            continue;
        }
 
        TF_DEV_AXIOM(w);
 
        // We'll try to do a quick local search from our last known
        // position, for speed.
        p.currentIndex = w->GetFirstDataIndex(index, p.currentIndex);
        if (p.currentIndex < w->GetSize()) {
            // We already found the next index that we care about, we
            // won't do any more work than is necessary.
            int wIndexVal = w->GetIndex(p.currentIndex);
            if (wIndexVal < nextExplicitIndex) {
                nextExplicitIndex = wIndexVal;
            }
        }
    }
 
    return nextExplicitIndex;
}
 
template<class IteratorType>
std::pair<bool, double>
VdfWeightedIterator<IteratorType>::GetExplicitWeight(
    size_t slot, double defWeight) const
{
    bool hasExplicitWeight = _GetExplicitWeight(slot, &defWeight);
    return std::pair<bool, double>(hasExplicitWeight, defWeight);
}
 
template<class IteratorType>
inline bool
VdfWeightedIterator<IteratorType>::_GetExplicitWeight(
    size_t slot, double *ret) const
{
    TF_DEV_AXIOM(ret);
 
    if (slot < _slots.size()) {
        const Vdf_WeightSlot &p = _slots[slot];
        if (p.weights &&
            p.currentIndex < p.weights->GetSize() &&
            p.weights->GetIndex(p.currentIndex) ==
                Vdf_GetIteratorIndex(_iterator)) {
            *ret = p.weights->GetData(p.currentIndex);
            return true;
        }
    }
    return false;
}
 
PXR_NAMESPACE_CLOSE_SCOPE
 
#endif