From 140bac6395610a3b9e5d5f92e2b366a27f6dcdc8 Mon Sep 17 00:00:00 2001
From: Luuk de Gram <luuk@degram.dev>
Date: Sat, 25 Jun 2022 18:36:56 +0200
Subject: [PATCH] link/wasm: Sort data segments

We now ensure the "bss" section is last, which allows us to not
emit this section and let the runtime initialize the memory with 0's instead.
This allows for smaller binaries.
The order of the other segments is arbitrary and does not matter, this may
change in the future.
---
 src/link/Wasm.zig | 42 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/src/link/Wasm.zig b/src/link/Wasm.zig
index 834ac3e1e1..e57f87d37c 100644
--- a/src/link/Wasm.zig
+++ b/src/link/Wasm.zig
@@ -1255,6 +1255,9 @@ fn parseAtom(self: *Wasm, atom: *Atom, kind: Kind) !void {
 }
 
 fn allocateAtoms(self: *Wasm) !void {
+    // first sort the data segments
+    try sortDataSegments(self);
+
     var it = self.atoms.iterator();
     while (it.next()) |entry| {
         const segment = &self.segments.items[entry.key_ptr.*];
@@ -1278,6 +1281,36 @@ fn allocateAtoms(self: *Wasm) !void {
     }
 }
 
+fn sortDataSegments(self: *Wasm) !void {
+    var new_mapping: std.StringArrayHashMapUnmanaged(u32) = .{};
+    try new_mapping.ensureUnusedCapacity(self.base.allocator, self.data_segments.count());
+    errdefer new_mapping.deinit(self.base.allocator);
+
+    const keys = try self.base.allocator.dupe([]const u8, self.data_segments.keys());
+    defer self.base.allocator.free(keys);
+
+    const SortContext = struct {
+        fn sort(_: void, lhs: []const u8, rhs: []const u8) bool {
+            return order(lhs) <= order(rhs);
+        }
+
+        fn order(name: []const u8) u8 {
+            if (mem.startsWith(u8, name, ".rodata")) return 0;
+            if (mem.startsWith(u8, name, ".data")) return 1;
+            if (mem.startsWith(u8, name, ".text")) return 2;
+            return 3;
+        }
+    };
+
+    std.sort.sort([]const u8, keys, {}, SortContext.sort);
+    for (keys) |key| {
+        const segment_index = self.data_segments.get(key).?;
+        new_mapping.putAssumeCapacity(key, segment_index);
+    }
+    self.data_segments.deinit(self.base.allocator);
+    self.data_segments = new_mapping;
+}
+
 fn setupImports(self: *Wasm) !void {
     log.debug("Merging imports", .{});
     var discarded_it = self.discarded.keyIterator();
@@ -2337,8 +2370,13 @@ fn emitNameSection(self: *Wasm, file: fs.File, arena: Allocator) !void {
         }
     }
     // data segments are already 'ordered'
-    for (self.data_segments.keys()) |key, index| {
-        segments.appendAssumeCapacity(.{ .index = @intCast(u32, index), .name = key });
+    var data_segment_index: u32 = 0;
+    for (self.data_segments.keys()) |key| {
+        // bss section is not emitted when this condition holds true, so we also
+        // do not output a name for it.
+        if (!self.base.options.import_memory and std.mem.eql(u8, key, ".bss")) continue;
+        segments.appendAssumeCapacity(.{ .index = data_segment_index, .name = key });
+        data_segment_index += 1;
     }
 
     std.sort.sort(Name, funcs.values(), {}, Name.lessThan);