MobileGL/tools/trace_replay/fixture_tools/realign_buffer_ranges.cpp

// MobileGL - tools/trace_replay/fixture_tools/realign_buffer_ranges.cpp
// Copyright (c) 2025-2026 MobileGL-Dev
// Licensed under the GNU Lesser General Public License v3.0:
//   https://www.gnu.org/licenses/gpl-3.0.txt
//   https://www.gnu.org/licenses/lgpl-3.0.txt
// SPDX-License-Identifier: LGPL-3.0-only
// End of Source File Header

// realign_buffer_ranges: make a GL trace's indexed uniform / storage buffer range binds portable.
//
// An app sizes glBindBufferRange offsets by the GL_{UNIFORM,SHADER_STORAGE}_BUFFER_OFFSET_ALIGNMENT
// of the driver it was captured on (Flywheel's staging scatter lists: 16 on desktop). A driver that
// requires more rejects those binds with GL_INVALID_VALUE. This tool rewrites every such bind whose
// offset is not a multiple of --align:
//   - the bind itself becomes glBindBuffer(target, buffer) (the generic binding it also set);
//   - right before the next draw / dispatch, the bytes the range holds at that point are uploaded into a
//     fresh --align-aligned slot of one relocation buffer, and that slot is bound at the same index.
// The bytes come from the trace itself (buffer data / subdata and the persistent-map memcpy commits), so a
// source buffer written any other way (GPU copies, clears) is refused rather than guessed.
//
// Inserting calls renumbers every later call: --map-call N prints where original call N landed (a case's
// target_call). The output is snappy-compressed; `apitrace repack --brotli` it before packaging.
//
// Build (A = 3rdparty/apitrace, B = a Release build of it, e.g. `cmake -S $A -B $B -G Ninja -DENABLE_GUI=OFF
// && ninja -C $B apitrace`):
//   g++ -std=c++17 -O2 -flto=auto -fno-rtti -fno-exceptions -I$A/lib/trace -I$A/lib/os -I$B -I$B/lib/trace \
//     realign_buffer_ranges.cpp -o realign_buffer_ranges $B/lib/trace/libcommon.a $B/lib/guids/libguids.a \
//     $B/lib/highlight/libhighlight.a $B/lib/os/libos.a $B/thirdparty/libbacktrace.a $B/thirdparty/libsnappy.a \
//     $B/thirdparty/libzstd_seekable.a -lbrotlidec -lbrotlienc -lz -lzstd -lpthread -ldl
// Usage: realign_buffer_ranges [--align N] [--map-call N]... in.trace out.trace

#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <map>
#include <set>
#include <string>
#include <unordered_map>
#include <vector>

#include "trace_model.hpp"
#include "trace_parser.hpp"
#include "trace_writer.hpp"

namespace {

constexpr unsigned kUniformBuffer = 0x8A11;
constexpr unsigned kShaderStorageBuffer = 0x90D2;
constexpr unsigned kDynamicStorageBit = 0x0100;
constexpr unsigned kRelocBufferName = 1000000;

[[noreturn]] void Fail(const trace::Call *call, const std::string &what) {
    std::fprintf(stderr, "error: call %u %s: %s\n", call ? call->no : 0u, call ? call->name() : "-", what.c_str());
    std::exit(1);
}

bool Is(const trace::Call *call, const char *name) { return std::strcmp(call->name(), name) == 0; }

bool StartsWith(const char *s, const char *prefix) { return std::strncmp(s, prefix, std::strlen(prefix)) == 0; }

bool ReadsIndexedBuffers(const trace::Call *call) {
    const char *n = call->name();
    return StartsWith(n, "glDraw") || StartsWith(n, "glMultiDraw") || StartsWith(n, "glDispatchCompute");
}

bool IsRangeTarget(unsigned target) { return target == kUniformBuffer || target == kShaderStorageBuffer; }

struct MapRegion {
    unsigned buffer;
    uint64_t offset;
    uint64_t length;
    uint64_t ptr;
};

struct Reloc {
    unsigned src = 0;
    uint64_t offset = 0;
    uint64_t size = 0;
    bool dirty = true;
    std::vector<char> last;
    const trace::EnumSig *targetSig = nullptr;
};

// Signatures for the calls this tool inserts; ids are allocated past the trace's own.
struct Sigs {
    const char *createNames[2] = {"n", "buffers"};
    const char *storageNames[4] = {"buffer", "size", "data", "flags"};
    const char *subDataNames[4] = {"buffer", "offset", "size", "data"};
    const char *bindRangeNames[5] = {"target", "index", "buffer", "offset", "size"};
    const char *bindNames[2] = {"target", "buffer"};
    trace::FunctionSig create{}, storage{}, subData{}, bindRange{}, bind{};

