MemoryPacking: Preserve segment names (#3458)

Also, avoid packing builtin llvm segments names so that
segments such as `__llvm_covfun` (use by llvm-cov) are
preserved in the final output.

6 files changed, 49 insertions, 22 deletions

diff --git a/src/binaryen-c.cpp b/src/binaryen-c.cpp
index 29724ec10..7b2ef9214 100644
--- a/src/binaryen-c.cpp
+++ b/src/binaryen-c.cpp
@@ -3274,7 +3274,8 @@ void BinaryenSetMemory(BinaryenModuleRef module,
     wasm->addExport(memoryExport.release());
   }
   for (BinaryenIndex i = 0; i < numSegments; i++) {
-    wasm->memory.segments.emplace_back(segmentPassive[i],
+    wasm->memory.segments.emplace_back(Name(),
+                                       segmentPassive[i],
                                        (Expression*)segmentOffsets[i],
                                        segments[i],
                                        segmentSizes[i]);
diff --git a/src/passes/MemoryPacking.cpp b/src/passes/MemoryPacking.cpp
index b6ee7acf9..7848118b8 100644
--- a/src/passes/MemoryPacking.cpp
+++ b/src/passes/MemoryPacking.cpp
@@ -156,6 +156,7 @@ void MemoryPacking::run(PassRunner* runner, Module* module) {
     auto& currReferrers = referrers[origIndex];
 
     std::vector<Range> ranges;
+
     if (canSplit(segment, currReferrers)) {
       calculateRanges(segment, currReferrers, ranges);
     } else {
@@ -252,6 +253,14 @@ bool MemoryPacking::canOptimize(const Memory& memory,
 
 bool MemoryPacking::canSplit(const Memory::Segment& segment,
                              const Referrers& referrers) {
+  // Don't mess with segments related to llvm coverage tools such as
+  // __llvm_covfun. There segments are expected/parsed by external downstream
+  // tools (llvm-cov) so they need to be left intact.
+  // See https://clang.llvm.org/docs/SourceBasedCodeCoverage.html
+  if (segment.name.is() && segment.name.startsWith("__llvm")) {
+    return false;
+  }
+
   if (segment.isPassive) {
     for (auto* referrer : referrers) {
       if (auto* init = referrer->dynCast<MemoryInit>()) {
@@ -262,10 +271,10 @@ bool MemoryPacking::canSplit(const Memory::Segment& segment,
       }
     }
     return true;
-  } else {
-    // Active segments can only be split if they have constant offsets
-    return segment.offset->is<Const>();
   }
+
+  // Active segments can only be split if they have constant offsets
+  return segment.offset->is<Const>();
 }
 
 void MemoryPacking::calculateRanges(const Memory::Segment& segment,
@@ -505,6 +514,7 @@ void MemoryPacking::createSplitSegments(Builder& builder,
                                         std::vector<Range>& ranges,
                                         std::vector<Memory::Segment>& packed,
                                         size_t segmentsRemaining) {
+  size_t segmentCount = 0;
   for (size_t i = 0; i < ranges.size(); ++i) {
     Range& range = ranges[i];
     if (range.isZero) {
@@ -528,7 +538,22 @@ void MemoryPacking::createSplitSegments(Builder& builder,
       range.end = lastNonzero->end;
       ranges.erase(ranges.begin() + i + 1, lastNonzero + 1);
     }
-    packed.emplace_back(segment.isPassive,
+    Name name;
+    if (segment.name.is()) {
+      // Name the first range after the original segment and all following
+      // ranges get numbered accordingly.  This means that for segments that
+      // canot be split (segments that contains a single range) the input and
+      // output segment have the same name.
+      if (!segmentCount) {
+        name = segment.name;
+      } else {
+        name = std::string(segment.name.c_str()) + "." +
+               std::to_string(segmentCount);
+      }
+      segmentCount++;
+    }
+    packed.emplace_back(name,
+                        segment.isPassive,
                         offset,
                         &segment.data[range.start],
                         range.end - range.start);
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 452906ab4..43841b873 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -288,7 +288,8 @@ private:
   void stringToBinary(const char* input, size_t size, std::vector<char>& data);
   void parseMemory(Element& s, bool preParseImport = false);
   void parseData(Element& s);
-  void parseInnerData(Element& s, Index i, Expression* offset, bool isPassive);
+  void parseInnerData(
+    Element& s, Index i, Name name, Expression* offset, bool isPassive);
   void parseExport(Element& s);
   void parseImport(Element& s);
   void parseGlobal(Element& s, bool preParseImport = false);
diff --git a/src/wasm.h b/src/wasm.h
index 255f18a16..1d5642625 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -1704,8 +1704,12 @@ public:
     Segment(Expression* offset, std::vector<char>& init) : offset(offset) {
       data.swap(init);
     }
-    Segment(bool isPassive, Expression* offset, const char* init, Address size)
-      : isPassive(isPassive), offset(offset) {
+    Segment(Name name,
+            bool isPassive,
+            Expression* offset,
+            const char* init,
+            Address size)
+      : name(name), isPassive(isPassive), offset(offset) {
       data.resize(size);
       std::copy_n(init, size, data.begin());
     }
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index 32b62d0e1..20e99436a 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -2337,7 +2337,7 @@ void SExpressionWasmBuilder::parseMemory(Element& s, bool preParseImport) {
       // (memory (data ..)) format
       auto j = parseMemoryIndex(inner, 1);
       auto offset = allocator.alloc<Const>()->set(Literal(int32_t(0)));
-      parseInnerData(inner, j, offset, false);
+      parseInnerData(inner, j, {}, offset, false);
       wasm.memory.initial = wasm.memory.segments[0].data.size();
       return;
     }
@@ -2399,16 +2399,11 @@ void SExpressionWasmBuilder::parseData(Element& s) {
   if (s.size() != 3 && s.size() != 4) {
     throw ParseException("Unexpected data items", s.line, s.col);
   }
-  parseInnerData(s, s.size() - 1, offset, isPassive);
-  if (name.is()) {
-    wasm.memory.segments.back().name = name;
-  }
+  parseInnerData(s, s.size() - 1, name, offset, isPassive);
 }
 
-void SExpressionWasmBuilder::parseInnerData(Element& s,
-                                            Index i,
-                                            Expression* offset,
-                                            bool isPassive) {
+void SExpressionWasmBuilder::parseInnerData(
+  Element& s, Index i, Name name, Expression* offset, bool isPassive) {
   std::vector<char> data;
   while (i < s.size()) {
     const char* input = s[i++]->c_str();
@@ -2417,7 +2412,7 @@ void SExpressionWasmBuilder::parseInnerData(Element& s,
     }
   }
   wasm.memory.segments.emplace_back(
-    isPassive, offset, data.data(), data.size());
+    name, isPassive, offset, data.data(), data.size());
 }
 
 void SExpressionWasmBuilder::parseExport(Element& s) {
diff --git a/test/unit/test_memory_packing.py b/test/unit/test_memory_packing.py
index 357752f10..e2468e689 100644
--- a/test/unit/test_memory_packing.py
+++ b/test/unit/test_memory_packing.py
@@ -13,7 +13,7 @@ class MemoryPackingTest(utils.BinaryenTestCase):
         module = '''
         (module
          (memory 256 256)
-         (data (i32.const 0) %s)
+         (data $d (i32.const 0) %s)
         )
         ''' % data
         opts = ['--memory-packing', '--disable-bulk-memory', '--print',
@@ -21,9 +21,10 @@ class MemoryPackingTest(utils.BinaryenTestCase):
         p = shared.run_process(shared.WASM_OPT + opts, input=module,
                                check=False, capture_output=True)
         output = [
-            '(data (i32.const 999970) "A")',
-            '(data (i32.const 999980) "A")',
-            '(data (i32.const 999990) "A' + ('\\00' * 9) + 'A")'
+            '(data $d (i32.const 0) "A")',
+            '(data $d.1 (i32.const 10) "A")',
+            '(data $d.99998 (i32.const 999980) "A")',
+            '(data $d.99999 (i32.const 999990) "A' + ('\\00' * 9) + 'A")'
         ]
         self.assertEqual(p.returncode, 0)
         for line in output: