tests.cpp

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>

#include "load_gguf.hpp"
#include "ops.hpp"
#include "reader.hpp"
#include "tensor.hpp"

struct TestFailure : std::runtime_error {
    using std::runtime_error::runtime_error;
};

struct TestSkip : std::runtime_error {
    using std::runtime_error::runtime_error;
};

static void check(bool cond, const std::string& msg) {
    if(!cond) {
        throw TestFailure(msg);
    }
}

static void skip_test(const std::string& msg) {
    throw TestSkip(msg);
}

template<typename T>
static void check_eq(const T& actual, const T& expected, const std::string& msg) {
    if(!(actual == expected)) {
        throw TestFailure(msg);
    }
}

static void check_near(float actual, float expected, float tol, const std::string& msg) {
    if(std::fabs(actual - expected) > tol) {
        throw TestFailure(
            msg + " actual=" + std::to_string(actual) + " expected=" + std::to_string(expected)
        );
    }
}

static void check_vec_eq(
    const std::vector<int64_t>& actual,
    const std::vector<int64_t>& expected,
    const std::string& msg
) {
    if(actual != expected) {
        throw TestFailure(msg);
    }
}

static void check_vec_near(
    const std::vector<float>& actual,
    const std::vector<float>& expected,
    float tol,
    const std::string& msg
) {
    check_eq(actual.size(), expected.size(), msg + " size");
    for(size_t i = 0; i < actual.size(); i++) {
        if(std::fabs(actual[i] - expected[i]) > tol) {
            throw TestFailure(
                msg + " idx=" + std::to_string(i) + " actual=" + std::to_string(actual[i]) +
                " expected=" + std::to_string(expected[i])
            );
        }
    }
}

static void fill_tensor(Tensor& t, const std::vector<float>& values) {
    check_eq(static_cast<size_t>(t.numel()), values.size(), "fill_tensor size mismatch");
    std::copy(values.begin(), values.end(), t.data_as<float>());
}

static std::vector<float> tensor_to_vec(const Tensor& t) {
    const float* p = t.data_as<float>();
    return std::vector<float>(p, p + t.numel());
}

template<typename T>
static void append_bytes(std::string& out, const T& value) {
    out.append(reinterpret_cast<const char*>(&value), sizeof(T));
}

static void append_gguf_string(std::string& out, const std::string& s) {
    append_bytes<uint64_t>(out, static_cast<uint64_t>(s.size()));
    out.append(s.data(), static_cast<std::streamsize>(s.size()));
}

static TensorInfo make_tensor_info_fixture(
    const std::string& name,
    const std::vector<uint64_t>& dims,
    ggml_type type,
    uint64_t offset
) {
    std::string bytes;
    append_gguf_string(bytes, name);
    append_bytes<uint32_t>(bytes, static_cast<uint32_t>(dims.size()));
    for(uint64_t d : dims) {
        append_bytes<uint64_t>(bytes, d);
    }
    append_bytes<uint32_t>(bytes, static_cast<uint32_t>(type));
    append_bytes<uint64_t>(bytes, offset);

    std::istringstream stream(bytes);
    Reader r(stream);
    return TensorInfo(r);
}

static uint32_t gguf_magic_le() {
    return static_cast<uint32_t>('G')
        | (static_cast<uint32_t>('G') << 8)
        | (static_cast<uint32_t>('U') << 16)
        | (static_cast<uint32_t>('F') << 24);
}

static std::vector<float> ref_matmul_2d(
    const std::vector<float>& a,
    const std::vector<float>& b,
    int64_t M,
    int64_t K,
    int64_t N
) {
    std::vector<float> out(static_cast<size_t>(M * N), 0.0f);
    for(int64_t m = 0; m < M; m++) {
        for(int64_t n = 0; n < N; n++) {
            float sum = 0.0f;
            for(int64_t k = 0; k < K; k++) {
                sum += a[static_cast<size_t>(m * K + k)] * b[static_cast<size_t>(k * N + n)];
            }
            out[static_cast<size_t>(m * N + n)] = sum;
        }
    }
    return out;
}