    void Init(trace::Id firstId) {
        create = {firstId + 0, "glCreateBuffers", 2, createNames};
        storage = {firstId + 1, "glNamedBufferStorage", 4, storageNames};
        subData = {firstId + 2, "glNamedBufferSubData", 4, subDataNames};
        bindRange = {firstId + 3, "glBindBufferRange", 5, bindRangeNames};
        bind = {firstId + 4, "glBindBuffer", 2, bindNames};
    }
};

class Rewriter {
public:
    uint64_t align = 1024;
    std::set<unsigned> mapCalls;

    int Run(const char *in, const char *out) {
        Scan(in);
        if (relocSrc_.empty()) {
            std::fprintf(stderr, "nothing to realign at %llu bytes\n", (unsigned long long)align);
            return 2;
        }
        sigs_.Init(maxSigId_ + 1);
        stride_ = (maxRelocSize_ + align - 1) / align * align;
        Pass(in, nullptr);  // dry run: count the slots
        slotCap_ = slotsUsed_;
        Pass(in, out);
        std::fprintf(stderr,
                     "realigned %llu binds of %zu source buffer(s) to %llu bytes: %llu uploads, relocation buffer %llu "
                     "bytes, %llu calls inserted\n",
                     (unsigned long long)relocBinds_, relocSrc_.size(), (unsigned long long)align,
                     (unsigned long long)slotsUsed_, (unsigned long long)(slotCap_ * stride_),
                     (unsigned long long)inserted_);
        for (const auto &m : callMap_) {
            std::printf("call %u -> %llu\n", m.first, (unsigned long long)m.second);
        }
        return 0;
    }

private:
    // Pass 0: which buffers need relocating, and the sizes involved.
    void Scan(const char *in) {
        trace::Parser p;
        if (!p.open(in)) Fail(nullptr, std::string("cannot open ") + in);
        trace::Call *call;
        while ((call = p.parse_call())) {
            maxSigId_ = std::max(maxSigId_, call->sig->id);
            if (Is(call, "glBindBufferRange")) {
                unsigned target = unsigned(call->arg(0).toSInt());
                uint64_t offset = uint64_t(call->arg(3).toSInt());
                if (IsRangeTarget(target) && offset % align != 0) {
                    relocSrc_.insert(unsigned(call->arg(2).toUInt()));
                    maxRelocSize_ = std::max(maxRelocSize_, uint64_t(call->arg(4).toSInt()));
                }
            } else if (Is(call, "glBindBuffersRange")) {
                unsigned target = unsigned(call->arg(0).toSInt());
                const trace::Array *offsets = call->arg(4).toArray();
                if (IsRangeTarget(target) && offsets) {
                    for (const trace::Value *v : offsets->values) {
                        if (uint64_t(v->toSInt()) % align != 0) Fail(call, "misaligned glBindBuffersRange is not handled");
                    }
                }
            }
            delete call;
        }
    }

    void Pass(const char *in, const char *out) {
        trace::Parser p;
        if (!p.open(in)) Fail(nullptr, std::string("cannot open ") + in);
        trace::Writer writer;
        writer_ = nullptr;
        if (out) {
            if (!writer.open(out, p.getVersion(), p.getProperties())) Fail(nullptr, std::string("cannot create ") + out);
            writer_ = &writer;
        }
        shadow_.clear();
        maps_.clear();
        generic_.clear();
        relocs_.clear();
        slotsUsed_ = 0;
        relocBinds_ = 0;
        inserted_ = 0;
        created_ = false;

        trace::Call *call;
        while ((call = p.parse_call())) {
            Process(call);
            delete call;
        }
        if (writer_) writer.close();
        writer_ = nullptr;
    }

    unsigned BufferArg(trace::Call *call, unsigned index) { return unsigned(call->arg(index).toUInt()); }

