//! Variables captured at patch location
// Copyright (C) 2025  Nguyễn Gia Phong
//
// This file is part of taosc.
//
// Taosc is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Taosc is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with taosc.  If not, see <https://www.gnu.org/licenses/>.

const Allocator = std.mem.Allocator;
const Writer = std.io.Writer;
const assert = std.debug.assert;
const bytesAsSlice = std.mem.bytesAsSlice;
const comptimePrint = std.fmt.comptimePrint;
const cwd = std.fs.cwd;
const divCeil = std.math.divCeil;
const eql = std.meta.eql;
const expect = std.testing.expect;
const maxInt = std.math.maxInt;
const minInt = std.math.minInt;
const std = @import("std");
const tags = std.meta.tags;

pub const RegisterEnum = enum(u4) {
    rax, rbx, rcx, rdx, rsi, rdi, rsp, rbp,
    r8, r9, r10, r11, r12, r13, r14, r15,
};
pub const Register = i64;
const Registers = [tags(RegisterEnum).len]Register;

pub const signed_integers = .{ i64, i32, i16, i8 };
const alignment = @alignOf(Register);

test alignment {
    inline for (signed_integers) |Int|
        try expect(alignment >= @alignOf(Int));
}

fn alignedSize(T: type, count: usize) !usize {
    return try divCeil(usize, @sizeOf(T) * count, @alignOf(T)) * @alignOf(T);
}

fn packedSize(T: type, container_size: usize, count: usize) !usize {
    return @divExact(container_size, @sizeOf(T)) * try alignedSize(T, count);
}

const Variables = @This();
bytes: []align(alignment) const u8,
samples: usize,

pub fn init(allocator: Allocator, stack_size: usize,
            path: []const u8) !Variables {
    var dir = try cwd().openDir(path, .{ .iterate = true });
    defer dir.close();
    var entries = dir.iterate();
    var count: usize = 0;
    while (try entries.next()) |entry| {
        assert(entry.kind == .file);
        count += 1;
    }
    var size = try packedSize(Register, @sizeOf(Registers), count);
    inline for (signed_integers) |Int|
        size += try packedSize(Int, stack_size, count);
    const bytes = try allocator.alignedAlloc(u8,
        .fromByteUnits(alignment), size);
    errdefer allocator.free(bytes);

    const file_size = @sizeOf(Registers) + stack_size;
    const buffer = try allocator.alignedAlloc(u8,
        .fromByteUnits(@alignOf(Registers)), file_size);
    defer allocator.free(buffer);
    const registers: *Registers = @ptrCast(buffer.ptr);
    entries.reset();
    for (0..count) |index| {
        const entry = try entries.next();
        const file = try dir.openFile(entry.?.name, .{});
        defer file.close();
        assert(try file.getEndPos() == file_size);
        assert(try file.read(buffer) == file_size);
        var offset: usize = 0;
        inline for (registers) |register| {
            const slice = bytesAsSlice(Register, bytes[offset..]);
            slice[index] = register;
            offset += try alignedSize(Register, count);
        }
        inline for (signed_integers) |Int| {
            const aligned_size = try alignedSize(Int, count);
            for (bytesAsSlice(Int, buffer[@sizeOf(Registers)..])) |value| {
                const slice = bytesAsSlice(Int, bytes[offset..]);
                slice[index] = value;
                offset += aligned_size;
            }
        }
    }
    return .{ .bytes = bytes, .samples = count };
}

pub fn deinit(variables: Variables, allocator: Allocator) void {
    allocator.free(variables.bytes);
}

const ConstantEnum = enum(i64) {
    max1 = maxInt(i1),
    min2 = minInt(i2),
    max2 = maxInt(i2),
    min3 = minInt(i3),
    max3 = maxInt(i3),
    min4 = minInt(i4),
    max4 = maxInt(i4),
    min5 = minInt(i5),
    max5 = maxInt(i5),
    min6 = minInt(i6),
    max6 = maxInt(i6),
    min7 = minInt(i7),
    max7 = maxInt(i7),
    min8 = minInt(i8),
    max8 = maxInt(i8),
    min9 = minInt(i9),
    min16 = minInt(i16),
    max16 = maxInt(i16),
    min17 = minInt(i17),
    min32 = minInt(i32),
    max32 = maxInt(i32),
    min33 = minInt(i33),
    min64 = minInt(i64),
    max64 = maxInt(i64),
};

pub const Query = union(enum) {
    constant: ConstantEnum,
    register: RegisterEnum,
    stack: usize,

    pub fn all(allocator: Allocator) Allocator.Error![]Query {
        const constants = tags(ConstantEnum);
        const registers = tags(RegisterEnum);
        const n = constants.len + registers.len;
        const queries = try allocator.alloc(Query, n);
        for (queries[0..constants.len], constants) |*dest, src|
            dest.* = .{ .constant = src };
        for (queries[constants.len..][0..registers.len],
             registers) |*dest, src|
            dest.* = .{ .register = src };
        return queries;
    }

    pub fn skip(q: Query, r: Query, s: Query) bool {
        return eql(q, r) or eql(r, s) or eql(s, q)
            or q == .constant and r == .constant and s == .constant
            or q != .constant and r != .constant and s != .constant;
    }

    pub fn format(query: Query, writer: *Writer) Writer.Error!void {
        switch (query) {
            inline .constant, .register => |tag|
                try writer.print("{s}", .{ @tagName(tag) }),
            .stack => unreachable,
        }
    }
};

pub fn get(variables: Variables, query: Query,
           tmp: []Register) ![]const Register {
    switch (query) {
        .constant => |constant| switch (constant) {
            inline else => |tag| {
                for (tmp) |*register|
                    register.* = @intFromEnum(tag);
                return tmp;
            },
        },
        .register => |tag| {
            const aligned_size = try alignedSize(Register, variables.samples);
            const offset = aligned_size * @intFromEnum(tag);
            const slice = bytesAsSlice(Register, variables.bytes[offset..]);
            return @alignCast(slice[0..variables.samples]);
        },
        .stack => unreachable,
    }
}