coff: resolve relocs on bytes buffer directly

2 files changed, 51 insertions, 59 deletions

diff --git a/src/link/Coff.zig b/src/link/Coff.zig
index df525bba90..7d853b3181 100644
--- a/src/link/Coff.zig
+++ b/src/link/Coff.zig
@@ -771,7 +771,7 @@ fn shrinkAtom(self: *Coff, atom_index: Atom.Index, new_block_size: u32) void {
     // capacity, insert a free list node for it.
 }
 
-fn writeAtom(self: *Coff, atom_index: Atom.Index, code: []const u8) !void {
+fn writeAtom(self: *Coff, atom_index: Atom.Index, code: []u8) !void {
     const atom = self.getAtom(atom_index);
     const sym = atom.getSymbol(self);
     const section = self.sections.get(@enumToInt(sym.section_number) - 1);
@@ -781,8 +781,8 @@ fn writeAtom(self: *Coff, atom_index: Atom.Index, code: []const u8) !void {
         file_offset,
         file_offset + code.len,
     });
+    self.resolveRelocs(atom_index, code);
     try self.base.file.?.pwriteAll(code, file_offset);
-    try self.resolveRelocs(atom_index);
 }
 
 fn writePtrWidthAtom(self: *Coff, atom_index: Atom.Index) !void {
@@ -820,14 +820,15 @@ fn markRelocsDirtyByAddress(self: *Coff, addr: u32) void {
     }
 }
 