    unsigned Generic(unsigned target) {
        auto it = generic_.find(target);
        return it == generic_.end() ? 0 : it->second;
    }

    std::vector<char> *Shadow(unsigned buffer) {
        auto it = shadow_.find(buffer);
        return it == shadow_.end() ? nullptr : &it->second;
    }

    void Specify(trace::Call *call, unsigned buffer, uint64_t size, trace::Value *data) {
        if (!relocSrc_.count(buffer)) return;
        std::vector<char> &s = shadow_[buffer];
        s.assign(size, 0);
        const trace::Blob *blob = data ? data->toBlob() : nullptr;
        if (blob) std::memcpy(s.data(), blob->buf, std::min<uint64_t>(size, blob->size));
        for (auto &r : relocs_) {
            if (r.second.src == buffer) r.second.dirty = true;
        }
        (void)call;
    }

    void Write(trace::Call *call, unsigned buffer, uint64_t offset, const char *bytes, uint64_t size) {
        std::vector<char> *s = Shadow(buffer);
        if (!s) return;
        if (offset + size > s->size()) Fail(call, "write past the end of a tracked buffer");
        std::memcpy(s->data() + offset, bytes, size);
    }

    void Refuse(trace::Call *call, unsigned buffer) {
        if (relocSrc_.count(buffer)) Fail(call, "writes a relocated source buffer in a way the trace does not carry the bytes of");
    }

    void Emit(trace::Call *call) {
        if (!writer_) return;
        if (mapCalls.count(call->no)) callMap_[call->no] = writtenCalls_;
        writer_->writeCall(call);
        ++writtenCalls_;
    }

    void EmitNew(trace::Call *call) {
        ++inserted_;
        if (writer_) {
            writer_->writeCall(call);
            ++writtenCalls_;
        }
        delete call;
    }

    trace::Call *NewCall(const trace::FunctionSig *sig, unsigned thread) {
        return new trace::Call(sig, trace::CallFlags(0), thread);
    }

    void EmitBind(unsigned thread, const trace::EnumSig *targetSig, unsigned target, unsigned buffer) {
        trace::Call *c = NewCall(&sigs_.bind, thread);
        c->args[0].value = targetSig ? static_cast<trace::Value *>(new trace::Enum(targetSig, target))
                                     : static_cast<trace::Value *>(new trace::SInt(target));
        c->args[1].value = new trace::UInt(buffer);
        EmitNew(c);
    }

    void CreateRelocBuffer(unsigned thread) {
        trace::Call *c = NewCall(&sigs_.create, thread);
        c->args[0].value = new trace::SInt(1);
        trace::Array *names = new trace::Array(1);
        names->values[0] = new trace::UInt(kRelocBufferName);
        c->args[1].value = names;
        EmitNew(c);

        c = NewCall(&sigs_.storage, thread);
        c->args[0].value = new trace::UInt(kRelocBufferName);
        c->args[1].value = new trace::SInt(int64_t(slotCap_ * stride_));
        c->args[2].value = new trace::Null();
        c->args[3].value = new trace::UInt(kDynamicStorageBit);
        EmitNew(c);
        created_ = true;
    }

    void Upload(trace::Call *before, unsigned target, unsigned index, Reloc &r) {
        std::vector<char> *s = Shadow(r.src);
        if (!s) Fail(before, "a relocated range reads a buffer whose storage the trace never specified");
        if (r.offset + r.size > s->size()) Fail(before, "a relocated range runs past its buffer");
        std::vector<char> bytes(s->begin() + r.offset, s->begin() + r.offset + r.size);
        if (!r.dirty && bytes == r.last) return;

        unsigned thread = before->thread_id;
        if (writer_ && slotsUsed_ >= slotCap_) Fail(before, "relocation slots exceeded the dry run");
        uint64_t slotOffset = slotsUsed_++ * stride_;
        if (!created_) CreateRelocBuffer(thread);

        trace::Call *c = NewCall(&sigs_.subData, thread);
        c->args[0].value = new trace::UInt(kRelocBufferName);
        c->args[1].value = new trace::SInt(int64_t(slotOffset));
        c->args[2].value = new trace::SInt(int64_t(r.size));
        trace::Blob *blob = new trace::Blob(r.size);
        std::memcpy(blob->buf, bytes.data(), r.size);
        c->args[3].value = blob;
        EmitNew(c);

        c = NewCall(&sigs_.bindRange, thread);
        c->args[0].value = r.targetSig ? static_cast<trace::Value *>(new trace::Enum(r.targetSig, target))
                                       : static_cast<trace::Value *>(new trace::SInt(target));
        c->args[1].value = new trace::UInt(index);
        c->args[2].value = new trace::UInt(kRelocBufferName);
        c->args[3].value = new trace::SInt(int64_t(slotOffset));
        c->args[4].value = new trace::SInt(int64_t(r.size));
        EmitNew(c);

        // glBindBufferRange also moved the generic binding; put back the app's.
        EmitBind(thread, r.targetSig, target, Generic(target));

        r.last = std::move(bytes);
        r.dirty = false;
    }

