[LLD][COFF] Add support for custom section layout

lld/COFF/Config.h

@@ -212,6 +212,8 @@ struct Configuration {
 
   // Used for /section=.name,{DEKPRSW} to set section attributes.
   std::map<StringRef, uint32_t> section;
+  // Used for /sectionlayout:@ to layout sections in specified order.
+  std::map<std::string, int> sectionOrder;
 
   // Options for manifest files.
   ManifestKind manifest = Default;

lld/COFF/Driver.cpp

@@ -2049,6 +2049,9 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   // Handle /section
   for (auto *arg : args.filtered(OPT_section))
     parseSection(arg->getValue());
+  // Handle /sectionlayout
+  if (auto *arg = args.getLastArg(OPT_sectionlayout))
+    parseSectionLayout(arg->getValue());
 
   // Handle /align
   if (auto *arg = args.getLastArg(OPT_align)) {

lld/COFF/Driver.h

@@ -213,6 +213,7 @@ class LinkerDriver {
   void parseMerge(StringRef);
   void parsePDBPageSize(StringRef);
   void parseSection(StringRef);
+  void parseSectionLayout(StringRef);
 
   void parseSameAddress(StringRef);
 

lld/COFF/DriverUtils.cpp

@@ -214,6 +214,48 @@ void LinkerDriver::parseSection(StringRef s) {
   ctx.config.section[name] = parseSectionAttributes(ctx, attrs);
 }
 
+// Parses /sectionlayout:@ option argument.
+void LinkerDriver::parseSectionLayout(StringRef path) {
+  if (path.starts_with("@"))
+    path = path.substr(1);
+  std::unique_ptr<MemoryBuffer> layoutFile =
+      CHECK(MemoryBuffer::getFile(path), "could not open " + path);
+  StringRef content = layoutFile->getBuffer();
+  int index = 0;
+  std::set<std::string> seenSections;
+
+  while (!content.empty()) {
+    size_t pos = content.find_first_of("\r\n");
+    StringRef line;
+
+    if (pos == StringRef::npos) {
+      line = content;
+      content = StringRef();
+    } else {
+      line = content.substr(0, pos);
+      content = content.substr(pos);
+      while (!content.empty() && (content[0] == '\r' || content[0] == '\n'))
+        content = content.substr(1);
+    }
+
+    line = line.trim();
+    if (line.empty())
+      continue;
+
+    StringRef sectionName = line.split(' ').first;
+    std::string sectionNameStr = sectionName.str();
+
+    if (seenSections.count(sectionNameStr)) {
+      Warn(ctx) << "duplicate section '" << sectionNameStr
+                << "' in section layout file, ignoring";
+      continue;
+    }
+
+    seenSections.insert(sectionNameStr);
+    ctx.config.sectionOrder[sectionNameStr] = index++;
+  }
+}
+
 void LinkerDriver::parseDosStub(StringRef path) {
   std::unique_ptr<MemoryBuffer> stub =
       CHECK(MemoryBuffer::getFile(path), "could not open " + path);

lld/COFF/Options.td

@@ -102,6 +102,7 @@ def pdbstream : Joined<["/", "-", "/?", "-?"], "pdbstream:">,
     MetaVarName<"<name>=<file>">,
     HelpText<"Embed the contents of <file> in the PDB as named stream <name>">;
 def section : P<"section", "Specify section attributes">;
+def sectionlayout : P<"sectionlayout", "Specifies the layout strategy for output sections">;
 def stack   : P<"stack", "Size of the stack">;
 def stub    : P<"stub", "Specify DOS stub file">;
 def subsystem : P<"subsystem", "Specify subsystem">;

lld/COFF/Writer.cpp

@@ -219,6 +219,7 @@ class Writer {
   void sortECChunks();
   void appendECImportTables();
   void removeUnusedSections();
+  void layoutSections();
   void assignAddresses();
   bool isInRange(uint16_t relType, uint64_t s, uint64_t p, int margin,
                  MachineTypes machine);
@@ -783,6 +784,7 @@ void Writer::run() {
     appendECImportTables();
     createDynamicRelocs();
     removeUnusedSections();
+    layoutSections();
     finalizeAddresses();
     removeEmptySections();
     assignOutputSectionIndices();
@@ -1413,6 +1415,34 @@ void Writer::removeUnusedSections() {
   llvm::erase_if(ctx.outputSections, isUnused);
 }
 
+void Writer::layoutSections() {
+  llvm::TimeTraceScope timeScope("Layout sections");
+  if (ctx.config.sectionOrder.empty())
+    return;
+
+  std::unordered_map<const OutputSection *, size_t> originalOrder;
+  for (size_t i = 0; i < ctx.outputSections.size(); ++i)
+    originalOrder[ctx.outputSections[i]] = i;
+
+  llvm::stable_sort(
+      ctx.outputSections,
+      [this, &originalOrder](const OutputSection *a, const OutputSection *b) {
+        auto itA = ctx.config.sectionOrder.find(a->name.str());
+        auto itB = ctx.config.sectionOrder.find(b->name.str());
+
+        if (itA != ctx.config.sectionOrder.end() &&
+            itB != ctx.config.sectionOrder.end())
+          return itA->second < itB->second;
+
+        if (itA == ctx.config.sectionOrder.end() &&
+            itB == ctx.config.sectionOrder.end())
+          return originalOrder[a] < originalOrder[b];
+
+        // Not found in layout file; respect the original order
+        return originalOrder[a] < originalOrder[b];
+      });
+}
+
 // The Windows loader doesn't seem to like empty sections,
 // so we remove them if any.
 void Writer::removeEmptySections() {

lld/test/COFF/Inputs/sectionlayout.yaml

@@ -0,0 +1,25 @@
+--- !COFF
+header:
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  Characteristics: []
+sections:
+  - Name:            '.text1'
+    Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ]
+    Alignment:       16
+    SectionData:     B82A000000C3
+  - Name:            '.text2'
+    Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ]
+    Alignment:       16
+    SectionData:     CC
+  - Name:            '.text3'
+    Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ]
+    Alignment:       16
+    SectionData:     CC
+symbols:
+  - Name:            main
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_FUNCTION
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...

lld/test/COFF/sectionlayout.test

@@ -0,0 +1,122 @@
+RUN: yaml2obj %p/Inputs/sectionlayout.yaml -o %t.obj
+
+## Error on non-exist input layout file
+RUN: not lld-link /entry:main /sectionlayout:doesnotexist.txt %t.obj
+
+## Order in 1 -> 3 -> 2
+RUN: echo ".text1" > %t.layout.txt
+RUN: echo ".text3" >> %t.layout.txt
+RUN: echo ".text2" >> %t.layout.txt
+RUN: lld-link /out:%t.exe /entry:main /sectionlayout:%t.layout.txt %t.obj
+RUN: llvm-readobj --sections %t.exe | FileCheck -check-prefix=CHECK1 %s
+
+## While /sectionlayout:abc is valid, /sectionlayout:@abc is also accepted (to align with MS link.exe)
+RUN: lld-link /out:%t.exe /entry:main /sectionlayout:@%t.layout.txt %t.obj
+RUN: llvm-readobj --sections %t.exe | FileCheck -check-prefix=CHECK1 %s
+
+CHECK1: Sections [
+CHECK1:   Section {
+CHECK1:     Number: 1
+CHECK1:     Name: .text1
+CHECK1:   }
+CHECK1:   Section {
+CHECK1:     Number: 2
+CHECK1:     Name: .text3
+CHECK1:   }
+CHECK1:   Section {
+CHECK1:     Number: 3
+CHECK1:     Name: .text2
+CHECK1:   }
+CHECK1: ]
+
+## Order in 3 -> 2 -> 1
+RUN: echo ".text3" > %t.layout.txt
+RUN: echo ".text2" >> %t.layout.txt
+RUN: echo ".text1" >> %t.layout.txt
+RUN: lld-link /out:%t.exe /entry:main /sectionlayout:%t.layout.txt %t.obj
+RUN: llvm-readobj --sections %t.exe | FileCheck -check-prefix=CHECK2 %s
+
+CHECK2: Sections [
+CHECK2:   Section {
+CHECK2:     Number: 1
+CHECK2:     Name: .text3
+CHECK2:   }
+CHECK2:   Section {
+CHECK2:     Number: 2
+CHECK2:     Name: .text2
+CHECK2:   }
+CHECK2:   Section {
+CHECK2:     Number: 3
+CHECK2:     Name: .text1
+CHECK2:   }
+CHECK2: ]
+
+## Put non-exisist section in layout file has no effect; original order is respected
+RUN: echo "notexist" > %t.layout.txt
+RUN: lld-link /out:%t.exe /entry:main /sectionlayout:%t.layout.txt %t.obj
+RUN: llvm-readobj --sections %t.exe | FileCheck -check-prefix=CHECK3 %s
+
+## Empty layout file has no effect
+RUN: echo "" > %t.layout.txt
+RUN: lld-link /out:%t.exe /entry:main /sectionlayout:%t.layout.txt %t.obj
+RUN: llvm-readobj --sections %t.exe | FileCheck -check-prefix=CHECK3 %s
+
+## Empty layout file has no effect
+RUN: echo " " > %t.layout.txt
+RUN: echo " " >> %t.layout.txt
+RUN: lld-link /out:%t.exe /entry:main /sectionlayout:%t.layout.txt %t.obj
+RUN: llvm-readobj --sections %t.exe | FileCheck -check-prefix=CHECK3 %s
+
+CHECK3: Sections [
+CHECK3:   Section {
+CHECK3:     Number: 1
+CHECK3:     Name: .text1
+CHECK3:   }
+CHECK3:   Section {
+CHECK3:     Number: 2
+CHECK3:     Name: .text2
+CHECK3:   }
+CHECK3:   Section {
+CHECK3:     Number: 3
+CHECK3:     Name: .text3
+CHECK3:   }
+CHECK3: ]
+
+## Order in 3 -> 1, but 2 remains unspecified
+RUN: echo ".text3" > %t.layout.txt
+RUN: echo ".text1" >> %t.layout.txt
+RUN: lld-link /out:%t.exe /entry:main /sectionlayout:%t.layout.txt %t.obj
+RUN: llvm-readobj --sections %t.exe | FileCheck -check-prefix=CHECK4 %s
+
+CHECK4: Sections [
+CHECK4:   Section {
+CHECK4:     Number: 1
+CHECK4:     Name: .text3
+CHECK4:   }
+CHECK4:   Section {
+CHECK4:     Number: 2
+CHECK4:     Name: .text1
+CHECK4:   }
+CHECK4:   Section {
+CHECK4:     Number: 3
+CHECK4:     Name: .text2
+CHECK4:   }
+CHECK4: ]
+
+## Order in 3 -> 2, but 1 remains unspecified.
+## 1 should be the first, as the original order (1 -> 2 -> 3) is respected
+RUN: echo ".text3" > %t.layout.txt
+RUN: echo ".text2" >> %t.layout.txt
+RUN: lld-link /out:%t.exe /entry:main /sectionlayout:%t.layout.txt %t.obj
+RUN: llvm-readobj --sections %t.exe | FileCheck -check-prefix=CHECK1 %s
+
+## Order in 3 -> 2 -> 1, multiple specification has no effect (the first one is used)
+RUN: echo ".text3" > %t.layout.txt
+RUN: echo ".text3" >> %t.layout.txt
+RUN: echo ".text3" >> %t.layout.txt
+RUN: echo ".text2" >> %t.layout.txt
+RUN: echo ".text2" >> %t.layout.txt
+RUN: echo ".text1" >> %t.layout.txt
+RUN: echo ".text3" >> %t.layout.txt
+RUN: lld-link /out:%t.exe /entry:main /sectionlayout:%t.layout.txt %t.obj
+RUN: llvm-readobj --sections %t.exe | FileCheck -check-prefix=CHECK2 %s