viewer/src/sample_tools/sample_atlas_packer.cpp

#ifdef RIVE_RENDERER_TESS
#include "viewer/sample_tools/sample_atlas_packer.hpp"
#include "utils/no_op_factory.hpp"
#include "viewer/tess/bitmap_decoder.hpp"
#include "rive/file.hpp"
#include "rive/assets/image_asset.hpp"
#include "rive/tess/sokol/sokol_tess_renderer.hpp"
#include <algorithm>

using namespace rive;

class AtlasRenderImage : public RenderImage {
private:
    std::vector<uint8_t> m_Pixels;

public:
    AtlasRenderImage(const uint8_t* pixels, uint32_t width, uint32_t height) :
        m_Pixels(pixels, pixels + (width * height * 4)) {
        m_Width = width;
        m_Height = height;
    }

    Span<const uint8_t> pixels() { return m_Pixels; }
};

class AtlasPackerFactory : public NoOpFactory {
    std::unique_ptr<RenderImage> decodeImage(Span<const uint8_t> bytes) override {
        auto bitmap = Bitmap::decode(bytes);
        if (bitmap) {
            // We have a bitmap, let's make an image.

            // For now only deal with RGBA.
            if (bitmap->pixelFormat() != Bitmap::PixelFormat::RGBA) {
                bitmap->pixelFormat(Bitmap::PixelFormat::RGBA);
            }

            return std::make_unique<AtlasRenderImage>(bitmap->bytes(),
                                                      bitmap->width(),
                                                      bitmap->height());
        }
        return nullptr;
    }
};

SampleAtlasPacker::SampleAtlasPacker(uint32_t maxWidth, uint32_t maxHeight) :
    m_maxWidth(maxWidth), m_maxHeight(maxHeight) {}

void SampleAtlasPacker::pack(Span<const uint8_t> rivBytes) {
    AtlasPackerFactory factory;
    if (auto file = rive::File::import(rivBytes, &factory)) {
        for (auto asset : file->assets()) {
            if (asset->is<ImageAsset>()) {
                Mat2D uvTransform;

                auto imageAsset = asset->as<ImageAsset>();
                auto renderImage = static_cast<AtlasRenderImage*>(imageAsset->renderImage());

                if (m_atlases.empty()) {
                    // Make the first atlas.
                    m_atlases.push_back(new SampleAtlas(m_maxWidth, m_maxHeight));
                }

                // Pack into the current atlas.
                if (!m_atlases.back()->pack(renderImage->pixels().data(),
                                            renderImage->width(),
                                            renderImage->height(),
                                            uvTransform))
                {
                    // Didn't fit in previous atlas, make a new one.
                    auto nextAtlas = new SampleAtlas(m_maxWidth, m_maxHeight);

                    // Try to pack into the new one.
                    if (!nextAtlas->pack(renderImage->pixels().data(),
                                         renderImage->width(),
                                         renderImage->height(),
                                         uvTransform))
                    {
                        // Still failed. Image must be larger than max atlas. Just push the whole
                        // image as an atlas.
                        m_atlases.push_back(new SampleAtlas(renderImage->pixels().data(),
                                                            renderImage->width(),
                                                            renderImage->height()));

                        // Clean up unused next atlas.
                        delete nextAtlas;
                    } else {
                        m_atlases.push_back(nextAtlas);
                    }
                }

                // Store where the image ended up.

                m_lookup[asset->assetId()] = {
                    .atlasIndex = m_atlases.size() - 1,
                    .transform = uvTransform,
                    .width = (uint32_t)renderImage->width(),
                    .height = (uint32_t)renderImage->height(),
                };
            }
        }
    }
}

SampleAtlasPacker::~SampleAtlasPacker() {
    for (auto atlas : m_atlases) {
        delete atlas;
    }
}

SampleAtlas* SampleAtlasPacker::atlas(std::size_t index) {
    assert(index < m_atlases.size());
    return m_atlases[index];
}

SampleAtlas::SampleAtlas(const uint8_t* pixels, uint32_t width, uint32_t height) :
    m_width(width), m_height(height), m_pixels(pixels, pixels + width * height * 4) {}

SampleAtlas::SampleAtlas(uint32_t width, uint32_t height) :
    m_width(width), m_height(height), m_pixels(width * height * 4) {}

bool SampleAtlas::pack(const uint8_t* sourcePixels,
                       uint32_t width,
                       uint32_t height,
                       Mat2D& packTransform) {
    if (m_x + width >= m_width) {
        m_x = 0;
        m_y = m_nextY;
    }
    // Check if we overflow vertically, we're done.
    if (m_y + height >= m_height) {
        return false;
    }

    // We fit, pack into m_x, m_y.
    for (uint32_t sy = 0; sy < height; sy++) {
        for (uint32_t sx = 0; sx < width; sx++) {
            for (uint8_t channel = 0; channel < 4; channel++) {
                m_pixels[((m_y + sy) * m_width + (m_x + sx)) * 4 + channel] =
                    sourcePixels[(sy * width + sx) * 4 + channel];
            }
        }
    }

    // Set the UV transform.
    packTransform = {
        width / (float)m_width,
        0,
        0,
        height / (float)m_height,
        m_x / (float)m_width,
        m_y / (float)m_height,
    };

    // Increment internal positions.
    m_x += width;
    if (m_y + height > m_nextY) {
        m_nextY = m_y + height;
    }

    return true;
}

bool SampleAtlasPacker::find(const ImageAsset& asset, SampleAtlasLocation* location) {
    auto assetId = asset.assetId();
    auto result = m_lookup.find(assetId);
    if (result != m_lookup.end()) {
        *location = result->second;
        return true;
    }
    return false;
}

SampleAtlasResolver::SampleAtlasResolver(SampleAtlasPacker* packer) : m_packer(packer) {}

void SampleAtlasResolver::loadContents(FileAsset& asset) {
    if (asset.is<ImageAsset>()) {
        SampleAtlasLocation location;
        auto imageAsset = asset.as<ImageAsset>();

        // Find which location this image got packed into.
        if (m_packer->find(*imageAsset, &location)) {
            // Determine if we've already loaded the a render image.
            auto sharedItr = m_sharedImageResources.find(location.atlasIndex);

            rive::rcp<rive::SokolRenderImageResource> imageResource;
            if (sharedItr != m_sharedImageResources.end()) {
                imageResource = sharedItr->second;
            } else {
                auto atlas = m_packer->atlas(location.atlasIndex);
                imageResource = rive::rcp<rive::SokolRenderImageResource>(
                    new SokolRenderImageResource(atlas->pixels().data(),
                                                 atlas->width(),
                                                 atlas->height()));
                m_sharedImageResources[location.atlasIndex] = imageResource;
            }

            // Make a render image from the atlas if we haven't yet. Our Factory
            // doesn't have an abstraction for creating a RenderImage from a
            // decoded set of pixels, we should either add that or live with/be
            // ok with the fact that most people writing a resolver doing
            // something like this will likely also have a custom/specific
            // renderer (and hence will know which RenderImage they need to
            // make).

            imageAsset->renderImage(std::make_unique<SokolRenderImage>(imageResource,
                                                                       location.width,
                                                                       location.height,
                                                                       location.transform));
        }
    }
}
#endif