    void Process(trace::Call *call) {
        const char *n = call->name();

        if (Is(call, "glBindBuffer")) {
            generic_[unsigned(call->arg(0).toSInt())] = BufferArg(call, 1);
        } else if (Is(call, "glBindBufferBase")) {
            unsigned target = unsigned(call->arg(0).toSInt());
            generic_[target] = BufferArg(call, 2);
            relocs_.erase({target, unsigned(call->arg(1).toUInt())});
        } else if (Is(call, "glBindBufferRange")) {
            unsigned target = unsigned(call->arg(0).toSInt());
            unsigned index = unsigned(call->arg(1).toUInt());
            unsigned buffer = BufferArg(call, 2);
            uint64_t offset = uint64_t(call->arg(3).toSInt());
            generic_[target] = buffer;
            if (IsRangeTarget(target) && offset % align != 0) {
                Reloc r;
                r.src = buffer;
                r.offset = offset;
                r.size = uint64_t(call->arg(4).toSInt());
                trace::Value &t = call->arg(0);
                trace::Enum *e = nullptr;
                // Enum values print by name in dumps; keep the trace's own enum signature when it has one.
                struct EnumProbe : trace::Visitor {
                    trace::Enum *found = nullptr;
                    void visit(trace::Enum *node) override { found = node; }
                } probe;
                t.visit(probe);
                e = probe.found;
                r.targetSig = e ? e->sig : nullptr;
                relocs_[{target, index}] = std::move(r);
                ++relocBinds_;
                if (mapCalls.count(call->no) && writer_) callMap_[call->no] = writtenCalls_;
                // Only the generic binding survives at this point; the indexed one is bound before its first read.
                EmitBind(call->thread_id, e ? e->sig : nullptr, target, buffer);
                --inserted_;  // a replacement, not an insertion
                return;
            }
            relocs_.erase({target, index});
        } else if (Is(call, "glBindBuffersBase") || Is(call, "glBindBuffersRange")) {
            unsigned target = unsigned(call->arg(0).toSInt());
            unsigned first = unsigned(call->arg(1).toUInt());
            unsigned count = unsigned(call->arg(2).toSInt());
            for (unsigned i = first; i < first + count; ++i) relocs_.erase({target, i});
        } else if (Is(call, "glNamedBufferStorage") || Is(call, "glNamedBufferData")) {
            Specify(call, BufferArg(call, 0), uint64_t(call->arg(1).toSInt()), &call->arg(2));
        } else if (Is(call, "glBufferStorage") || Is(call, "glBufferData")) {
            Specify(call, Generic(unsigned(call->arg(0).toSInt())), uint64_t(call->arg(1).toSInt()), &call->arg(2));
        } else if (Is(call, "glNamedBufferSubData") || Is(call, "glBufferSubData")) {
            unsigned buffer = Is(call, "glNamedBufferSubData") ? BufferArg(call, 0) : Generic(unsigned(call->arg(0).toSInt()));
            if (relocSrc_.count(buffer)) {
                const trace::Blob *blob = call->arg(3).toBlob();
                if (!blob) Fail(call, "subdata without bytes");
                Write(call, buffer, uint64_t(call->arg(1).toSInt()), blob->buf,
                      std::min<uint64_t>(uint64_t(call->arg(2).toSInt()), blob->size));
            }
        } else if (Is(call, "glCopyNamedBufferSubData")) {
            Refuse(call, BufferArg(call, 1));
        } else if (Is(call, "glCopyBufferSubData")) {
            Refuse(call, Generic(unsigned(call->arg(1).toSInt())));
        } else if (StartsWith(n, "glClearNamedBuffer")) {
            Refuse(call, BufferArg(call, 0));
        } else if (StartsWith(n, "glClearBuffer") && (StartsWith(n, "glClearBufferData") || StartsWith(n, "glClearBufferSubData"))) {
            Refuse(call, Generic(unsigned(call->arg(0).toSInt())));
        } else if (Is(call, "glMapNamedBufferRange") || Is(call, "glMapBufferRange")) {
            unsigned buffer = Is(call, "glMapNamedBufferRange") ? BufferArg(call, 0) : Generic(unsigned(call->arg(0).toSInt()));
            if (relocSrc_.count(buffer) && call->ret) {
                maps_.push_back({buffer, uint64_t(call->arg(1).toSInt()), uint64_t(call->arg(2).toSInt()),
                                 call->ret->toUIntPtr()});
            }
        } else if (Is(call, "glMapNamedBuffer") || Is(call, "glMapBuffer")) {
            unsigned buffer = Is(call, "glMapNamedBuffer") ? BufferArg(call, 0) : Generic(unsigned(call->arg(0).toSInt()));
            std::vector<char> *s = Shadow(buffer);
            if (s && call->ret) maps_.push_back({buffer, 0, s->size(), call->ret->toUIntPtr()});
        } else if (Is(call, "glUnmapNamedBuffer") || Is(call, "glUnmapBuffer")) {
            unsigned buffer = Is(call, "glUnmapNamedBuffer") ? BufferArg(call, 0) : Generic(unsigned(call->arg(0).toSInt()));
            maps_.erase(std::remove_if(maps_.begin(), maps_.end(), [&](const MapRegion &m) { return m.buffer == buffer; }),
                        maps_.end());
        } else if (Is(call, "memcpy")) {
            uint64_t dest = call->arg(0).toUIntPtr();
            const trace::Blob *blob = call->arg(1).toBlob();
            uint64_t size = call->arg(2).toUInt();
            if (blob) {
                size = std::min<uint64_t>(size, blob->size);
                for (const MapRegion &m : maps_) {
                    uint64_t lo = std::max(dest, m.ptr);
                    uint64_t hi = std::min(dest + size, m.ptr + m.length);
                    if (lo >= hi) continue;
                    Write(call, m.buffer, m.offset + (lo - m.ptr), blob->buf + (lo - dest), hi - lo);
                }
            }
        } else if (ReadsIndexedBuffers(call)) {
            for (auto &r : relocs_) Upload(call, r.first.first, r.first.second, r.second);
        }

        Emit(call);
    }

