NAryTreeTraversal.h

Go to the documentation of this file.

 
#ifndef PBAT_COMMON_NARYTREETRAVERSAL_H
#define PBAT_COMMON_NARYTREETRAVERSAL_H
 
#include "ConstexprFor.h"
#include "Stack.h"
#include "pbat/Aliases.h"
#include "pbat/HostDevice.h"
 
namespace pbat::common {

template <class FVisit, class FChild, class TIndex = Index, auto N = 2, auto kStackDepth = 64>
PBAT_HOST_DEVICE void TraverseNAryTreePseudoPreOrder(FVisit fVisit, FChild fChild, TIndex root = 0);

template <class FVisit, class FChild, class TIndex = Index, auto N = 2, auto kStackDepth = 64>
PBAT_HOST_DEVICE void
TraverseNAryTreePseudoPostOrder(FVisit fVisit, FChild fChild, TIndex root = 0);
 
template <class FVisit, class FChild, class TIndex, auto N, auto kStackDepth>
PBAT_HOST_DEVICE void TraverseNAryTreePseudoPreOrder(FVisit fVisit, FChild fChild, TIndex root)
{
    Stack<TIndex, kStackDepth> dfs{};
    dfs.Push(root);
    while (not dfs.IsEmpty())
    {
        TIndex const node = dfs.Pop();
        if (not fVisit(node))
            continue;
 
        ForRange<0, N>([&]<auto i> {
            auto const child = fChild.template operator()<N - i - 1>(node);
            if (child >= 0)
                dfs.Push(child);
        });
    }
}
 
template <class FVisit, class FChild, class TIndex, auto N, auto kStackDepth>
PBAT_HOST_DEVICE void TraverseNAryTreePseudoPostOrder(FVisit fVisit, FChild fChild, TIndex root)
{
    struct Visit
    {
        TIndex node;
        bool bChildrenVisited;
    };
    Stack<Visit, N * kStackDepth> dfs{};
    dfs.Push({root, false});
    while (not dfs.IsEmpty())
    {
        auto const visit = dfs.Top();
        dfs.Pop();
        if (visit.bChildrenVisited)
        {
            fVisit(visit.node);
        }
        else
        {
            dfs.Push({visit.node, true});
            ForRange<0, N>([&]<auto i> {
                auto const child = fChild.template operator()<N - i - 1>(visit.node);
                if (child >= 0)
                    dfs.Push({child, false});
            });
        }
    }
}
 
} // namespace pbat::common
 
#endif // PBAT_COMMON_NARYTREETRAVERSAL_H