//! Patch's predicate synthesizer
// 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 ArenaAllocator = std.heap.ArenaAllocator;
const CompareOperator = std.math.CompareOperator;
const Writer = std.io.Writer;
const argsAlloc = std.process.argsAlloc;
const argsFree = std.process.argsFree;
const assert = std.debug.assert;
const compare = std.math.compare;
const divTrunc = std.math.divTrunc;
const exit = std.process.exit;
const page_allocator = std.heap.page_allocator;
const parseUnsigned = std.fmt.parseUnsigned;
const print = std.debug.print;
const rem = std.math.rem;
const shl = std.math.shl;
const shr = std.math.shr;
const std = @import("std");
const stdout = std.fs.File.stdout;
const tags = std.meta.tags;

const Register = Variables.Register;
const RegisterEnum = Variables.RegisterEnum;
const Variables = @import("Variables.zig");
const signed_integers = Variables.signed_integers;

comptime {
    _ = Variables; // for test discovery
}

const UnaryOperator = enum {
    @"+",
    @"-",
    @"~",

    fn apply(op: UnaryOperator, val: Register) Register {
        return switch (op) {
            .@"+" => val,
            .@"-" => -%val,
            .@"~" => ~val,
        };
    }
};

const Unary = struct {
    op: UnaryOperator,
    val: Variables.Query,

    pub fn format(un: Unary, writer: *Writer) Writer.Error!void {
        try writer.print("{s}{f}", .{ @tagName(un.op), un.val });
    }
};

const ArithmeticOperator = enum {
    @"+", @"-", @"*", @"/", @"%",
    @"&", @"|", @"^", @"<<", @">>",

    fn apply(op: ArithmeticOperator, lhs: Register, rhs: Register) !Register {
        return switch (op) {
            .@"+" => lhs +% rhs,
            .@"-" => lhs -% rhs,
            .@"*" => lhs *% rhs,
            .@"/" => try divTrunc(Register, lhs, rhs),
            .@"%" => try rem(Register, lhs, rhs),
            .@"&" => lhs & rhs,
            .@"|" => lhs | rhs,
            .@"^" => lhs ^ rhs,
            .@"<<" => shl(Register, lhs, rhs),
            .@">>" => shr(Register, lhs, rhs),
        };
    }
};

const Arithmetic = struct {
    lhs: Variables.Query,
    op: ArithmeticOperator,
    rhs: Variables.Query,

    pub fn format(arith: Arithmetic, writer: *Writer) Writer.Error!void {
        try writer.print("{f} {s} {f}", .{
            arith.lhs,
            @tagName(arith.op),
            arith.rhs,
        });
    }
};

const Comparison = struct {
    lhs: Arithmetic,
    op: CompareOperator,
    rhs: Unary,

    fn check(cmp: Comparison, vars: Variables, expected: bool,
             tmp: Temporaries) !bool {
        for (try vars.get(cmp.lhs.lhs, tmp[0]),
             try vars.get(cmp.lhs.rhs, tmp[1]),
             try vars.get(cmp.rhs.val, tmp[2])) |v0, v1, v2| {
            const arith = cmp.lhs.op.apply(v0, v1) catch return false;
            const actual = compare(arith, cmp.op, cmp.rhs.op.apply(v2));
            if (if (expected) !actual else actual)
                return false;
        }
        return true;
    }

    pub fn format(cmp: Comparison, writer: *Writer) Writer.Error!void {
        try writer.print("{f} {s} {f}", .{
            cmp.lhs,
            switch (cmp.op) {
                .lt => "<",
                .lte => "<=",
                .eq => "==",
                .gte => ">=",
                .gt => ">",
                .neq => "!=",
            },
            cmp.rhs,
        });
    }
};

const Temporaries = [3][]Register;

fn process(writer: *Writer, bot: Variables, top: Variables,
           bot_tmp: Temporaries, top_tmp: Temporaries,
           v0: Variables.Query, arith: ArithmeticOperator,
           v1: Variables.Query, cmp: CompareOperator,
           unary: UnaryOperator, v2: Variables.Query) !void {
    const predicate = Comparison{
        .lhs = .{ .lhs = v0, .op = arith, .rhs = v1 },
        .op = cmp,
        .rhs = .{ .op = unary, .val = v2 },
    };
    if (try predicate.check(bot, false, bot_tmp))
        return;
    if (try predicate.check(top, true, top_tmp))
        return;
    try writer.print("{f}\n", .{ predicate });
}

fn mineOperator(writer: *Writer, bot: Variables, top: Variables,
                bot_tmp: Temporaries, top_tmp: Temporaries,
                v0: Variables.Query, v1: Variables.Query,
                v2: Variables.Query) !void {
    for (tags(ArithmeticOperator)) |arith|
        for (tags(UnaryOperator)) |unary|
            for (tags(CompareOperator)) |cmp|
                try process(writer, bot, top, bot_tmp, top_tmp,
                            v0, arith, v1, cmp, unary, v2);
}

fn alloc_tmp(allocator: Allocator, len: usize) Allocator.Error!Temporaries {
    var tmp = [_][]Register{ &.{} } ** 3;
    errdefer free_tmp(allocator, tmp);
    for (&tmp) |*variable|
        variable.* = try allocator.alloc(Register, len);
    return tmp;
}

fn free_tmp(allocator: Allocator, tmp: Temporaries) void {
    for (tmp) |variable|
        allocator.free(variable);
}

fn mine(allocator: Allocator, writer: *Writer,
        bot: Variables, top: Variables) !void {
    const bot_tmp = try alloc_tmp(allocator, bot.samples);
    defer free_tmp(allocator, bot_tmp);
    const top_tmp = try alloc_tmp(allocator, top.samples);
    defer free_tmp(allocator, top_tmp);

    const variables = try Variables.allQueries(allocator);
    defer allocator.free(variables);
    inline for (comptime Variables.constantQueries()) |v2|
        for (variables) |v0| for (variables) |v1| if (!v2.skip(v0, v1))
            try mineOperator(writer, bot, top, bot_tmp, top_tmp, v0, v1, v2);
    try writer.flush();
}

pub fn main() !void {
    var buffer: [80]u8 = undefined;
    var writer = stdout().writer(&buffer);
    var arena = ArenaAllocator.init(page_allocator);
    defer arena.deinit();
    const allocator = arena.allocator();
    const args = try argsAlloc(allocator);
    defer argsFree(allocator, args);
    if (args.len != 4) {
        try writer.interface.print("Usage: taosc-synth {s} {s} {s}\n", .{
            "STACK_SIZE", "BOTTOM_DIR", "TOP_DIR",
        });
        try writer.interface.flush();
        exit(1);
    }
    const stack_size = try parseUnsigned(usize, args[1], 0);
    inline for (signed_integers) |SignedInt|
        assert(stack_size % @sizeOf(SignedInt) == 0);
    const bot = try Variables.init(allocator, stack_size, args[2]);
    defer bot.deinit(allocator);
    const top = try Variables.init(allocator, stack_size, args[3]);
    defer top.deinit(allocator);
    try mine(allocator, &writer.interface, bot, top);
}