//! ZigModule encapsulates the state of the incrementally compiled Zig module.
//! It stores the associated input local and global symbols, allocated atoms,
//! and any relocations that may have been emitted.
//! Think about this as fake in-memory Object file for the Zig module.

/// Path is owned by Module and lives as long as *Module.
path: []const u8,
index: File.Index,

local_esyms: std.ArrayListUnmanaged(elf.Elf64_Sym) = .{},
global_esyms: std.ArrayListUnmanaged(elf.Elf64_Sym) = .{},
local_symbols: std.ArrayListUnmanaged(Symbol.Index) = .{},
global_symbols: std.ArrayListUnmanaged(Symbol.Index) = .{},
globals_lookup: std.AutoHashMapUnmanaged(u32, Symbol.Index) = .{},

atoms: std.AutoArrayHashMapUnmanaged(Atom.Index, void) = .{},
relocs: std.ArrayListUnmanaged(std.ArrayListUnmanaged(elf.Elf64_Rela)) = .{},

output_symtab_size: Elf.SymtabSize = .{},

pub fn deinit(self: *ZigModule, allocator: Allocator) void {
    self.local_esyms.deinit(allocator);
    self.global_esyms.deinit(allocator);
    self.local_symbols.deinit(allocator);
    self.global_symbols.deinit(allocator);
    self.globals_lookup.deinit(allocator);
    self.atoms.deinit(allocator);
    for (self.relocs.items) |*list| {
        list.deinit(allocator);
    }
    self.relocs.deinit(allocator);
}

pub fn addLocalEsym(self: *ZigModule, allocator: Allocator) !Symbol.Index {
    try self.local_esyms.ensureUnusedCapacity(allocator, 1);
    const index = @as(Symbol.Index, @intCast(self.local_esyms.items.len));
    const esym = self.local_esyms.addOneAssumeCapacity();
    esym.* = Elf.null_sym;
    esym.st_info = elf.STB_LOCAL << 4;
    return index;
}

pub fn addGlobalEsym(self: *ZigModule, allocator: Allocator) !Symbol.Index {
    try self.global_esyms.ensureUnusedCapacity(allocator, 1);
    const index = @as(Symbol.Index, @intCast(self.global_esyms.items.len));
    const esym = self.global_esyms.addOneAssumeCapacity();
    esym.* = Elf.null_sym;
    esym.st_info = elf.STB_GLOBAL << 4;
    return index | 0x10000000;
}

pub fn addAtom(self: *ZigModule, output_section_index: u16, elf_file: *Elf) !Symbol.Index {
    const gpa = elf_file.base.allocator;

    const atom_index = try elf_file.addAtom();
    try self.atoms.putNoClobber(gpa, atom_index, {});
    const atom_ptr = elf_file.atom(atom_index).?;
    atom_ptr.file_index = self.index;
    atom_ptr.output_section_index = output_section_index;

    const symbol_index = try elf_file.addSymbol();
    try self.local_symbols.append(gpa, symbol_index);
    const symbol_ptr = elf_file.symbol(symbol_index);
    symbol_ptr.file_index = self.index;
    symbol_ptr.atom_index = atom_index;
    symbol_ptr.output_section_index = output_section_index;

    const esym_index = try self.addLocalEsym(gpa);
    const esym = &self.local_esyms.items[esym_index];
    esym.st_shndx = atom_index;
    symbol_ptr.esym_index = esym_index;

    const relocs_index = @as(Atom.Index, @intCast(self.relocs.items.len));
    const relocs = try self.relocs.addOne(gpa);
    relocs.* = .{};
    atom_ptr.relocs_section_index = relocs_index;

    return symbol_index;
}

pub fn resolveSymbols(self: *ZigModule, elf_file: *Elf) void {
    for (self.globals(), 0..) |index, i| {
        const esym_index = @as(Symbol.Index, @intCast(i)) | 0x10000000;
        const esym = self.global_esyms.items[i];

        if (esym.st_shndx == elf.SHN_UNDEF) continue;

        if (esym.st_shndx != elf.SHN_ABS and esym.st_shndx != elf.SHN_COMMON) {
            const atom_index = esym.st_shndx;
            const atom = elf_file.atom(atom_index) orelse continue;
            if (!atom.alive) continue;
        }

        const global = elf_file.symbol(index);
        if (self.asFile().symbolRank(esym, false) < global.symbolRank(elf_file)) {
            const atom_index = switch (esym.st_shndx) {
                elf.SHN_ABS, elf.SHN_COMMON => 0,
                else => esym.st_shndx,
            };
            const output_section_index = if (elf_file.atom(atom_index)) |atom|
                atom.output_section_index
            else
                0;
            global.value = esym.st_value;
            global.atom_index = atom_index;
            global.esym_index = esym_index;
            global.file_index = self.index;
            global.output_section_index = output_section_index;
            global.version_index = elf_file.default_sym_version;
            if (esym.st_bind() == elf.STB_WEAK) global.flags.weak = true;
        }
    }
}

pub fn claimUnresolved(self: *ZigModule, elf_file: *Elf) void {
    for (self.globals(), 0..) |index, i| {
        const esym_index = @as(Symbol.Index, @intCast(i)) | 0x10000000;
        const esym = self.global_esyms.items[i];

        if (esym.st_shndx != elf.SHN_UNDEF) continue;

        const global = elf_file.symbol(index);
        if (global.file(elf_file)) |_| {
            if (global.elfSym(elf_file).st_shndx != elf.SHN_UNDEF) continue;
        }

        const is_import = blk: {
            if (!elf_file.isDynLib()) break :blk false;
            const vis = @as(elf.STV, @enumFromInt(esym.st_other));
            if (vis == .HIDDEN) break :blk false;
            break :blk true;
        };

        global.value = 0;
        global.atom_index = 0;
        global.esym_index = esym_index;
        global.file_index = self.index;
        global.version_index = if (is_import) elf.VER_NDX_LOCAL else elf_file.default_sym_version;
        global.flags.import = is_import;
    }
}