static void test_tensor_basic() {
    Tensor t({2, 3, 4});
    check_vec_eq(t.shape(), {2, 3, 4}, "tensor shape");
    check_vec_eq(t.strides(), {12, 4, 1}, "tensor strides");
    check_eq(t.numel(), static_cast<int64_t>(24), "tensor numel");
    check_eq(t.nbytes(), static_cast<size_t>(96), "tensor nbytes");

    for(int64_t i = 0; i < t.numel(); i++) {
        t.data_as<float>()[i] = static_cast<float>(i);
    }

    check_eq(t.offset({1, 2, 3}), static_cast<int64_t>(23), "tensor offset");
    check_near(t.at<float>({1, 2, 3}), 23.0f, 1e-6f, "tensor at");

    Tensor v = t.view({4, 6});
    check(v.data() == t.data(), "view should share storage");
    check_eq(v.numel(), static_cast<int64_t>(24), "view numel");
    check_near(v.at<float>({3, 5}), 23.0f, 1e-6f, "view indexing");
}

static void test_broadcast_helpers() {
    check_vec_eq(broadcast_shapes({2, 3}, {3}), {2, 3}, "broadcast shape vec");
    check_vec_eq(broadcast_shapes({1, 3}, {2, 3}), {2, 3}, "broadcast shape leading one");

    bool threw = false;
    try {
        (void)broadcast_shapes({2, 3}, {4, 3});
    } catch(const std::runtime_error&) {
        threw = true;
    }
    check(threw, "broadcast_shapes should throw on incompatible shapes");

    check_vec_eq(broadcast_strides({3}, {1}, {2, 3}), {0, 1}, "broadcast strides rank expand");
    check_vec_eq(
        broadcast_strides({2, 1}, Tensor::compute_strides({2, 1}), {2, 3}),
        {1, 0},
        "broadcast strides singleton dim"
    );
}

static void test_add_ops() {
    Tensor a({2, 3});
    Tensor b({2, 3});
    Tensor out({2, 3});
    fill_tensor(a, {1, 2, 3, 4, 5, 6});
    fill_tensor(b, {10, 20, 30, 40, 50, 60});
    add(out, a, b);
    check_vec_near(tensor_to_vec(out), {11, 22, 33, 44, 55, 66}, 1e-6f, "add same shape");

    Tensor scalar({1});
    Tensor t({2, 3});
    Tensor out2({2, 3});
    fill_tensor(scalar, {0.5f});
    fill_tensor(t, {1, 2, 3, 4, 5, 6});
    add(out2, scalar, t);
    check_vec_near(tensor_to_vec(out2), {1.5f, 2.5f, 3.5f, 4.5f, 5.5f, 6.5f}, 1e-6f, "add scalar");

    Tensor c({2, 3});
    Tensor d({2, 1});
    Tensor out3({2, 3});
    fill_tensor(c, {1, 2, 3, 10, 20, 30});
    fill_tensor(d, {100, 200});
    add(out3, c, d);
    check_vec_near(tensor_to_vec(out3), {101, 102, 103, 210, 220, 230}, 1e-6f, "add broadcast");
}

static void test_matmul_ops() {
    Tensor a({2, 3});
    Tensor b({3, 2});
    Tensor out({2, 2});
    fill_tensor(a, {1, 2, 3, 4, 5, 6});
    fill_tensor(b, {7, 8, 9, 10, 11, 12});
    matmul(out, a, b);
    check_vec_near(tensor_to_vec(out), {58, 64, 139, 154}, 1e-6f, "matmul 2d");

    Tensor ba({2, 2, 3});
    Tensor bb({3, 2});
    Tensor bout({2, 2, 2});
    fill_tensor(ba, {
        1, 2, 3,
        4, 5, 6,
        10, 20, 30,
        40, 50, 60
    });
    fill_tensor(bb, {1, 2, 3, 4, 5, 6});
    matmul(bout, ba, bb);

    std::vector<float> expected;
    auto avals = tensor_to_vec(ba);
    auto bvals = tensor_to_vec(bb);
    auto e0 = ref_matmul_2d(std::vector<float>(avals.begin(), avals.begin() + 6), bvals, 2, 3, 2);
    auto e1 = ref_matmul_2d(std::vector<float>(avals.begin() + 6, avals.end()), bvals, 2, 3, 2);
    expected.insert(expected.end(), e0.begin(), e0.end());
    expected.insert(expected.end(), e1.begin(), e1.end());
    check_vec_near(tensor_to_vec(bout), expected, 1e-6f, "matmul batched");

    Tensor bad_a({2, 3});
    Tensor bad_b({4, 2});
    Tensor bad_out({2, 2});
    bool threw = false;
    try {
        matmul(bad_out, bad_a, bad_b);
    } catch(const std::runtime_error&) {
        threw = true;
    }
    check(threw, "matmul should throw on inner dim mismatch");
}

