MapLibreParticleLab/src/mbgl/mtl/offscreen_texture.cpp

#include <mbgl/mtl/offscreen_texture.hpp>
#include <mbgl/mtl/context.hpp>
#include <mbgl/mtl/renderable_resource.hpp>
#include <mbgl/mtl/renderer_backend.hpp>
#include <mbgl/mtl/texture2d.hpp>

#include <Metal/Metal.hpp>
#include <stdexcept>

namespace mbgl {
namespace mtl {

class OffscreenTextureResource final : public RenderableResource {
public:
    OffscreenTextureResource(Context& context_,
                             const Size size_,
                             const gfx::TextureChannelDataType type_,
                             bool depth,
                             bool stencil)
        : context(context_),
          size(size_),
          type(type_),
          needsDepth(depth),
          needsStencil(stencil) {
        assert(!size.isEmpty());
        colorTexture = context.createTexture2D();
        colorTexture->setSize(size);
        colorTexture->setFormat(gfx::TexturePixelType::RGBA, type);
        colorTexture->setSamplerConfiguration({.filter = gfx::TextureFilterType::Linear,
                                               .wrapU = gfx::TextureWrapType::Clamp,
                                               .wrapV = gfx::TextureWrapType::Clamp});
        static_cast<Texture2D*>(colorTexture.get())
            ->setUsage(MTL::TextureUsageShaderRead | MTL::TextureUsageShaderWrite | MTL::TextureUsageRenderTarget);

#if TARGET_OS_SIMULATOR || defined(__x86_64__)
        constexpr bool combinedDepthStencil = true;
#else
        constexpr bool combinedDepthStencil = false;
#endif
        if (depth || (stencil && combinedDepthStencil)) {
            depthTexture = context.createTexture2D();
            depthTexture->setSize(size);
            depthTexture->setFormat(gfx::TexturePixelType::Depth, gfx::TextureChannelDataType::Float);
            depthTexture->setSamplerConfiguration({.filter = gfx::TextureFilterType::Linear,
                                                   .wrapU = gfx::TextureWrapType::Clamp,
                                                   .wrapV = gfx::TextureWrapType::Clamp});
            static_cast<Texture2D*>(depthTexture.get())
                ->setUsage(MTL::TextureUsageShaderRead | MTL::TextureUsageShaderWrite | MTL::TextureUsageRenderTarget);
        }

        // On iOS simulator, the depth target is PixelFormatDepth32Float_Stencil8
#if !TARGET_OS_SIMULATOR && !defined(__x86_64__)
        if (stencil) {
            stencilTexture = context.createTexture2D();
            stencilTexture->setSize(size);
            stencilTexture->setFormat(gfx::TexturePixelType::Stencil, gfx::TextureChannelDataType::UnsignedByte);
            stencilTexture->setSamplerConfiguration({.filter = gfx::TextureFilterType::Linear,
                                                     .wrapU = gfx::TextureWrapType::Clamp,
                                                     .wrapV = gfx::TextureWrapType::Clamp});
            static_cast<Texture2D*>(stencilTexture.get())
                ->setUsage(MTL::TextureUsageShaderRead | MTL::TextureUsageShaderWrite | MTL::TextureUsageRenderTarget);
        }
#endif

        context.renderingStats().numFrameBuffers++;
    }

    ~OffscreenTextureResource() noexcept override { context.renderingStats().numFrameBuffers--; }

