[lld][WebAssembly] Add an `--initial-heap` option (#75594)

It is beneficial to preallocate a certain number of pages in the linear
memory (i. e. use the "minimum" field of WASM memories) so that fewer
"memory.grow"s are needed at startup.

So far, the way to do that has been to pass the "--initial-memory"
option to the linker. It works, but has the very significant downside of
requiring the user to know the size of static data beforehand, as it
must not exceed the number of bytes passed-in as "--initial-memory".

The new "--initial-heap" option avoids this downside by simply appending
the specified number of pages to static data (and stack), regardless of
how large they already are.

Ref: https://github.com/emscripten-core/emscripten/issues/20888.
GitOrigin-RevId: b2cdf3cc4c08729d0ff582d55e40793a20bbcdcc
diff --git a/docs/WebAssembly.rst b/docs/WebAssembly.rst
index dad3177..3f554de 100644
--- a/docs/WebAssembly.rst
+++ b/docs/WebAssembly.rst
@@ -123,9 +123,13 @@
    is not possible for undefined data symbols.  Undefined data symbols will
    still be reported as normal (in accordance with ``--unresolved-symbols``).
 
+.. option:: --initial-heap=<value>
+
+  Initial size of the heap. Default: zero.
+
 .. option:: --initial-memory=<value>
 
-  Initial size of the linear memory. Default: static data size.
+  Initial size of the linear memory. Default: the sum of stack, static data and heap sizes.
 
 .. option:: --max-memory=<value>
 
diff --git a/test/wasm/initial-heap.test b/test/wasm/initial-heap.test
new file mode 100644
index 0000000..3e8bbd3
--- /dev/null
+++ b/test/wasm/initial-heap.test
@@ -0,0 +1,27 @@
+RUN: llvm-mc -filetype=obj -triple=wasm32-unknown-unknown %p/Inputs/start.s -o %t.o
+
+; The initial heap size will be added to the stack size
+RUN: wasm-ld %t.o -o %t1.wasm --stack-first -z stack-size=65536 --initial-heap=131072
+RUN: obj2yaml %t1.wasm | FileCheck %s --check-prefixes=CHECK,CHECK-2P
+
+; Also test that we can parse and process a large size correctly
+RUN: wasm-ld %t.o -o %t2.wasm --stack-first -z stack-size=65536 --initial-heap=4294901760
+RUN: obj2yaml %t2.wasm | FileCheck %s --check-prefixes=CHECK,CHECK-4G
+
+CHECK:      - Type:            MEMORY
+CHECK-NEXT:   Memories:
+CHECK-2P-NEXT:    Minimum:         0x3
+CHECK-4G-NEXT:    Minimum:         0x10000
+
+; Test various error cases.
+RUN: not wasm-ld %t.o -o %t3.wasm --initial-heap=131073 2>&1 | FileCheck %s --check-prefix NOT-PAGE-MULTIPLE
+RUN: not wasm-ld %t.o -o %t4.wasm --stack-first -z stack-size=65536 --initial-heap=4295032832 2>&1 | FileCheck %s --check-prefix TOO-LARGE-BY-ITSELF
+RUN: not wasm-ld %t.o -o %t5.wasm --stack-first -z stack-size=131072 --initial-heap=4294901760 2>&1 | FileCheck %s --check-prefix TOO-LARGE-WITH-STACK
+RUN: not wasm-ld %t.o -o %t6.wasm --stack-first -z stack-size=65536 --initial-heap=131072 --initial-memory=131072 2>&1 | FileCheck %s --check-prefix INITIAL-MEMORY-TOO-SMALL
+RUN: not wasm-ld %t.o -o %t7.wasm --stack-first -z stack-size=65536 --initial-heap=131072 --max-memory=131072 2>&1 | FileCheck %s --check-prefix MAX-MEMORY-TOO-SMALL
+
+NOT-PAGE-MULTIPLE: initial heap must be 65536-byte aligned
+TOO-LARGE-BY-ITSELF: initial heap too large, cannot be greater than 4294901760
+TOO-LARGE-WITH-STACK: initial heap too large, cannot be greater than 4294836224
+INITIAL-MEMORY-TOO-SMALL: initial memory too small, 196608 bytes needed
+MAX-MEMORY-TOO-SMALL: maximum memory too small, 196608 bytes needed
diff --git a/wasm/Config.h b/wasm/Config.h
index d76d438..104d470 100644
--- a/wasm/Config.h
+++ b/wasm/Config.h
@@ -68,6 +68,7 @@
   bool isStatic = false;
   bool trace;
   uint64_t globalBase;
+  uint64_t initialHeap;
   uint64_t initialMemory;
   uint64_t maxMemory;
   uint64_t zStackSize;
diff --git a/wasm/Driver.cpp b/wasm/Driver.cpp
index c68fe33..a354260 100644
--- a/wasm/Driver.cpp
+++ b/wasm/Driver.cpp
@@ -502,9 +502,10 @@
   errorHandler().verbose = args.hasArg(OPT_verbose);
   LLVM_DEBUG(errorHandler().verbose = true);
 
-  config->initialMemory = args::getInteger(args, OPT_initial_memory, 0);
-  config->globalBase = args::getInteger(args, OPT_global_base, 0);
   config->tableBase = args::getInteger(args, OPT_table_base, 0);
+  config->globalBase = args::getInteger(args, OPT_global_base, 0);
+  config->initialHeap = args::getInteger(args, OPT_initial_heap, 0);
+  config->initialMemory = args::getInteger(args, OPT_initial_memory, 0);
   config->maxMemory = args::getInteger(args, OPT_max_memory, 0);
   config->zStackSize =
       args::getZOptionValue(args, OPT_z, "stack-size", WasmPageSize);
diff --git a/wasm/Options.td b/wasm/Options.td
index 2df6196..95ebc20 100644
--- a/wasm/Options.td
+++ b/wasm/Options.td
@@ -215,6 +215,9 @@
 def import_table: FF<"import-table">,
   HelpText<"Import function table from the environment">;
 
+def initial_heap: JJ<"initial-heap=">,
+  HelpText<"Initial size of the heap">;
+
 def initial_memory: JJ<"initial-memory=">,
   HelpText<"Initial size of the linear memory">;
 
diff --git a/wasm/Writer.cpp b/wasm/Writer.cpp
index 0576bf2..805018c 100644
--- a/wasm/Writer.cpp
+++ b/wasm/Writer.cpp
@@ -439,6 +439,16 @@
     maxMemorySetting = 1ULL << 34;
   }
 
+  if (config->initialHeap != 0) {
+    if (config->initialHeap != alignTo(config->initialHeap, WasmPageSize))
+      error("initial heap must be " + Twine(WasmPageSize) + "-byte aligned");
+    uint64_t maxInitialHeap = maxMemorySetting - memoryPtr;
+    if (config->initialHeap > maxInitialHeap)
+      error("initial heap too large, cannot be greater than " +
+            Twine(maxInitialHeap));
+    memoryPtr += config->initialHeap;
+  }
+
   if (config->initialMemory != 0) {
     if (config->initialMemory != alignTo(config->initialMemory, WasmPageSize))
       error("initial memory must be " + Twine(WasmPageSize) + "-byte aligned");