static void test_reader_primitives() {
    std::string bytes;
    append_bytes<uint32_t>(bytes, 0x78563412u);
    append_bytes<int16_t>(bytes, static_cast<int16_t>(-7));
    float f = 3.5f;
    append_bytes<float>(bytes, f);
    append_bytes<uint8_t>(bytes, 1);
    append_bytes<uint64_t>(bytes, 2);
    bytes.append("hi", 2);

    std::istringstream stream(bytes);
    Reader r(stream);

    check_eq(r.read_u32(), static_cast<uint32_t>(0x78563412u), "reader u32");
    check_eq(r.read_i16(), static_cast<int16_t>(-7), "reader i16");
    check_near(r.read_f32(), 3.5f, 1e-6f, "reader f32");
    check(r.read_bool(), "reader bool");
    check_eq(r.read_string(), std::string("hi"), "reader string");
}

static void test_reader_truncated() {
    std::istringstream stream(std::string(3, '\0'));
    Reader r(stream);

    bool threw = false;
    try {
        (void)r.read_u32();
    } catch(const std::runtime_error&) {
        threw = true;
    }
    check(threw, "truncated read should throw");
}

static void test_fp16_conversion() {
    check_near(fp16_to_fp32(0x3C00), 1.0f, 1e-6f, "fp16 1.0");
    check_near(fp16_to_fp32(0xC000), -2.0f, 1e-6f, "fp16 -2.0");
    check_near(fp16_to_fp32(0x0000), 0.0f, 1e-6f, "fp16 zero");
    check(std::isinf(fp16_to_fp32(0x7C00)), "fp16 inf");
}

static void test_ggml_dims_to_shape() {
    check_vec_eq(ggml_dims_to_shape({3, 2}), {2, 3}, "ggml dims reverse to tensor shape");
    check_vec_eq(ggml_dims_to_shape({4, 3, 2}), {2, 3, 4}, "ggml dims reverse rank3");
}

static void test_load_tensor_f32_from_f32() {
    TensorInfo ti = make_tensor_info_fixture("w_f32", {3, 2}, GGML_TYPE_F32, 8);

    const uint64_t tensor_data_start = 16;
    const uint64_t abs = tensor_data_start + ti.offset;

    std::string bytes(static_cast<size_t>(abs), '\0');
    for(float v : std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}) {
        append_bytes<float>(bytes, v);
    }

    std::istringstream stream(bytes);
    Reader r(stream);
    Tensor t = load_tensor_f32(r, ti, tensor_data_start);

    check_vec_eq(t.shape(), {2, 3}, "load_tensor_f32 shape from f32");
    check_vec_near(tensor_to_vec(t), {1, 2, 3, 4, 5, 6}, 1e-6f, "load_tensor_f32 values f32");
}

static void test_load_tensor_f32_from_f16() {
    TensorInfo ti = make_tensor_info_fixture("w_f16", {2, 2}, GGML_TYPE_F16, 4);

    const uint64_t tensor_data_start = 20;
    const uint64_t abs = tensor_data_start + ti.offset;

    std::string bytes(static_cast<size_t>(abs), '\0');
    for(uint16_t h : std::vector<uint16_t>{0x3C00, 0xC000, 0x3800, 0x0000}) { // 1, -2, 0.5, 0
        append_bytes<uint16_t>(bytes, h);
    }

    std::istringstream stream(bytes);
    Reader r(stream);
    Tensor t = load_tensor_f32(r, ti, tensor_data_start);

    check_vec_eq(t.shape(), {2, 2}, "load_tensor_f32 shape from f16");
    check_vec_near(tensor_to_vec(t), {1.0f, -2.0f, 0.5f, 0.0f}, 1e-6f, "load_tensor_f32 values f16");
}