    void bind() override {
        assert(context.getBackend().getCommandQueue());
        commandBuffer = NS::RetainPtr(context.getBackend().getCommandQueue()->commandBuffer());
        colorTexture->create();
        auto* colorMetalTexture = static_cast<Texture2D*>(colorTexture.get())->getMetalTexture();
        if (!commandBuffer || !colorMetalTexture) throw std::runtime_error("Metal offscreen color attachment unavailable");

        renderPassDescriptor = NS::TransferPtr(MTL::RenderPassDescriptor::alloc()->init());
        if (!renderPassDescriptor) throw std::runtime_error("Metal offscreen render pass descriptor unavailable");
        if (auto* colorTarget = renderPassDescriptor->colorAttachments()->object(0)) {
            colorTarget->setTexture(colorMetalTexture);
        } else {
            throw std::runtime_error("Metal offscreen color attachment missing");
        }

        if (depthTexture) {
            depthTexture->create();
            auto* depthMetalTexture = static_cast<Texture2D*>(depthTexture.get())->getMetalTexture();
            if (!depthMetalTexture) throw std::runtime_error("Metal offscreen depth attachment unavailable");
            if (auto* depthTarget = renderPassDescriptor->depthAttachment()) {
                depthTarget->setTexture(depthMetalTexture);
            } else {
                throw std::runtime_error("Metal offscreen depth attachment missing");
            }
        }
        if (stencilTexture) {
            stencilTexture->create();
            auto* stencilMetalTexture = static_cast<Texture2D*>(stencilTexture.get())->getMetalTexture();
            if (!stencilMetalTexture) throw std::runtime_error("Metal offscreen stencil attachment unavailable");
            if (auto* stencilTarget = renderPassDescriptor->stencilAttachment()) {
                stencilTarget->setTexture(stencilMetalTexture);
            } else {
                throw std::runtime_error("Metal offscreen stencil attachment missing");
            }
        } else if (needsStencil && depthTexture) {
            // Simulator/x86 Metal uses one Depth32Float_Stencil8 texture for both attachments.
            if (auto* stencilTarget = renderPassDescriptor->stencilAttachment()) {
                stencilTarget->setTexture(static_cast<Texture2D*>(depthTexture.get())->getMetalTexture());
            } else {
                throw std::runtime_error("Metal offscreen combined stencil attachment missing");
            }
        } else if (needsStencil) {
            throw std::runtime_error("Metal offscreen stencil attachment unavailable");
        }
    }

    void swap() override {
        assert(commandBuffer);
        commandBuffer->commit();
        commandBuffer->waitUntilCompleted();
        commandBuffer.reset();
        renderPassDescriptor.reset();
    }

    PremultipliedImage readStillImage() {
        assert(static_cast<Texture2D*>(colorTexture.get())->getMetalTexture());

        auto data = std::make_unique<uint8_t[]>(colorTexture->getDataSize());
        MTL::Region region = MTL::Region::Make2D(0, 0, size.width, size.height);
        NS::UInteger bytesPerRow = size.width * colorTexture->getPixelStride();

        static_cast<Texture2D*>(colorTexture.get())->getMetalTexture()->getBytes(data.get(), bytesPerRow, region, 0);

        return {size, std::move(data)};
    }

    gfx::Texture2DPtr& getTexture() {
        assert(colorTexture);
        return colorTexture;
    }

    bool isRenderable() {
        try {
            colorTexture->create();
            if (!static_cast<Texture2D*>(colorTexture.get())->getMetalTexture()) return false;
            if (needsDepth || (needsStencil && !stencilTexture)) {
                if (!depthTexture) return false;
                depthTexture->create();
                if (!static_cast<Texture2D*>(depthTexture.get())->getMetalTexture()) return false;
            }
            if (stencilTexture) {
                stencilTexture->create();
                if (!static_cast<Texture2D*>(stencilTexture.get())->getMetalTexture()) return false;
            }
            return true;
        } catch (...) {
            return false;
        }
    }

    const RendererBackend& getBackend() const override { return context.getBackend(); }

    const MTLCommandBufferPtr& getCommandBuffer() const override { return commandBuffer; }

    MTLBlitPassDescriptorPtr getUploadPassDescriptor() const override {
        return NS::TransferPtr(MTL::BlitPassDescriptor::alloc()->init());
    }

    const MTLRenderPassDescriptorPtr& getRenderPassDescriptor() const override {
        assert(renderPassDescriptor);
        return renderPassDescriptor;
    }

private:
    Context& context;
    const Size size;
    const gfx::TextureChannelDataType type;
    const bool needsDepth;
    const bool needsStencil;
    gfx::Texture2DPtr colorTexture;
    gfx::Texture2DPtr depthTexture;
    gfx::Texture2DPtr stencilTexture;
    MTLCommandBufferPtr commandBuffer;
    MTLRenderPassDescriptorPtr renderPassDescriptor;
};

OffscreenTexture::OffscreenTexture(
    Context& context, const Size size_, const gfx::TextureChannelDataType type, bool depth, bool stencil)
    : gfx::OffscreenTexture(size_, std::make_unique<OffscreenTextureResource>(context, size_, type, depth, stencil)) {}

bool OffscreenTexture::isRenderable() {
    return getResource<OffscreenTextureResource>().isRenderable();
}

PremultipliedImage OffscreenTexture::readStillImage() {
    return getResource<OffscreenTextureResource>().readStillImage();
}

const gfx::Texture2DPtr& OffscreenTexture::getTexture() {
    return getResource<OffscreenTextureResource>().getTexture();
}

} // namespace mtl
} // namespace mbgl