omni/extras/OutArrayUtils.h

File members: omni/extras/OutArrayUtils.h

// Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
#pragma once

#include "../core/IObject.h" // for omni::core::Result

#include <type_traits>
#include <vector>

namespace omni
{
namespace extras
{

template <typename T, typename Callable, typename SizeType>
omni::core::Result fillOutArray(T* outArray, SizeType* outArrayCount, SizeType requiredCount, Callable&& fillFn)
{
    static_assert(std::is_unsigned<SizeType>::value, "SizeType must be an unsigned integer type!");

    if (!outArrayCount)
    {
        return omni::core::kResultInvalidArgument;
    }

    if (!outArray)
    {
        *outArrayCount = requiredCount;
        return omni::core::kResultSuccess;
    }

    if (*outArrayCount < requiredCount)
    {
        *outArrayCount = requiredCount;
        return omni::core::kResultInsufficientBuffer;
    }

    *outArrayCount = requiredCount;

    fillFn(outArray, requiredCount);

    return omni::core::kResultSuccess;
}

template <typename T, typename GetCallable, typename FillCallable>
omni::core::Result getOutArray(GetCallable&& getFn,
                               FillCallable&& fillFn,
                               uint32_t stackCount = (4096 / sizeof(T)),
                               uint32_t maxRetryCount = (UINT32_MAX - 1))
{
    // constructors won't run when we call alloca, make sure the type doesn't need a constructor to run.
    static_assert(std::is_trivially_default_constructible<T>::value, "T must be trivially default constructible");

    T* stack = CARB_STACK_ALLOC(T, stackCount);
    std::vector<T> heap;
    T* buffer = stack;
    uint32_t count = stackCount;

    omni::core::Result result = getFn(buffer, &count);
    uint32_t retryCount = maxRetryCount + 1;
    while (--retryCount)
    {
        switch (result)
        {
            case omni::core::kResultSuccess:
                if (buffer)
                {
                    fillFn(buffer, count);
                    return omni::core::kResultSuccess;
                }
                CARB_FALLTHROUGH; // (alloca returned nullptr, count is now correct and we should alloc on heap)
            case omni::core::kResultInsufficientBuffer:
                heap.resize(count);
                buffer = heap.data();
                result = getFn(buffer, &count);
                break;
            default:
                return result;
        }
    }

    return omni::core::kResultTryAgain;
}

} // namespace extras
} // namespace omni