[wasm-split] Do not split out functions referring to segments

src/ir/module-splitting.cpp

@@ -282,7 +282,7 @@ struct ModuleSplitter {
   // Initialization helpers
   static std::unique_ptr<Module> initSecondary(const Module& primary);
   static std::pair<std::set<Name>, std::set<Name>>
-  classifyFunctions(const Module& primary, const Config& config);
+  classifyFunctions(Module& primary, const Config& config);
   static std::map<Name, Name> initExportedPrimaryFuncs(const Module& primary);
 
   // Other helpers
@@ -343,22 +343,80 @@ std::unique_ptr<Module> ModuleSplitter::initSecondary(const Module& primary) {
 }
 
 std::pair<std::set<Name>, std::set<Name>>
-ModuleSplitter::classifyFunctions(const Module& primary, const Config& config) {
+ModuleSplitter::classifyFunctions(Module& primary, const Config& config) {
+  // Find functions that refer to data or element segments. These functions must
+  // remain in the primary module because segments cannot be exported to be
+  // accessed from the secondary module.
+  //
+  // TODO: Investigate other options, such as moving the segments to the
+  // secondary module or replacing the segment-using instructions in the
+  // secondary module with calls to imports.
+  ModuleUtils::ParallelFunctionAnalysis<std::vector<Name>>
+    segmentReferrerCollector(
+      primary, [&](Function* func, std::vector<Name>& segmentReferrers) {
+        if (func->imported()) {
+          return;
+        }
+
+        struct SegmentReferrerCollector
+          : PostWalker<SegmentReferrerCollector,
+                       UnifiedExpressionVisitor<SegmentReferrerCollector>> {
+          bool hasSegmentReference = false;
+
+          void visitExpression(Expression* curr) {
+
+#define DELEGATE_ID curr->_id
+
+#define DELEGATE_START(id) [[maybe_unused]] auto* cast = curr->cast<id>();
+#define DELEGATE_GET_FIELD(id, field) cast->field
+#define DELEGATE_FIELD_TYPE(id, field)
+#define DELEGATE_FIELD_HEAPTYPE(id, field)
+#define DELEGATE_FIELD_CHILD(id, field)
+#define DELEGATE_FIELD_OPTIONAL_CHILD(id, field)
+#define DELEGATE_FIELD_INT(id, field)
+#define DELEGATE_FIELD_LITERAL(id, field)
+#define DELEGATE_FIELD_NAME(id, field)
+#define DELEGATE_FIELD_SCOPE_NAME_DEF(id, field)
+#define DELEGATE_FIELD_SCOPE_NAME_USE(id, field)
+#define DELEGATE_FIELD_ADDRESS(id, field)
+
+#define DELEGATE_FIELD_NAME_KIND(id, field, kind)                              \
+  if (kind == ModuleItemKind::DataSegment ||                                   \
+      kind == ModuleItemKind::ElementSegment) {                                \
+    hasSegmentReference = true;                                                \
+  }
+
+#include "wasm-delegations-fields.def"
+          }
+        };
+        SegmentReferrerCollector collector;
+        collector.walkFunction(func);
+        if (collector.hasSegmentReference) {
+          segmentReferrers.push_back(func->name);
+        }
+      });
+
+  std::unordered_set<Name> segmentReferrers;
+  for (auto& [_, referrers] : segmentReferrerCollector.map) {
+    segmentReferrers.insert(referrers.begin(), referrers.end());
+  }
+
   std::set<Name> primaryFuncs, secondaryFuncs;
   for (auto& func : primary.functions) {
     // In JSPI mode exported functions cannot be moved to the secondary
     // module since that would make them async when they may not have the JSPI
     // wrapper. Exported JSPI functions can still benefit from splitting though
     // since only the JSPI wrapper stub will remain in the primary module.
     if (func->imported() || config.primaryFuncs.count(func->name) ||
-        (config.jspi && ExportUtils::isExported(primary, *func))) {
+        (config.jspi && ExportUtils::isExported(primary, *func)) ||
+        segmentReferrers.count(func->name)) {
       primaryFuncs.insert(func->name);
     } else {
       assert(func->name != primary.start && "The start function must be kept");
       secondaryFuncs.insert(func->name);
     }
   }
-  return std::make_pair(primaryFuncs, secondaryFuncs);
+  return std::make_pair(std::move(primaryFuncs), std::move(secondaryFuncs));
 }
 
 std::map<Name, Name>

test/lit/wasm-split/segments.wast

@@ -0,0 +1,94 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-split %s -all --keep-funcs=foo -g -o1 %t.1.wasm -o2 %t.2.wasm
+;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY
+;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY
+
+(module
+
+ ;; PRIMARY:      (type $0 (func))
+
+ ;; PRIMARY:      (type $data-array (array i8))
+ (type $data-array (array i8))
+
+ ;; PRIMARY:      (type $elem-array (array externref))
+ (type $elem-array (array externref))
+
+ ;; PRIMARY:      (memory $mem 0)
+ (memory $mem 0)
+
+ ;; PRIMARY:      (data $data "hello world")
+ (data $data "hello world")
+
+ ;; PRIMARY:      (elem $elem externref)
+ (elem $elem externref)
+
+ ;; PRIMARY:      (export "memory" (memory $mem))
+
+ ;; PRIMARY:      (func $data.drop
+ ;; PRIMARY-NEXT:  (data.drop $data)
+ ;; PRIMARY-NEXT: )
+ (func $data.drop
+  (data.drop $data)
+ )
+
+ ;; PRIMARY:      (func $memory.init
+ ;; PRIMARY-NEXT:  (memory.init $data
+ ;; PRIMARY-NEXT:   (i32.const 0)
+ ;; PRIMARY-NEXT:   (i32.const 0)
+ ;; PRIMARY-NEXT:   (i32.const 0)
+ ;; PRIMARY-NEXT:  )
+ ;; PRIMARY-NEXT: )
+ (func $memory.init
+  (memory.init $mem $data
+   (i32.const 0)
+   (i32.const 0)
+   (i32.const 0)
+  )
+ )
+
+ ;; PRIMARY:      (func $array.new_data
+ ;; PRIMARY-NEXT:  (drop
+ ;; PRIMARY-NEXT:   (array.new_data $data-array $data
+ ;; PRIMARY-NEXT:    (i32.const 0)
+ ;; PRIMARY-NEXT:    (i32.const 0)
+ ;; PRIMARY-NEXT:   )
+ ;; PRIMARY-NEXT:  )
+ ;; PRIMARY-NEXT: )
+ (func $array.new_data
+  (drop
+   (array.new_data $data-array $data
+    (i32.const 0)
+    (i32.const 0)
+   )
+  )
+ )
+
+ ;; PRIMARY:      (func $array.new_elem
+ ;; PRIMARY-NEXT:  (drop
+ ;; PRIMARY-NEXT:   (array.new_elem $elem-array $elem
+ ;; PRIMARY-NEXT:    (i32.const 0)
+ ;; PRIMARY-NEXT:    (i32.const 0)
+ ;; PRIMARY-NEXT:   )
+ ;; PRIMARY-NEXT:  )
+ ;; PRIMARY-NEXT: )
+ (func $array.new_elem
+  (drop
+   (array.new_elem $elem-array $elem
+    (i32.const 0)
+    (i32.const 0)
+   )
+  )
+ )
+
+ ;; SECONDARY:      (type $0 (func))
+
+ ;; SECONDARY:      (import "primary" "memory" (memory $mem 0))
+
+ ;; SECONDARY:      (func $no-segment
+ ;; SECONDARY-NEXT:  (nop)
+ ;; SECONDARY-NEXT: )
+ (func $no-segment
+  (nop)
+ )
+)