finalize: Strip segments that contains on EM_ASM/EM_JS data

If we find a data segment whose entire contents is EM_JS or EM_ASM
strings then strip it from the binary.

See: emscripten-core/emscripten#13443

Author: sbc100

src/binaryen-c.h

@@ -1625,9 +1625,9 @@ BINARYEN_API void BinaryenMemoryInitSetSize(BinaryenExpressionRef expr,
 
 // DataDrop
 
-// Gets the index of the segment being dropped by a `memory.drop` expression.
+// Gets the index of the segment being dropped by a `data.drop` expression.
 BINARYEN_API uint32_t BinaryenDataDropGetSegment(BinaryenExpressionRef expr);
-// Sets the index of the segment being dropped by a `memory.drop` expression.
+// Sets the index of the segment being dropped by a `data.drop` expression.
 BINARYEN_API void BinaryenDataDropSetSegment(BinaryenExpressionRef expr,
                                              uint32_t segmentIndex);
 

src/wasm/wasm-emscripten.cpp

@@ -111,13 +111,16 @@ class StringConstantTracker {
     : wasm(wasm) { calcSegmentOffsets(); }
 
   std::string codeForConstAddr(int64_t address) {
-    const char* str = stringAtAddr(address);
+    char* str = stringAtAddr(address);
     if (!str) {
       Fatal() << "unable to find data for ASM/EM_JS const at: " << address;
     }
     return escape(str);
   }
 
+  std::set<Index> segmentsUsed;
+  std::vector<Address> segmentOffsets; // segment index => address offset
+
 private:
   void calcSegmentOffsets() {
     std::unordered_map<Index, Address> passiveOffsets;
@@ -173,20 +176,20 @@ class StringConstantTracker {
     }
   }
 
-  const char* stringAtAddr(Address address) {
+  char* stringAtAddr(Address address) {
     for (unsigned i = 0; i < wasm.memory.segments.size(); ++i) {
       Memory::Segment& segment = wasm.memory.segments[i];
       Address offset = segmentOffsets[i];
       if (offset != UNKNOWN_OFFSET && address >= offset &&
           address < offset + segment.data.size()) {
+        segmentsUsed.insert(i);
         return &segment.data[address - offset];
       }
     }
     return nullptr;
   }
 
   Module& wasm;
-  std::vector<Address> segmentOffsets; // segment index => address offset
 };
 
 enum class Proxying {
@@ -381,6 +384,7 @@ struct EmJsWalker : public PostWalker<EmJsWalker> {
   std::vector<Export> toRemove;
 
   std::map<std::string, std::string> codeByName;
+  std::map<Address, size_t> codeAddresses; // map from address to string len
 
   EmJsWalker(Module& _wasm)
     : wasm(_wasm), stringTracker(_wasm) {}
@@ -406,9 +410,34 @@ struct EmJsWalker : public PostWalker<EmJsWalker> {
     int64_t address = addrConst->value.getInteger();
     auto code = stringTracker.codeForConstAddr(address);
     codeByName[funcName] = code;
+    codeAddresses[address] = code.size() + 1;
   }
 };
 
+struct SegmentRemover : WalkerPass<PostWalker<SegmentRemover>> {
+  SegmentRemover(Index segment) : segment(segment) {}
+
+  bool isFunctionParallel() override { return true; }
+
+  Pass* create() override { return new SegmentRemover(segment); }
+
+  void visitMemoryInit(MemoryInit* curr) {
+    if (segment == curr->segment) {
+      Builder builder(*getModule());
+      replaceCurrent(builder.makeNop());
+    }
+  }
+
+  void visitDataDrop(DataDrop* curr) {
+    if (segment == curr->segment) {
+      Builder builder(*getModule());
+      replaceCurrent(builder.makeNop());
+    }
+  }
+
+  Index segment;
+};
+
 EmJsWalker fixEmJsFuncsAndReturnWalker(Module& wasm) {
   EmJsWalker walker(wasm);
   walker.walkModule(&wasm);
@@ -417,6 +446,31 @@ EmJsWalker fixEmJsFuncsAndReturnWalker(Module& wasm) {
     wasm.removeExport(exp.name);
     wasm.removeFunction(exp.value);
   }
+
+
+  for (Index segIndex: walker.stringTracker.segmentsUsed) {
+    Address start = walker.stringTracker.segmentOffsets[0];
+    Address cur = start;
+
+    // With newer versions of emscripten/llvm we pack all EM_JS strings into
+    // single segment.
+    // We can detect this by checking if the strings found cover the entire
+    // segment
+    while (cur < start + wasm.memory.segments[segIndex].data.size()) {
+      if (walker.codeAddresses.count(cur) == 0)
+        break;
+      cur.addr += walker.codeAddresses[cur];
+    }
+
+    if (cur == start + wasm.memory.segments[segIndex].data.size()) {
+      PassRunner runner(&wasm);
+      SegmentRemover(segIndex).run(&runner, &wasm);
+      // Resize the segment to zero.  In theory we should completely remove it
+      // but that would mean re-numbering the segments that follow which would
+      // mean renumbering.
+      wasm.memory.segments[segIndex].data.resize(0);
+    }
+  }
   return walker;
 }
 

test/lit/wasm-emscripten-finalize/em_js.wat

@@ -3,20 +3,32 @@
 
 ;; RUN: wasm-emscripten-finalize %s -S | filecheck %s
 
-;; Both functions should be stripped from the binary
+;; All functions should be stripped from the binary, regardless
+;; of internal name
 ;; CHECK-NOT:  (func
 
+;; The data section that contains only em_js strings should
+;; be stripped.
+;; CHECK-NOT: (i32.const 512) "Only em_js strings here\00")
+
+;; Data sections that also contain other stuff should not be stripped
+;; CHECK: (data (i32.const 1024) "some JS string data\00xxx")
+;; CHECK: (data (i32.const 2048) "more JS string data\00yyy")
+
 ;;      CHECK:  "emJsFuncs": {
-;; CHECK-NEXT:    "bar": "more JS string dara",
-;; CHECK-NEXT:    "foo": "some JS string"
+;; CHECK-NEXT:    "bar": "more JS string data",
+;; CHECK-NEXT:    "baz": "Only em_js strings here
+;; CHECK-NEXT:    "foo": "some JS string data"
 ;; CHECK-NEXT:  },
 
 (module
  (memory 1 1)
- (data (i32.const 1024) "some JS string\00")
- (data (i32.const 2048) "more JS string dara\00")
+ (data (i32.const 512) "Only em_js strings here\00")
+ (data (i32.const 1024) "some JS string data\00xxx")
+ (data (i32.const 2048) "more JS string data\00yyy")
  (export "__em_js__foo" (func $__em_js__foo))
  (export "__em_js__bar" (func $bar))
+ (export "__em_js__baz" (func $baz))
  ;; Name matches export name
  (func $__em_js__foo (result i32)
   (i32.const 1024)
@@ -25,4 +37,7 @@
  (func $bar (result i32)
   (i32.const 2048)
  )
+ (func $baz (result i32)
+  (i32.const 512)
+ )
 )