[wasm-ld] Allow importing/exporting the output module's memory with arbitrary names

This adds an `--export-memory` option to wasm-ld which allows passing
a name to give to the exported memory, and extends `--import-memory` to
allow passing a <module>,<name> pair specifying where the memory should
be imported from.

This is based on https://reviews.llvm.org/D131376, with the main
difference being that it only supports exporting memory by one name.

Differential Revision: https://reviews.llvm.org/D135898

9 files changed, 147 insertions, 13 deletions

diff --git a/lld/test/wasm/memory-naming.test b/lld/test/wasm/memory-naming.test
new file mode 100644
index 000000000000..b4aabaeeac35
--- /dev/null
+++ b/lld/test/wasm/memory-naming.test
@@ -0,0 +1,88 @@
+# RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown %p/Inputs/start.s -o %t.start.o
+# RUN: wasm-ld --export-memory=foo -o %t.wasm %t.start.o
+# RUN: obj2yaml %t.wasm | FileCheck %s
+
+# Verify that the --export-memory=<name> option changes the exported name of the module's memory
+
+# CHECK:       - Type:            EXPORT
+# CHECK-NEXT:    Exports:
+# CHECK-NEXT:      - Name:            foo
+# CHECK-NEXT:        Kind:            MEMORY
+# CHECK-NEXT:        Index:           0
+# CHECK-NEXT:      - Name:            _start
+# CHECK-NEXT:        Kind:            FUNCTION
+# CHECK-NEXT:        Index:           0
+# CHECK-NEXT:  - Type:
+
+# RUN:wasm-ld --export-memory --export-memory -o %t.unnamed.wasm %t.start.o
+# RUN: obj2yaml %t.unnamed.wasm | FileCheck -check-prefix=CHECK-UNNAMED %s
+
+# Verify that the --export-memory option without a parameter exports the memory
+# as "memory"
+
+# CHECK-UNNAMED:       - Type:            EXPORT
+# CHECK-UNNAMED-NEXT:    Exports:
+# CHECK-UNNAMED-NEXT:      - Name:            memory
+# CHECK-UNNAMED-NEXT:        Kind:            MEMORY
+# CHECK-UNNAMED-NEXT:        Index:           0
+# CHECK-UNNAMED-NEXT:      - Name:            _start
+# CHECK-UNNAMED-NEXT:        Kind:            FUNCTION
+# CHECK-UNNAMED-NEXT:        Index:           0
+# CHECK-UNNAMED-NEXT:  - Type:
+
+# RUN:wasm-ld --export-memory=foo --export-memory=foo -o %t.duplicate.wasm %t.start.o
+# RUN: obj2yaml %t.duplicate.wasm | FileCheck -check-prefix=CHECK-DUPLICATE %s
+
+# Verify that passing --export-memory with the same name twice works
+
+# CHECK-DUPLICATE:       - Type:            EXPORT
+# CHECK-DUPLICATE-NEXT:    Exports:
+# CHECK-DUPLICATE-NEXT:      - Name:            foo
+# CHECK-DUPLICATE-NEXT:        Kind:            MEMORY
+# CHECK-DUPLICATE-NEXT:        Index:           0
+# CHECK-DUPLICATE-NEXT:      - Name:            _start
+# CHECK-DUPLICATE-NEXT:        Kind:            FUNCTION
+# CHECK-DUPLICATE-NEXT:        Index:           0
+# CHECK-DUPLICATE-NEXT:  - Type:
+
+# RUN:wasm-ld --import-memory=foo,bar -o %t.import.wasm %t.start.o
+# RUN: obj2yaml %t.import.wasm | FileCheck -check-prefix=CHECK-IMPORT %s
+
+# Verify that memory imports can be renamed, and that no memory is exported by
+# default when memory is being imported
+
+# CHECK-IMPORT:       - Type:            IMPORT
+# CHECK-IMPORT-NEXT:    Imports:
+# CHECK-IMPORT-NEXT:      - Module:          foo
+# CHECK-IMPORT-NEXT:        Field:           bar
+# CHECK-IMPORT-NEXT:        Kind:            MEMORY
+# CHECK-IMPORT-NEXT:        Memory:
+# CHECK-IMPORT-NEXT:          Minimum:         0x2
+# CHECK-IMPORT:      - Type:            EXPORT
+# CHECK-IMPORT-NEXT:    Exports:
+# CHECK-IMPORT-NEXT:      - Name:            _start
+# CHECK-IMPORT-NEXT:        Kind:            FUNCTION
+# CHECK-IMPORT-NEXT:        Index:           0
+# CHECK-IMPORT-NEXT:  - Type:
+
+# RUN:wasm-ld --import-memory=foo,bar --export-memory=qux -o %t.both.wasm %t.start.o
+# RUN: obj2yaml %t.both.wasm | FileCheck -check-prefix=CHECK-BOTH %s
+
+# Verify that memory can be both imported and exported from a module
+
+# CHECK-BOTH:       - Type:            IMPORT
+# CHECK-BOTH-NEXT:    Imports:
+# CHECK-BOTH-NEXT:      - Module:          foo
+# CHECK-BOTH-NEXT:        Field:           bar
+# CHECK-BOTH-NEXT:        Kind:            MEMORY
+# CHECK-BOTH-NEXT:        Memory:
+# CHECK-BOTH-NEXT:          Minimum:         0x2
+# CHECK-BOTH:       - Type:            EXPORT
+# CHECK-BOTH-NEXT:    Exports:
+# CHECK-BOTH-NEXT:      - Name:            qux
+# CHECK-BOTH-NEXT:        Kind:            MEMORY
+# CHECK-BOTH-NEXT:        Index:           0
+# CHECK-BOTH-NEXT:      - Name:            _start
+# CHECK-BOTH-NEXT:        Kind:            FUNCTION
+# CHECK-BOTH-NEXT:        Index:           0
+# CHECK-BOTH-NEXT:  - Type:
diff --git a/lld/wasm/Config.h b/lld/wasm/Config.h
index 35171b5afd8f..71d6bd99f649 100644
--- a/lld/wasm/Config.h
+++ b/lld/wasm/Config.h
@@ -39,7 +39,8 @@ struct Configuration {
   bool extendedConst;
   bool growableTable;
   bool gcSections;
-  bool importMemory;
+  llvm::Optional<std::pair<llvm::StringRef, llvm::StringRef>> memoryImport;
+  llvm::Optional<llvm::StringRef> memoryExport;
   bool sharedMemory;
   bool importTable;
   bool importUndefined;
diff --git a/lld/wasm/Driver.cpp b/lld/wasm/Driver.cpp
index b0d37b9fc3a4..68e217f7e987 100644
--- a/lld/wasm/Driver.cpp
+++ b/lld/wasm/Driver.cpp
@@ -393,7 +393,27 @@ static void readConfigs(opt::InputArgList &args) {
   config->exportAll = args.hasArg(OPT_export_all);
   config->exportTable = args.hasArg(OPT_export_table);
   config->growableTable = args.hasArg(OPT_growable_table);
-  config->importMemory = args.hasArg(OPT_import_memory);
+
+  if (args.hasArg(OPT_import_memory_with_name)) {
+    config->memoryImport =
+        args.getLastArgValue(OPT_import_memory_with_name).split(",");
+  } else if (args.hasArg(OPT_import_memory)) {
+    config->memoryImport =
+        std::pair<llvm::StringRef, llvm::StringRef>(defaultModule, memoryName);
+  } else {
+    config->memoryImport =
+        llvm::Optional<std::pair<llvm::StringRef, llvm::StringRef>>();
+  }
+
+  if (args.hasArg(OPT_export_memory_with_name)) {
+    config->memoryExport =
+        args.getLastArgValue(OPT_export_memory_with_name);
+  } else if (args.hasArg(OPT_export_memory)) {
+    config->memoryExport = memoryName;
+  } else {
+    config->memoryExport = llvm::Optional<llvm::StringRef>();
+  }
+
   config->sharedMemory = args.hasArg(OPT_shared_memory);
   config->importTable = args.hasArg(OPT_import_table);
   config->importUndefined = args.hasArg(OPT_import_undefined);
@@ -510,9 +530,21 @@ static void setConfigs() {
   }
 
   if (config->shared) {
-    config->importMemory = true;
+    if (config->memoryExport.hasValue()) {
+      error("--export-memory is incompatible with --shared");
+    }
+    if (!config->memoryImport.hasValue()) {
+      config->memoryImport =
+          std::pair<llvm::StringRef, llvm::StringRef>(defaultModule, memoryName);
+    }
     config->importUndefined = true;
   }
+
+  // If neither export-memory nor import-memory is specified, default to
+  // exporting memory under its default name.
+  if (!config->memoryExport.hasValue() && !config->memoryImport.hasValue()) {
+    config->memoryExport = memoryName;
+  }
 }
 
 // Some command line options or some combinations of them are not allowed.
diff --git a/lld/wasm/Options.td b/lld/wasm/Options.td
index ec452a38ae53..2871ae6d576b 100644
--- a/lld/wasm/Options.td
+++ b/lld/wasm/Options.td
@@ -189,7 +189,15 @@ def global_base: JJ<"global-base=">,
   HelpText<"Where to start to place global data">;
 
 def import_memory: FF<"import-memory">,
-  HelpText<"Import memory from the environment">;
+  HelpText<"Import the module's memory from the default module of \"env\" with the name \"memory\".">;
+def import_memory_with_name: JJ<"import-memory=">,
+  HelpText<"Import the module's memory from the passed module with the passed name.">,
+  MetaVarName<"<module>,<name>">;
+
+def export_memory: FF<"export-memory">,
+  HelpText<"Export the module's memory with the default name of \"memory\"">;
+def export_memory_with_name: JJ<"export-memory=">,
+  HelpText<"Export the module's memory with the passed name">;
 
 def shared_memory: FF<"shared-memory">,
   HelpText<"Use shared linear memory">;
diff --git a/lld/wasm/Symbols.cpp b/lld/wasm/Symbols.cpp
index c6f6e04ad4c8..5ef92dde7cc5 100644
--- a/lld/wasm/Symbols.cpp
+++ b/lld/wasm/Symbols.cpp
@@ -457,6 +457,7 @@ void printTraceSymbol(Symbol *sym) {
 
 const char *defaultModule = "env";
 const char *functionTableName = "__indirect_function_table";
+const char *memoryName = "memory";
 
 } // namespace wasm
 } // namespace lld
diff --git a/lld/wasm/Symbols.h b/lld/wasm/Symbols.h
index ee6c5db181a1..ac0e42a2496b 100644
--- a/lld/wasm/Symbols.h
+++ b/lld/wasm/Symbols.h
@@ -26,6 +26,9 @@ extern const char *defaultModule;
 // The name under which to import or export the wasm table.
 extern const char *functionTableName;
 
+// The name under which to import or export the wasm memory.
+extern const char *memoryName;
+
 using llvm::wasm::WasmSymbolType;
 
 class InputFile;
diff --git a/lld/wasm/SyntheticSections.cpp b/lld/wasm/SyntheticSections.cpp
index c39d86cee87e..3be2beb73578 100644
--- a/lld/wasm/SyntheticSections.cpp
+++ b/lld/wasm/SyntheticSections.cpp
@@ -162,7 +162,7 @@ void TypeSection::writeBody() {
 uint32_t ImportSection::getNumImports() const {
   assert(isSealed);
   uint32_t numImports = importedSymbols.size() + gotSymbols.size();
-  if (config->importMemory)
+  if (config->memoryImport.hasValue())
     ++numImports;
   return numImports;
 }
@@ -234,10 +234,10 @@ void ImportSection::writeBody() {
 
   bool is64 = config->is64.value_or(false);
 
-  if (config->importMemory) {
+  if (config->memoryImport) {
     WasmImport import;
-    import.Module = defaultModule;
-    import.Field = "memory";
+    import.Module = config->memoryImport.getValue().first;
+    import.Field = config->memoryImport.getValue().second;
     import.Kind = WASM_EXTERNAL_MEMORY;
     import.Memory.Flags = 0;
     import.Memory.Minimum = out.memorySec->numMemoryPages;
diff --git a/lld/wasm/SyntheticSections.h b/lld/wasm/SyntheticSections.h
index 06579054b630..fe2d3879b6e2 100644
--- a/lld/wasm/SyntheticSections.h
+++ b/lld/wasm/SyntheticSections.h
@@ -230,7 +230,7 @@ class MemorySection : public SyntheticSection {
 public:
   MemorySection() : SyntheticSection(llvm::wasm::WASM_SEC_MEMORY) {}
 
-  bool isNeeded() const override { return !config->importMemory; }
+  bool isNeeded() const override { return !config->memoryImport.hasValue(); }
   void writeBody() override;
 
   uint64_t numMemoryPages = 0;
diff --git a/lld/wasm/Writer.cpp b/lld/wasm/Writer.cpp
index 9937c009f62f..bd0147435333 100644
--- a/lld/wasm/Writer.cpp
+++ b/lld/wasm/Writer.cpp
@@ -577,7 +577,7 @@ done:
   // memory is not being imported then we can assume its zero initialized.
   // In the case the memory is imported, and we can use the memory.fill
   // instruction, then we can also avoid including the segments.
-  if (config->importMemory && !allowed.count("bulk-memory"))
+  if (config->memoryImport.hasValue() && !allowed.count("bulk-memory"))
     config->emitBssSegments = true;
 
   if (allowed.count("extended-const"))
@@ -671,9 +671,10 @@ void Writer::calculateExports() {
   if (config->relocatable)
     return;
 
-  if (!config->relocatable && !config->importMemory)
+  if (!config->relocatable && config->memoryExport.hasValue()) {
     out.exportSec->exports.push_back(
-        WasmExport{"memory", WASM_EXTERNAL_MEMORY, 0});
+        WasmExport{*config->memoryExport, WASM_EXTERNAL_MEMORY, 0});
+  }
 
   unsigned globalIndex =
       out.importSec->getNumImportedGlobals() + out.globalSec->numGlobals();
@@ -1007,7 +1008,7 @@ static void createFunction(DefinedFunction *func, StringRef bodyContent) {
 bool Writer::needsPassiveInitialization(const OutputSegment *segment) {
   // If bulk memory features is supported then we can perform bss initialization
   // (via memory.fill) during `__wasm_init_memory`.
-  if (config->importMemory && !segment->requiredInBinary())
+  if (config->memoryImport.hasValue() && !segment->requiredInBinary())
     return true;
   return segment->initFlags & WASM_DATA_SEGMENT_IS_PASSIVE;
 }