    Sigs sigs_;
    trace::Id maxSigId_ = 0;
    std::set<unsigned> relocSrc_;
    uint64_t maxRelocSize_ = 0;
    uint64_t stride_ = 0;
    uint64_t slotCap_ = 0;
    uint64_t slotsUsed_ = 0;
    uint64_t relocBinds_ = 0;
    uint64_t inserted_ = 0;
    uint64_t writtenCalls_ = 0;
    bool created_ = false;
    trace::Writer *writer_ = nullptr;
    std::unordered_map<unsigned, std::vector<char>> shadow_;
    std::vector<MapRegion> maps_;
    std::unordered_map<unsigned, unsigned> generic_;
    std::map<std::pair<unsigned, unsigned>, Reloc> relocs_;
    std::map<unsigned, uint64_t> callMap_;
};

}  // namespace

int main(int argc, char **argv) {
    Rewriter rw;
    std::vector<const char *> files;
    for (int i = 1; i < argc; ++i) {
        if (!std::strcmp(argv[i], "--align") && i + 1 < argc) {
            rw.align = std::strtoull(argv[++i], nullptr, 0);
        } else if (!std::strcmp(argv[i], "--map-call") && i + 1 < argc) {
            rw.mapCalls.insert(unsigned(std::strtoul(argv[++i], nullptr, 0)));
        } else {
            files.push_back(argv[i]);
        }
    }
    if (files.size() != 2 || rw.align == 0) {
        std::fprintf(stderr, "usage: %s [--align N] [--map-call N]... in.trace out.trace\n", argv[0]);
        return 1;
    }
    return rw.Run(files[0], files[1]);
}