pickle_parser.cpp

#include "pickle_parser.h"

#include <algorithm>
#include <cctype>
#include <cstdint>
#include <iomanip>
#include <limits>
#include <map>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <variant>
#include <vector>

namespace gd {
namespace {

struct PickleMark {};

struct PickleGlobal {
    std::string module;
    std::string name;
};

struct PickleStorage {
    ggml_type type = GGML_TYPE_F32;
    std::string key;
    std::string location;
    uint64_t size = 0;
};

struct PickleTuple;
struct PickleDict;

using PickleValue = std::variant<std::monostate,
                                 PickleMark,
                                 int64_t,
                                 bool,
                                 std::string,
                                 PickleGlobal,
                                 PickleStorage,
                                 TensorDescriptor,
                                 std::shared_ptr<PickleTuple>,
                                 std::shared_ptr<PickleDict>>;

struct PickleTuple {
    std::vector<PickleValue> items;
};

struct PickleDict {
    std::map<std::string, PickleValue> items;
};

uint8_t read_u8(std::istream & input) {
    char byte = 0;
    input.read(&byte, 1);
    return static_cast<uint8_t>(byte);
}

uint16_t read_u16(std::istream & input) {
    uint8_t bytes[2] = {0, 0};
    input.read(reinterpret_cast<char *>(bytes), sizeof(bytes));
    return static_cast<uint16_t>(bytes[0]) |
           (static_cast<uint16_t>(bytes[1]) << 8);
}

uint32_t read_u32(std::istream & input) {
    uint8_t bytes[4] = {0, 0, 0, 0};
    input.read(reinterpret_cast<char *>(bytes), sizeof(bytes));
    return static_cast<uint32_t>(bytes[0]) |
           (static_cast<uint32_t>(bytes[1]) << 8) |
           (static_cast<uint32_t>(bytes[2]) << 16) |
           (static_cast<uint32_t>(bytes[3]) << 24);
}

uint64_t read_u64(std::istream & input) {
    uint8_t bytes[8] = {0, 0, 0, 0, 0, 0, 0, 0};
    input.read(reinterpret_cast<char *>(bytes), sizeof(bytes));
    return static_cast<uint64_t>(bytes[0]) |
           (static_cast<uint64_t>(bytes[1]) << 8) |
           (static_cast<uint64_t>(bytes[2]) << 16) |
           (static_cast<uint64_t>(bytes[3]) << 24) |
           (static_cast<uint64_t>(bytes[4]) << 32) |
           (static_cast<uint64_t>(bytes[5]) << 40) |
           (static_cast<uint64_t>(bytes[6]) << 48) |
           (static_cast<uint64_t>(bytes[7]) << 56);
}

int32_t read_i32(std::istream & input) {
    return static_cast<int32_t>(read_u32(input));
}

ggml_type storage_type_to_ggml(const std::string & storage_name) {
    if (storage_name == "FloatStorage") {
        return GGML_TYPE_F32;
    }
    if (storage_name == "HalfStorage") {
        return GGML_TYPE_F16;
    }
    throw std::runtime_error("Unsupported storage type: " + storage_name);
}

bool is_contiguous_shape(const std::vector<int64_t> & shape, const std::vector<int64_t> & stride) {
    if (shape.size() != stride.size()) {
        return false;
    }

    int64_t expected = 1;
    for (int64_t index = static_cast<int64_t>(shape.size()) - 1; index >= 0; --index) {
        if (stride[static_cast<size_t>(index)] != expected) {
            return false;
        }
        expected *= shape[static_cast<size_t>(index)];
    }
    return true;
}

template <typename T>
const T * get_if(const PickleValue & value) {
    return std::get_if<T>(&value);
}

std::vector<PickleValue> pop_until_mark(std::vector<PickleValue> & stack) {
    std::vector<PickleValue> values;
    while (!stack.empty()) {
        PickleValue value = std::move(stack.back());
        stack.pop_back();
        if (std::holds_alternative<PickleMark>(value)) {
            std::reverse(values.begin(), values.end());
            return values;
        }
        values.push_back(std::move(value));
    }
    throw std::runtime_error("Malformed pickle: missing MARK");
}

std::shared_ptr<PickleTuple> make_tuple_from_values(std::vector<PickleValue> values) {
    auto tuple = std::make_shared<PickleTuple>();
    tuple->items = std::move(values);
    return tuple;
}

std::string require_string(const PickleValue & value, const char * context) {
    if (const std::string * string = get_if<std::string>(value)) {
        return *string;
    }
    throw std::runtime_error(std::string("Expected string in ") + context);
}

int64_t require_int(const PickleValue & value, const char * context) {
    if (const int64_t * integer = get_if<int64_t>(value)) {
        return *integer;
    }
    throw std::runtime_error(std::string("Expected integer in ") + context);
}

bool require_bool(const PickleValue & value, const char * context) {
    if (const bool * boolean = get_if<bool>(value)) {
        return *boolean;
    }
    throw std::runtime_error(std::string("Expected bool in ") + context);
}

const PickleGlobal & require_global(const PickleValue & value, const char * context) {
    if (const PickleGlobal * global = get_if<PickleGlobal>(value)) {
        return *global;
    }
    throw std::runtime_error(std::string("Expected global in ") + context);
}

const PickleStorage & require_storage(const PickleValue & value, const char * context) {
    if (const PickleStorage * storage = get_if<PickleStorage>(value)) {
        return *storage;
    }
    throw std::runtime_error(std::string("Expected storage in ") + context);
}

const std::shared_ptr<PickleTuple> & require_tuple(const PickleValue & value, const char * context) {
    if (const std::shared_ptr<PickleTuple> * tuple = get_if<std::shared_ptr<PickleTuple>>(value)) {
        return *tuple;
    }
    throw std::runtime_error(std::string("Expected tuple in ") + context);
}

const std::shared_ptr<PickleDict> & require_dict(const PickleValue & value, const char * context) {
    if (const std::shared_ptr<PickleDict> * dict = get_if<std::shared_ptr<PickleDict>>(value)) {
        return *dict;
    }
    throw std::runtime_error(std::string("Expected dict in ") + context);
}

std::vector<int64_t> tuple_to_ints(const PickleValue & value, const char * context) {
    const std::shared_ptr<PickleTuple> & tuple = require_tuple(value, context);
    std::vector<int64_t> values;
    values.reserve(tuple->items.size());
    for (const PickleValue & item : tuple->items) {
        values.push_back(require_int(item, context));
    }
    return values;
}

PickleValue reduce_value(const PickleValue & callable, const PickleValue & args) {
    const PickleGlobal & global = require_global(callable, "REDUCE callable");
    const std::shared_ptr<PickleTuple> & tuple = require_tuple(args, "REDUCE args");

    if (global.module == "collections" && global.name == "OrderedDict") {
        return std::make_shared<PickleDict>();
    }

    if (global.module == "torch._utils" && global.name == "_rebuild_parameter") {
        if (tuple->items.empty()) {
            throw std::runtime_error("Unexpected _rebuild_parameter argument count");
        }
        return tuple->items[0];
    }

    if (global.module == "torch._utils" && global.name == "_rebuild_tensor") {
        if (tuple->items.size() != 4) {
            throw std::runtime_error("Unexpected _rebuild_tensor argument count");
        }

        const PickleStorage & storage = require_storage(tuple->items[0], "tensor storage");
        const uint64_t storage_offset = static_cast<uint64_t>(require_int(tuple->items[1], "storage offset"));
        const std::vector<int64_t> shape = tuple_to_ints(tuple->items[2], "shape tuple");
        const std::vector<int64_t> stride = tuple_to_ints(tuple->items[3], "stride tuple");

        if (!is_contiguous_shape(shape, stride)) {
            throw std::runtime_error("Non-contiguous tensor storage is not supported");
        }

        TensorDescriptor tensor;
        tensor.storage_key = storage.key;
        tensor.type = storage.type;
        tensor.storage_offset = storage_offset;
        tensor.storage_size = storage.size;
        tensor.shape = shape;
        tensor.stride = stride;
        return tensor;
    }

    if (global.module == "torch._utils" && global.name == "_rebuild_tensor_v2") {
        if (tuple->items.size() != 6) {
            throw std::runtime_error("Unexpected _rebuild_tensor_v2 argument count");
        }

        const PickleStorage & storage = require_storage(tuple->items[0], "tensor storage");
        const uint64_t storage_offset = static_cast<uint64_t>(require_int(tuple->items[1], "storage offset"));
        const std::vector<int64_t> shape = tuple_to_ints(tuple->items[2], "shape tuple");
        const std::vector<int64_t> stride = tuple_to_ints(tuple->items[3], "stride tuple");
        (void) require_bool(tuple->items[4], "requires_grad");
        (void) require_dict(tuple->items[5], "backward hooks");

        if (!is_contiguous_shape(shape, stride)) {
            throw std::runtime_error("Non-contiguous tensor storage is not supported");
        }

        TensorDescriptor tensor;
        tensor.storage_key = storage.key;
        tensor.type = storage.type;
        tensor.storage_offset = storage_offset;
        tensor.storage_size = storage.size;
        tensor.shape = shape;
        tensor.stride = stride;
        return tensor;
    }

    throw std::runtime_error("Unsupported REDUCE callable: " + global.module + "." + global.name);
}

int64_t decode_long(const std::vector<uint8_t> & bytes) {
    if (bytes.empty()) {
        return 0;
    }
    if (bytes.size() > 8) {
        throw std::runtime_error("Unsupported LONG integer width");
    }

    uint64_t value = 0;
    for (size_t i = 0; i < bytes.size(); ++i) {
        value |= static_cast<uint64_t>(bytes[i]) << (8 * i);
    }

    const bool negative = (bytes.back() & 0x80U) != 0U;
    if (!negative) {
        return static_cast<int64_t>(value);
    }

    const uint64_t mask = ~0ULL << (bytes.size() * 8);
    value |= mask;
    return static_cast<int64_t>(value);
}

void collect_tensors(const PickleValue & value, std::vector<TensorDescriptor> & tensors) {
    if (const TensorDescriptor * tensor = get_if<TensorDescriptor>(value)) {
        tensors.push_back(*tensor);
        return;
    }

    if (const std::shared_ptr<PickleDict> * dict = get_if<std::shared_ptr<PickleDict>>(value)) {
        for (const auto & item : (*dict)->items) {
            collect_tensors(item.second, tensors);
        }
        return;
    }

    if (const std::shared_ptr<PickleTuple> * tuple = get_if<std::shared_ptr<PickleTuple>>(value)) {
        for (const PickleValue & item : (*tuple)->items) {
            collect_tensors(item, tensors);
        }
    }
}

}  // namespace

std::vector<TensorDescriptor> parse_tensor_descriptors(const std::vector<uint8_t> & payload) {
    std::vector<PickleValue> stack;
    std::unordered_map<uint32_t, PickleValue> memo;
    uint32_t memo_next = 0;
    std::istringstream input(std::string(reinterpret_cast<const char *>(payload.data()), payload.size()), std::ios::binary);

    auto memo_store = [&](uint32_t index) {
        if (stack.empty()) {
            throw std::runtime_error("Malformed pickle: empty stack during memo store");
        }
        memo[index] = stack.back();
        memo_next = std::max<uint32_t>(memo_next, index + 1);
    };

    auto memo_store_auto = [&]() {
        if (stack.empty()) {
            throw std::runtime_error("Malformed pickle: empty stack during memo store");
        }
        memo[memo_next++] = stack.back();
    };

    auto memo_load = [&](uint32_t index) {
        auto found = memo.find(index);
        if (found == memo.end()) {
            throw std::runtime_error("Malformed pickle: missing memo entry");
        }
        stack.push_back(found->second);
    };

    while (input.peek() != EOF) {
        const char opcode = static_cast<char>(read_u8(input));
        switch (opcode) {
            case '\x80':
                (void) read_u8(input);
                break;
            case '\x95':
                (void) read_u64(input);
                break;
            case '\x94':
                memo_store_auto();
                break;
            case '(':
                stack.push_back(PickleMark{});
                break;
            case 'N':
                stack.push_back(std::monostate{});
                break;
            case 'X': {
                const uint32_t size = read_u32(input);
                std::string value(size, '\0');
                input.read(value.data(), static_cast<std::streamsize>(size));
                stack.push_back(std::move(value));
            } break;
            case '\x8c': {
                const uint8_t size = read_u8(input);
                std::string value(size, '\0');
                input.read(value.data(), static_cast<std::streamsize>(size));
                stack.push_back(std::move(value));
            } break;
            case '\x8d': {
                const uint64_t size = read_u64(input);
                if (size > static_cast<uint64_t>(std::numeric_limits<std::streamsize>::max())) {
                    throw std::runtime_error("BINUNICODE8 payload too large");
                }
                std::string value(static_cast<size_t>(size), '\0');
                input.read(value.data(), static_cast<std::streamsize>(size));
                stack.push_back(std::move(value));
            } break;
            case 'c': {
                std::string module;
                std::string name;
                std::getline(input, module, '\n');
                std::getline(input, name, '\n');
                stack.push_back(PickleGlobal{module, name});
            } break;
            case '\x93': {
                const std::string name = require_string(stack.back(), "STACK_GLOBAL name");
                stack.pop_back();
                const std::string module = require_string(stack.back(), "STACK_GLOBAL module");
                stack.pop_back();
                stack.push_back(PickleGlobal{module, name});
            } break;
            case 'q':
                memo_store(read_u8(input));
                break;
            case 'r':
                memo_store(read_u32(input));
                break;
            case 'h':
                memo_load(read_u8(input));
                break;
            case 'j':
                memo_load(read_u32(input));
                break;
            case ')':
                stack.push_back(make_tuple_from_values({}));
                break;
            case ']':
                stack.push_back(make_tuple_from_values({}));
                break;
            case '}':
                stack.push_back(std::make_shared<PickleDict>());
                break;
            case 'K':
                stack.push_back(static_cast<int64_t>(read_u8(input)));
                break;
            case 'M':
                stack.push_back(static_cast<int64_t>(read_u16(input)));
                break;
            case 'J':
                stack.push_back(static_cast<int64_t>(read_i32(input)));
                break;
            case '\x8a': {
                const uint8_t n = read_u8(input);
                std::vector<uint8_t> bytes(n);
                input.read(reinterpret_cast<char *>(bytes.data()), static_cast<std::streamsize>(n));
                stack.push_back(decode_long(bytes));
            } break;
            case '\x8b': {
                const uint32_t n = read_u32(input);
                std::vector<uint8_t> bytes(n);
                input.read(reinterpret_cast<char *>(bytes.data()), static_cast<std::streamsize>(n));
                stack.push_back(decode_long(bytes));
            } break;
            case 'Q': {
                PickleValue persistent = std::move(stack.back());
                stack.pop_back();
                const std::shared_ptr<PickleTuple> & tuple = require_tuple(persistent, "BINPERSID");
                if (tuple->items.size() != 5) {
                    throw std::runtime_error("Unexpected persistent storage tuple size");
                }

                const std::string marker = require_string(tuple->items[0], "storage marker");
                if (marker != "storage") {
                    throw std::runtime_error("Unsupported persistent id marker: " + marker);
                }

                const PickleGlobal & type_global = require_global(tuple->items[1], "storage type");
                if (type_global.module.rfind("torch", 0) != 0) {
                    throw std::runtime_error("Unsupported storage module: " + type_global.module);
                }

                PickleStorage storage;
                storage.type = storage_type_to_ggml(type_global.name);
                storage.key = require_string(tuple->items[2], "storage key");
                storage.location = require_string(tuple->items[3], "storage location");
                storage.size = static_cast<uint64_t>(require_int(tuple->items[4], "storage size"));
                stack.push_back(storage);
            } break;
            case '\x88':
                stack.push_back(true);
                break;
            case '\x89':
                stack.push_back(false);
                break;
            case 't':
                stack.push_back(make_tuple_from_values(pop_until_mark(stack)));
                break;
            case '\x85': {
                PickleValue item = std::move(stack.back());
                stack.pop_back();
                stack.push_back(make_tuple_from_values({std::move(item)}));
            } break;
            case '\x86': {
                PickleValue b = std::move(stack.back());
                stack.pop_back();
                PickleValue a = std::move(stack.back());
                stack.pop_back();
                stack.push_back(make_tuple_from_values({std::move(a), std::move(b)}));
            } break;
            case '\x87': {
                PickleValue c = std::move(stack.back());
                stack.pop_back();
                PickleValue b = std::move(stack.back());
                stack.pop_back();
                PickleValue a = std::move(stack.back());
                stack.pop_back();
                stack.push_back(make_tuple_from_values({std::move(a), std::move(b), std::move(c)}));
            } break;
            case '\x81': {
                PickleValue args = std::move(stack.back());
                stack.pop_back();
                PickleValue callable = std::move(stack.back());
                stack.pop_back();
                stack.push_back(reduce_value(callable, args));
            } break;
            case 'R': {
                PickleValue args = std::move(stack.back());
                stack.pop_back();
                PickleValue callable = std::move(stack.back());
                stack.pop_back();
                stack.push_back(reduce_value(callable, args));
            } break;
            case 's': {
                PickleValue value = std::move(stack.back());
                stack.pop_back();
                const std::string key = require_string(stack.back(), "SETITEM key");
                stack.pop_back();
                PickleValue dict_value = std::move(stack.back());
                stack.pop_back();
                std::shared_ptr<PickleDict> dict = require_dict(dict_value, "SETITEM dict");
                if (TensorDescriptor * tensor = std::get_if<TensorDescriptor>(&value)) {
                    tensor->name = key;
                }
                dict->items[key] = std::move(value);
                stack.push_back(dict);
            } break;
            case 'u': {
                std::vector<PickleValue> items = pop_until_mark(stack);
                PickleValue dict_value = std::move(stack.back());
                stack.pop_back();
                std::shared_ptr<PickleDict> dict = require_dict(dict_value, "SETITEMS dict");
                for (size_t index = 0; index + 1 < items.size(); index += 2) {
                    const std::string key = require_string(items[index], "SETITEMS key");
                    PickleValue value = std::move(items[index + 1]);
                    if (TensorDescriptor * tensor = std::get_if<TensorDescriptor>(&value)) {
                        tensor->name = key;
                    }
                    dict->items[key] = std::move(value);
                }
                stack.push_back(dict);
            } break;
            case 'e': {
                std::vector<PickleValue> items = pop_until_mark(stack);
                PickleValue list_value = std::move(stack.back());
                stack.pop_back();
                std::shared_ptr<PickleTuple> list = require_tuple(list_value, "APPENDS list");
                for (PickleValue & item : items) {
                    list->items.push_back(std::move(item));
                }
                stack.push_back(list);
            } break;
            case 'b': {
                if (stack.size() < 2) {
                    throw std::runtime_error("Malformed pickle: BUILD requires instance and state");
                }
                PickleValue state = std::move(stack.back());
                stack.pop_back();
                PickleValue instance = std::move(stack.back());
                stack.pop_back();
                (void) state;
                stack.push_back(std::move(instance));
            } break;
            case '.': {
                if (stack.empty()) {
                    throw std::runtime_error("Malformed pickle: empty stack at STOP");
                }
                std::vector<TensorDescriptor> tensors;
                collect_tensors(stack.back(), tensors);
                return tensors;
            }
            default:
                throw std::runtime_error("Unsupported pickle opcode: 0x" +
                                         [&]() {
                                             std::ostringstream out;
                                             out << std::hex << std::uppercase << std::setw(2) << std::setfill('0')
                                                 << static_cast<int>(static_cast<uint8_t>(opcode));
                                             return out.str();
                                         }());
        }
    }

    throw std::runtime_error("Malformed pickle: missing STOP opcode");
}

}  // namespace gd