[Mach-O] wip

macho/mold.h

@@ -125,50 +125,43 @@ class OutputLazyBindSection : public OutputSection {
 
 class ExportEncoder {
 public:
-  void add(std::string_view name, u64 addr, u32 flags);
-  void finish();
-
-  std::vector<u8> contents;
+  void add(std::string_view name, u32 flags, u64 addr);
+  i64 finish();
+  void write_trie(u8 *buf);
 
 private:
   struct Entry {
     std::string_view name;
-    u64 addr;
     u32 flags;
+    u64 addr;
   };
 
   struct TrieNode {
     std::string_view prefix;
     bool is_leaf = false;
-    u64 addr = 0;
     u32 flags = 0;
+    u64 addr = 0;
     u32 offset = 0;
     u32 size = UINT32_MAX;
     std::unique_ptr<TrieNode> children[256] = {};
   };
 
   static void construct_trie(TrieNode &parent, std::span<Entry> entries, i64 len);
   static i64 common_prefix_len(std::span<Entry> entries, i64 len);
+  static i64 set_offset(TrieNode &node, i64 offset);
+  void write_trie(u8 *buf, TrieNode &node);
 
+  TrieNode root;
   std::vector<Entry> entries;
 };
 
 class OutputExportSection : public OutputSection {
 public:
-  OutputExportSection(OutputSegment &parent) : OutputSection(parent) {
-    is_hidden = true;
-    hdr.size = contents.size();
-  }
-
+  OutputExportSection(OutputSegment &parent);
   void copy_buf(Context &ctx) override;
 
-  std::vector<u8> contents = {
-    0x00, 0x01, 0x5f, 0x00, 0x05, 0x00, 0x03, 0x5f, 0x6d, 0x68, 0x5f, 0x65,
-    0x78, 0x65, 0x63, 0x75, 0x74, 0x65, 0x5f, 0x68, 0x65, 0x61, 0x64, 0x65,
-    0x72, 0x00, 0x28, 0x68, 0x65, 0x6c, 0x6c, 0x6f, 0x00, 0x2c, 0x6d, 0x61,
-    0x69, 0x6e, 0x00, 0x31, 0x02, 0x00, 0x00, 0x00, 0x03, 0x00, 0xd0, 0x7e,
-    0x00, 0x03, 0x00, 0xf0, 0x7e, 0x00, 0x00, 0x00,
-  };
+private:
+  ExportEncoder enc;
 };
 
 class OutputFunctionStartsSection : public OutputSection {

macho/output-chunks.cc

@@ -395,7 +395,7 @@ void BindEncoder::add(i64 dylib_idx, std::string_view sym, i64 flags,
     assert(flags < 16);
     buf.push_back(BIND_OPCODE_SET_SYMBOL_TRAILING_FLAGS_IMM | flags);
     buf.insert(buf.end(), (u8 *)sym.data(), (u8 *)(sym.data() + sym.size()));
-    buf.push_back(0);
+    buf.push_back('\0');
   }
 
   if (last_seg != seg_idx || last_off != offset) {
@@ -449,17 +449,25 @@ void OutputLazyBindSection::copy_buf(Context &ctx) {
   write_vector(ctx.buf + hdr.offset, contents);
 }
 
-void ExportEncoder::add(std::string_view name, u64 addr, u32 flags) {
-  entries.push_back({name, addr, flags});
+void ExportEncoder::add(std::string_view name, u32 flags, u64 addr) {
+  entries.push_back({name, flags, addr});
 }
 
-void ExportEncoder::finish() {
+i64 ExportEncoder::finish() {
   sort(entries, [](const Entry &a, const Entry &b) {
     return a.name < b.name;
   });
 
-  TrieNode root;
   construct_trie(root, entries, 0);
+
+  i64 size = set_offset(root, 0);
+  for (;;) {
+    i64 size2 = set_offset(root, 0);
+    if (size == size2)
+      break;
+    size = size2;
+  }
+  return size;
 }
 
 void
@@ -469,8 +477,8 @@ ExportEncoder::construct_trie(TrieNode &parent, std::span<Entry> entries, i64 le
 
   if (entries[0].name.size() == len) {
     parent.is_leaf = true;
-    parent.addr = entries[0].addr;
     parent.flags = entries[0].flags;
+    parent.addr = entries[0].addr;
     entries = entries.subspan(1);
   }
 
@@ -504,8 +512,73 @@ i64 ExportEncoder::common_prefix_len(std::span<Entry> entries, i64 len) {
   }
 }
 
+i64 ExportEncoder::set_offset(TrieNode &node, i64 offset) {
+  node.offset = offset;
+
+  if (node.is_leaf) {
+    node.size = uleb_size(node.flags) + uleb_size(node.addr);
+    node.size += uleb_size(node.size);
+  } else {
+    node.size = 1;
+  }
+
+  node.size++; // # of children
+
+  for (std::unique_ptr<TrieNode> &child : node.children) {
+    if (child) {
+      // +1 for NUL byte
+      node.size += child->prefix.size() + 1 + uleb_size(child->offset);
+    }
+  }
+
+  offset += node.size;
+
+  for (std::unique_ptr<TrieNode> &child : node.children)
+    if (child)
+      offset = set_offset(*child, offset);
+  return offset;
+}
+
+void ExportEncoder::write_trie(u8 *start) {
+  write_trie(start, root);
+}
+
+void ExportEncoder::write_trie(u8 *start, TrieNode &node) {
+  u8 *buf = start;
+
+  if (node.is_leaf) {
+    buf += write_uleb(buf, uleb_size(node.flags) + uleb_size(node.addr));
+    buf += write_uleb(buf, node.flags);
+    buf += write_uleb(buf, node.addr);
+  } else {
+    *buf++ = 1;
+  }
+
+  u8 *num_children = buf++;
+  *num_children = 0;
+
+  for (std::unique_ptr<TrieNode> &child : node.children) {
+    if (child) {
+      *num_children += 1;
+      buf += write_string(buf, child->prefix);
+      buf += write_uleb(buf, child->offset);
+    }
+  }
+
+  for (std::unique_ptr<TrieNode> &child : node.children)
+    if (child)
+      write_trie(start, *child);
+}
+
+OutputExportSection::OutputExportSection(OutputSegment &parent)
+  : OutputSection(parent) {
+  is_hidden = true;
+  enc.add("__mh_execute_header", 0, 0x100000000);
+  hdr.size = enc.finish();
+}
+
 void OutputExportSection::copy_buf(Context &ctx) {
-  write_vector(ctx.buf + hdr.offset, contents);
+  enc.write_trie(ctx.buf + hdr.offset);
 }
 
 void OutputFunctionStartsSection::copy_buf(Context &ctx) {

mold.h

@@ -194,14 +194,17 @@ inline void write32be(u8 *buf, u32 val) {
   buf[3] = val;
 }
 
-inline void write_string(u8 *buf, std::string_view str) {
+inline i64 write_string(u8 *buf, std::string_view str) {
   memcpy(buf, str.data(), str.size());
   buf[str.size()] = '\0';
+  return str.size() + 1;
 }
 
 template <typename T>
-inline void write_vector(u8 *buf, const std::vector<T> &vec) {
-  memcpy(buf, vec.data(), vec.size() * sizeof(T));
+inline i64 write_vector(u8 *buf, const std::vector<T> &vec) {
+  i64 sz = vec.size() * sizeof(T);
+  memcpy(buf, vec.data(), sz);
+  return sz;
 }
 
 inline void encode_uleb(std::vector<u8> &vec, u64 val) {
@@ -212,6 +215,25 @@ inline void encode_uleb(std::vector<u8> &vec, u64 val) {
   } while (val);
 }
 
+inline i64 write_uleb(u8 *buf, u64 val) {
+  i64 i = 0;
+  do {
+    u8 byte = val & 0x7f;
+    val >>= 7;
+    buf[i++] = val ? (byte | 0x80) : byte;
+  } while (val);
+  return i;
+}
+
+inline i64 uleb_size(u64 val) {
+  i64 i = 0;
+  do {
+    i++;
+    val >>= 7;
+  } while (val);
+  return i;
+}
+
 //
 // Concurrent Map
 //