pub fn scanRelocs(self: *ZigModule, elf_file: *Elf, undefs: anytype) !void {
    for (self.atoms.keys()) |atom_index| {
        const atom = elf_file.atom(atom_index) orelse continue;
        if (!atom.alive) continue;
        try atom.scanRelocs(elf_file, undefs);
    }
}

pub fn updateSymtabSize(self: *ZigModule, elf_file: *Elf) void {
    for (self.locals()) |local_index| {
        const local = elf_file.symbol(local_index);
        const esym = local.elfSym(elf_file);
        switch (esym.st_type()) {
            elf.STT_SECTION, elf.STT_NOTYPE => {
                local.flags.output_symtab = false;
                continue;
            },
            else => {},
        }
        local.flags.output_symtab = true;
        self.output_symtab_size.nlocals += 1;
    }

    for (self.globals()) |global_index| {
        const global = elf_file.symbol(global_index);
        if (global.file(elf_file)) |file| if (file.index() != self.index) {
            global.flags.output_symtab = false;
            continue;
        };
        global.flags.output_symtab = true;
        if (global.isLocal()) {
            self.output_symtab_size.nlocals += 1;
        } else {
            self.output_symtab_size.nglobals += 1;
        }
    }
}

pub fn writeSymtab(self: *ZigModule, elf_file: *Elf, ctx: anytype) void {
    var ilocal = ctx.ilocal;
    for (self.locals()) |local_index| {
        const local = elf_file.symbol(local_index);
        if (!local.flags.output_symtab) continue;
        local.setOutputSym(elf_file, &ctx.symtab[ilocal]);
        ilocal += 1;
    }

    var iglobal = ctx.iglobal;
    for (self.globals()) |global_index| {
        const global = elf_file.symbol(global_index);
        if (global.file(elf_file)) |file| if (file.index() != self.index) continue;
        if (!global.flags.output_symtab) continue;
        if (global.isLocal()) {
            global.setOutputSym(elf_file, &ctx.symtab[ilocal]);
            ilocal += 1;
        } else {
            global.setOutputSym(elf_file, &ctx.symtab[iglobal]);
            iglobal += 1;
        }
    }
}

pub fn symbol(self: *ZigModule, index: Symbol.Index) Symbol.Index {
    const is_global = index & 0x10000000 != 0;
    const actual_index = index & 0x0fffffff;
    if (is_global) return self.global_symbols.items[actual_index];
    return self.local_symbols.items[actual_index];
}

pub fn elfSym(self: *ZigModule, index: Symbol.Index) *elf.Elf64_Sym {
    const is_global = index & 0x10000000 != 0;
    const actual_index = index & 0x0fffffff;
    if (is_global) return &self.global_esyms.items[actual_index];
    return &self.local_esyms.items[actual_index];
}

pub fn locals(self: *ZigModule) []const Symbol.Index {
    return self.local_symbols.items;
}

pub fn globals(self: *ZigModule) []const Symbol.Index {
    return self.global_symbols.items;
}

pub fn asFile(self: *ZigModule) File {
    return .{ .zig_module = self };
}

pub fn fmtSymtab(self: *ZigModule, elf_file: *Elf) std.fmt.Formatter(formatSymtab) {
    return .{ .data = .{
        .self = self,
        .elf_file = elf_file,
    } };
}

const FormatContext = struct {
    self: *ZigModule,
    elf_file: *Elf,
};

fn formatSymtab(
    ctx: FormatContext,
    comptime unused_fmt_string: []const u8,
    options: std.fmt.FormatOptions,
    writer: anytype,
) !void {
    _ = unused_fmt_string;
    _ = options;
    try writer.writeAll("  locals\n");
    for (ctx.self.locals()) |index| {
        const local = ctx.elf_file.symbol(index);
        try writer.print("    {}\n", .{local.fmt(ctx.elf_file)});
    }
    try writer.writeAll("  globals\n");
    for (ctx.self.globals()) |index| {
        const global = ctx.elf_file.symbol(index);
        try writer.print("    {}\n", .{global.fmt(ctx.elf_file)});
    }
}

pub fn fmtAtoms(self: *ZigModule, elf_file: *Elf) std.fmt.Formatter(formatAtoms) {
    return .{ .data = .{
        .self = self,
        .elf_file = elf_file,
    } };
}

fn formatAtoms(
    ctx: FormatContext,
    comptime unused_fmt_string: []const u8,
    options: std.fmt.FormatOptions,
    writer: anytype,
) !void {
    _ = unused_fmt_string;
    _ = options;
    try writer.writeAll("  atoms\n");
    for (ctx.self.atoms.keys()) |atom_index| {
        const atom = ctx.elf_file.atom(atom_index) orelse continue;
        try writer.print("    {}\n", .{atom.fmt(ctx.elf_file)});
    }
}

const assert = std.debug.assert;
const std = @import("std");
const elf = std.elf;

const Allocator = std.mem.Allocator;
const Atom = @import("Atom.zig");
const Elf = @import("../Elf.zig");
const File = @import("file.zig").File;
const Module = @import("../../Module.zig");
const Symbol = @import("Symbol.zig");
const ZigModule = @This();