static void test_load_tensors_multiple() {
    std::vector<TensorInfo> infos;
    infos.push_back(make_tensor_info_fixture("a", {2}, GGML_TYPE_F32, 0));
    infos.push_back(make_tensor_info_fixture("b", {2, 1}, GGML_TYPE_F16, 8));

    const uint64_t tensor_data_start = 32;
    std::string bytes(static_cast<size_t>(tensor_data_start), '\0');

    // tensor 0 at offset 0: two f32s
    append_bytes<float>(bytes, 10.0f);
    append_bytes<float>(bytes, 20.0f);

    // tensor 1 at offset 8: two f16s (shape [1,2] after dim reversal)
    append_bytes<uint16_t>(bytes, 0x4000); // 2.0
    append_bytes<uint16_t>(bytes, 0x4200); // 3.0

    std::istringstream stream(bytes);
    Reader r(stream);
    auto tensors = load_tensors(r, infos, tensor_data_start);

    check_eq(tensors.size(), static_cast<size_t>(2), "load_tensors size");
    check_vec_eq(tensors[0].shape(), {2}, "load_tensors tensor0 shape");
    check_vec_near(tensor_to_vec(tensors[0]), {10.0f, 20.0f}, 1e-6f, "load_tensors tensor0 values");
    check_vec_eq(tensors[1].shape(), {1, 2}, "load_tensors tensor1 shape");
    check_vec_near(tensor_to_vec(tensors[1]), {2.0f, 3.0f}, 1e-6f, "load_tensors tensor1 values");
}

static void test_gguf_minimal_fixture() {
    std::string bytes;
    append_bytes<uint32_t>(bytes, gguf_magic_le());
    append_bytes<uint32_t>(bytes, 3u);
    append_bytes<uint64_t>(bytes, 0u);
    append_bytes<uint64_t>(bytes, 0u);

    std::istringstream stream(bytes);
    Reader r(stream);
    GGUFFile gguf(r);

    check_eq(gguf.magic, gguf_magic_le(), "fixture gguf magic");
    check_eq(gguf.version, 3u, "fixture gguf version");
    check_eq(gguf.tensor_count, static_cast<uint64_t>(0), "fixture tensor_count");
    check_eq(gguf.metadata_kv_count, static_cast<uint64_t>(0), "fixture metadata_kv_count");
    check(gguf.metadata.empty(), "fixture metadata empty");
    check(gguf.tensor_infos.empty(), "fixture tensor infos empty");
}

static void test_gguf_header_real_file() {
    namespace fs = std::filesystem;
    const fs::path path = "qwen2.5-0.5b-instruct-fp16.gguf";
    if(!fs::exists(path)) {
        skip_test("model file not found");
    }

    std::ifstream file(path, std::ios::binary);
    check(static_cast<bool>(file), "failed to open gguf file");
    Reader r(file);
    const uint32_t magic = r.read_u32();
    const uint32_t version = r.read_u32();
    const uint64_t tensor_count = r.read_u64();
    const uint64_t metadata_kv_count = r.read_u64();

    check_eq(magic, gguf_magic_le(), "gguf magic");
    check(version >= 2, "gguf version");
    check(tensor_count > 0, "gguf tensor count");
    check(metadata_kv_count > 0, "gguf metadata count");
}

int main() {
    struct TestCase {
        const char* name;
        void (*fn)();
    };

    const std::vector<TestCase> tests = {
        {"tensor_basic", test_tensor_basic},
        {"broadcast_helpers", test_broadcast_helpers},
        {"add_ops", test_add_ops},
        {"matmul_ops", test_matmul_ops},
        {"reader_primitives", test_reader_primitives},
        {"reader_truncated", test_reader_truncated},
        {"fp16_conversion", test_fp16_conversion},
        {"ggml_dims_to_shape", test_ggml_dims_to_shape},
        {"load_tensor_f32_from_f32", test_load_tensor_f32_from_f32},
        {"load_tensor_f32_from_f16", test_load_tensor_f32_from_f16},
        {"load_tensors_multiple", test_load_tensors_multiple},
        {"gguf_minimal_fixture", test_gguf_minimal_fixture},
        {"gguf_header_real_file", test_gguf_header_real_file},
    };

    int passed = 0;
    int skipped = 0;
    int failed = 0;

    for(const auto& test : tests) {
        try {
            test.fn();
            passed++;
            std::cout << "[PASS] " << test.name << "\n";
        } catch(const TestSkip& e) {
            skipped++;
            std::cout << "[SKIP] " << test.name << " - " << e.what() << "\n";
        } catch(const std::exception& e) {
            failed++;
            std::cout << "[FAIL] " << test.name << " - " << e.what() << "\n";
        } catch(...) {
            failed++;
            std::cout << "[FAIL] " << test.name << " - unknown exception\n";
        }
    }

    std::cout << "Summary: " << passed << " passed, "
              << skipped << " skipped, "
              << failed << " failed\n";

    return failed == 0 ? 0 : 1;
}