-fn resolveRelocs(self: *Coff, atom_index: Atom.Index) !void {
+fn resolveRelocs(self: *Coff, atom_index: Atom.Index, code: []u8) void {
     const relocs = self.relocs.get(atom_index) orelse return;
 
     log.debug("relocating '{s}'", .{self.getAtom(atom_index).getName(self)});
 
     for (relocs.items) |*reloc| {
         if (!reloc.dirty) continue;
-        try reloc.resolve(atom_index, self);
+        reloc.resolve(atom_index, code, self);
+        reloc.dirty = false;
     }
 }
 
@@ -944,7 +945,7 @@ pub fn updateFunc(self: *Coff, module: *Module, func: *Module.Fn, air: Air, live
         &code_buffer,
         .none,
     );
-    const code = switch (res) {
+    var code = switch (res) {
         .ok => code_buffer.items,
         .fail => |em| {
             decl.analysis = .codegen_failure;
@@ -994,7 +995,7 @@ pub fn lowerUnnamedConst(self: *Coff, tv: TypedValue, decl_index: Module.Decl.In
     const res = try codegen.generateSymbol(&self.base, decl.srcLoc(), tv, &code_buffer, .none, .{
         .parent_atom_index = self.getAtom(atom_index).getSymbolIndex().?,
     });
-    const code = switch (res) {
+    var code = switch (res) {
         .ok => code_buffer.items,
         .fail => |em| {
             decl.analysis = .codegen_failure;
@@ -1057,7 +1058,7 @@ pub fn updateDecl(self: *Coff, module: *Module, decl_index: Module.Decl.Index) !
     }, &code_buffer, .none, .{
         .parent_atom_index = atom.getSymbolIndex().?,
     });
-    const code = switch (res) {
+    var code = switch (res) {
         .ok => code_buffer.items,
         .fail => |em| {
             decl.analysis = .codegen_failure;
@@ -1110,7 +1111,7 @@ fn getDeclOutputSection(self: *Coff, decl_index: Module.Decl.Index) u16 {
     return index;
 }
 
-fn updateDeclCode(self: *Coff, decl_index: Module.Decl.Index, code: []const u8, complex_type: coff.ComplexType) !void {
+fn updateDeclCode(self: *Coff, decl_index: Module.Decl.Index, code: []u8, complex_type: coff.ComplexType) !void {
     const gpa = self.base.allocator;
     const mod = self.base.options.module.?;
     const decl = mod.declPtr(decl_index);
@@ -1424,8 +1425,28 @@ pub fn flushModule(self: *Coff, comp: *Compilation, prog_node: *std.Progress.Nod
     try self.writeImportTables();
     {
         var it = self.relocs.keyIterator();
-        while (it.next()) |atom| {
-            try self.resolveRelocs(atom.*);
+        while (it.next()) |atom_index_ptr| {
+            const atom_index = atom_index_ptr.*;
+            const relocs = self.relocs.get(atom_index).?;
+            const needs_update = for (relocs.items) |reloc| {
+                if (reloc.dirty) break true;
+            } else false;
+
+            if (!needs_update) continue;
+
+            const atom = self.getAtom(atom_index);
+            const sym = atom.getSymbol(self);
+            const section = self.sections.get(@enumToInt(sym.section_number) - 1).header;
+            const file_offset = section.pointer_to_raw_data + sym.value - section.virtual_address;
+
+            var code = std.ArrayList(u8).init(gpa);
+            defer code.deinit();
+            try code.resize(math.cast(usize, atom.size) orelse return error.Overflow);
+
+            const amt = try self.base.file.?.preadAll(code.items, file_offset);
+            if (amt != code.items.len) return error.InputOutput;
+
+            try self.writeAtom(atom_index, code.items);
         }
     }
     try self.writeBaseRelocations();
@@ -1642,6 +1663,7 @@ fn writeImportTables(self: *Coff) !void {
         const sect_vm_capacity = self.allocatedVirtualSize(header.virtual_address);
         if (needed_size > sect_vm_capacity) {
             try self.growSectionVM(self.idata_section_index.?, needed_size);
+            self.markRelocsDirtyByAddress(header.virtual_address + needed_size);
         }
 
         header.virtual_size = @max(header.virtual_size, needed_size);
diff --git a/src/link/Coff/Relocation.zig b/src/link/Coff/Relocation.zig
index 45f3a97052..c0fd5562e4 100644
--- a/src/link/Coff/Relocation.zig
+++ b/src/link/Coff/Relocation.zig
@@ -72,62 +72,46 @@ pub fn getTargetAddress(self: Relocation, coff_file: *const Coff) ?u32 {
     }
 }
 
-pub fn resolve(self: *Relocation, atom_index: Atom.Index, coff_file: *Coff) !void {
+pub fn resolve(self: *Relocation, atom_index: Atom.Index, code: []u8, coff_file: *Coff) void {
     const atom = coff_file.getAtom(atom_index);
     const source_sym = atom.getSymbol(coff_file);
-    const source_section = coff_file.sections.get(@enumToInt(source_sym.section_number) - 1).header;
     const source_vaddr = source_sym.value + self.offset;
 
-    const file_offset = source_section.pointer_to_raw_data + source_sym.value - source_section.virtual_address;
-
     const target_vaddr = self.getTargetAddress(coff_file) orelse return;
     const target_vaddr_with_addend = target_vaddr + self.addend;
 
-    log.debug("  ({x}: [() => 0x{x} ({s})) ({s}) (in file at 0x{x})", .{
+    log.debug("  ({x}: [() => 0x{x} ({s})) ({s}) ", .{
         source_vaddr,
         target_vaddr_with_addend,
         coff_file.getSymbolName(self.target),
         @tagName(self.type),
-        file_offset + self.offset,
     });
 
     const ctx: Context = .{
         .source_vaddr = source_vaddr,
         .target_vaddr = target_vaddr_with_addend,
-        .file_offset = file_offset,
         .image_base = coff_file.getImageBase(),
+        .code = code,
+        .ptr_width = coff_file.ptr_width,
     };
 
     switch (coff_file.base.options.target.cpu.arch) {
-        .aarch64 => try self.resolveAarch64(ctx, coff_file),
-        .x86, .x86_64 => try self.resolveX86(ctx, coff_file),
+        .aarch64 => self.resolveAarch64(ctx),
+        .x86, .x86_64 => self.resolveX86(ctx),
         else => unreachable, // unhandled target architecture
     }
-
-    self.dirty = false;
 }
 
 const Context = struct {
     source_vaddr: u32,
     target_vaddr: u32,
-    file_offset: u32,
     image_base: u64,
+    code: []u8,
+    ptr_width: Coff.PtrWidth,
 };
 
-fn resolveAarch64(self: Relocation, ctx: Context, coff_file: *Coff) !void {
-    var buffer: [@sizeOf(u64)]u8 = undefined;
-    switch (self.length) {
-        2 => {
-            const amt = try coff_file.base.file.?.preadAll(buffer[0..4], ctx.file_offset + self.offset);
-            if (amt != 4) return error.InputOutput;
-        },
-        3 => {
-            const amt = try coff_file.base.file.?.preadAll(&buffer, ctx.file_offset + self.offset);
-            if (amt != 8) return error.InputOutput;
-        },
-        else => unreachable,
-    }
-
+fn resolveAarch64(self: Relocation, ctx: Context) void {
+    var buffer = ctx.code[self.offset..];
     switch (self.type) {
         .got_page, .import_page, .page => {
             const source_page = @intCast(i32, ctx.source_vaddr >> 12);
@@ -188,7 +172,7 @@ fn resolveAarch64(self: Relocation, ctx: Context, coff_file: *Coff) !void {
                     buffer[0..4],
                     @truncate(u32, ctx.target_vaddr + ctx.image_base),
                 ),
-                3 => mem.writeIntLittle(u64, &buffer, ctx.target_vaddr + ctx.image_base),
+                3 => mem.writeIntLittle(u64, buffer[0..8], ctx.target_vaddr + ctx.image_base),
                 else => unreachable,
             }
         },
@@ -196,15 +180,10 @@ fn resolveAarch64(self: Relocation, ctx: Context, coff_file: *Coff) !void {
         .got => unreachable,
         .import => unreachable,
     }
-
-    switch (self.length) {
-        2 => try coff_file.base.file.?.pwriteAll(buffer[0..4], ctx.file_offset + self.offset),
-        3 => try coff_file.base.file.?.pwriteAll(&buffer, ctx.file_offset + self.offset),
-        else => unreachable,
-    }
 }
 
-fn resolveX86(self: Relocation, ctx: Context, coff_file: *Coff) !void {
+fn resolveX86(self: Relocation, ctx: Context) void {
+    var buffer = ctx.code[self.offset..];
     switch (self.type) {
         .got_page => unreachable,
         .got_pageoff => unreachable,
@@ -216,26 +195,17 @@ fn resolveX86(self: Relocation, ctx: Context, coff_file: *Coff) !void {
         .got, .import => {
             assert(self.pcrel);
             const disp = @intCast(i32, ctx.target_vaddr) - @intCast(i32, ctx.source_vaddr) - 4;
-            try coff_file.base.file.?.pwriteAll(mem.asBytes(&disp), ctx.file_offset + self.offset);
+            mem.writeIntLittle(i32, buffer[0..4], disp);
         },
         .direct => {
             if (self.pcrel) {
                 const disp = @intCast(i32, ctx.target_vaddr) - @intCast(i32, ctx.source_vaddr) - 4;
-                try coff_file.base.file.?.pwriteAll(mem.asBytes(&disp), ctx.file_offset + self.offset);
-            } else switch (coff_file.ptr_width) {
-                .p32 => try coff_file.base.file.?.pwriteAll(
-                    mem.asBytes(&@intCast(u32, ctx.target_vaddr + ctx.image_base)),
-                    ctx.file_offset + self.offset,
-                ),
+                mem.writeIntLittle(i32, buffer[0..4], disp);
+            } else switch (ctx.ptr_width) {
+                .p32 => mem.writeIntLittle(u32, buffer[0..4], @intCast(u32, ctx.target_vaddr + ctx.image_base)),
                 .p64 => switch (self.length) {
-                    2 => try coff_file.base.file.?.pwriteAll(
-                        mem.asBytes(&@truncate(u32, ctx.target_vaddr + ctx.image_base)),
-                        ctx.file_offset + self.offset,
-                    ),
-                    3 => try coff_file.base.file.?.pwriteAll(
-                        mem.asBytes(&(ctx.target_vaddr + ctx.image_base)),
-                        ctx.file_offset + self.offset,
-                    ),
+                    2 => mem.writeIntLittle(u32, buffer[0..4], @truncate(u32, ctx.target_vaddr + ctx.image_base)),
+                    3 => mem.writeIntLittle(u64, buffer[0..8], ctx.target_vaddr + ctx.image_base),
                     else => unreachable,
                 },
             }