[wasm-split] Split active segments (part 1)

src/ir/module-splitting.cpp

@@ -77,6 +77,7 @@
 #include "ir/find_all.h"
 #include "ir/module-utils.h"
 #include "ir/names.h"
+#include "support/stdckdint.h"
 #include "wasm-builder.h"
 #include "wasm.h"
 
@@ -363,6 +364,8 @@ struct ModuleSplitter {
     std::unordered_set<Name> memories;
     std::unordered_set<Name> tables;
     std::unordered_set<Name> tags;
+    std::unordered_set<Name> dataSegments;
+    std::unordered_set<Name> elementSegments;
   };
   using PrimarySecondaryUsedNames =
     std::pair<UsedNames, std::vector<UsedNames>>;
@@ -647,23 +650,6 @@ void ModuleSplitter::thunkExportedSecondaryFunctions() {
 }
 
 ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() {
-  auto walkSegments = [](auto& walker, Module* module) {
-    walker.setModule(module);
-    for (auto& curr : module->elementSegments) {
-      if (curr->offset) {
-        walker.walk(curr->offset);
-      }
-      for (auto* item : curr->data) {
-        walker.walk(item);
-      }
-    }
-    for (auto& curr : module->dataSegments) {
-      if (curr->offset) {
-        walker.walk(curr->offset);
-      }
-    }
-  };
-
   struct NameCollector
     : public PostWalker<NameCollector,
                         UnifiedExpressionVisitor<NameCollector>> {
@@ -699,9 +685,13 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() {
       case ModuleItemKind::Tag:                                                \
         used.tags.insert(cast->field);                                         \
         break;                                                                 \
-      case ModuleItemKind::Function:                                           \
       case ModuleItemKind::DataSegment:                                        \
+        used.dataSegments.insert(cast->field);                                 \
+        break;                                                                 \
       case ModuleItemKind::ElementSegment:                                     \
+        used.elementSegments.insert(cast->field);                              \
+        break;                                                                 \
+      case ModuleItemKind::Function:                                           \
       case ModuleItemKind::Invalid:                                            \
         break;                                                                 \
     }                                                                          \
@@ -726,32 +716,13 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() {
       used.memories.insert(funcUsed.memories.begin(), funcUsed.memories.end());
       used.tables.insert(funcUsed.tables.begin(), funcUsed.tables.end());
       used.tags.insert(funcUsed.tags.begin(), funcUsed.tags.end());
+      used.dataSegments.insert(funcUsed.dataSegments.begin(),
+                               funcUsed.dataSegments.end());
+      used.elementSegments.insert(funcUsed.elementSegments.begin(),
+                                  funcUsed.elementSegments.end());
     }
 
     NameCollector collector(used);
-    // We shouldn't use collector.walkModuleCode here, because we don't want to
-    // walk global initializers. At this point, all globals are still in the
-    // primary module, so if we walk global initializers here, other globals
-    // appearing in their initializers will all be marked as used in the primary
-    // module, which is not what we want.
-    //
-    // For example, we have (global $a i32 (global.get $b)). Because $a is at
-    // this point still in the primary module, $b will be marked as "used" in
-    // the primary module. But $a can be moved to a secondary module later if it
-    // is used exclusively by that module. Then $b can be also moved, in case it
-    // doesn't have other uses. But if it is marked as "used" in the primary
-    // module, it can't.
-    walkSegments(collector, &module);
-    for (auto& segment : module.dataSegments) {
-      if (segment->isActive()) {
-        used.memories.insert(segment->memory);
-      }
-    }
-    for (auto& segment : module.elementSegments) {
-      if (segment->isActive()) {
-        used.tables.insert(segment->table);
-      }
-    }
 
     return used;
   };
@@ -849,6 +820,101 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() {
     }
   }
 
+  auto mayTrap = [&](auto* segment) {
+    if constexpr (std::is_same_v<decltype(segment), ElementSegment*>) {
+      if (primary.features.hasCustomDescriptors()) {
+        for (auto* item : segment->data) {
+          if (EffectAnalyzer(config.passOptions, primary, item)
+                .hasUnremovableSideEffects()) {
+            return true;
+          }
+        }
+      }
+    }
+
+    // Check for out-of-bounds offset. This is adapted from maybeRootSegment
+    // function in RemoveUnusedModuleElements pass.
+    if (!config.passOptions.trapsNeverHappen) {
+      Index segmentSize;
+      Index parentSize;
+      if constexpr (std::is_same_v<decltype(segment), DataSegment*>) {
+        segmentSize = segment->data.size();
+        auto* memory = primary.getMemory(segment->memory);
+        parentSize = memory->initial << memory->pageSizeLog2;
+      } else {
+        segmentSize = segment->data.size();
+        auto* table = primary.getTable(segment->table);
+        parentSize = table->initial;
+      }
+
+      // Check if this might trap. If it is obviously in bounds then it cannot.
+      auto* c = segment->offset->template dynCast<Const>();
+      // Check for overflow in the largest possible space of addresses.
+      uint64_t maxWritten;
+      // If there is no integer, or if there is and the addition overflows, or
+      // if the addition leads to a too-large value, then we may trap.
+      if (!c ||
+          std::ckd_add(&maxWritten,
+                       (uint64_t)segmentSize,
+                       (uint64_t)c->value.getInteger()) ||
+          maxWritten > parentSize) {
+        return true;
+      }
+    }
+    return false;
+  };
+
+  // Iterate on active data and element segments. If its table or memory is
+  // used by a single secondary module, mark it "used" there. Only scan its
+  // 'offset' or 'data'(in case of ElementSegment) and add it to that module's
+  // used only when it is a sole secondary owner. If not assign it to the
+  // primary module and scan it there.
+  ModuleUtils::iterActiveDataSegments(primary, [&](DataSegment* segment) {
+    UsedNames* owner = getOwner(segment->memory, &UsedNames::memories);
+    if (mayTrap(segment)) {
+      owner = &primaryUsed;
+    }
+    if (!owner) {
+      return;
+    }
+    owner->dataSegments.insert(segment->name);
+    owner->memories.insert(segment->memory);
+    if (segment->offset) {
+      NameCollector(*owner).walk(segment->offset);
+    }
+  });
+
+  ModuleUtils::iterActiveElementSegments(primary, [&](ElementSegment* segment) {
+    UsedNames* owner = getOwner(segment->table, &UsedNames::tables);
+    if (mayTrap(segment)) {
+      owner = &primaryUsed;
+    }
+    if (!owner) {
+      return;
+    }
+    owner->elementSegments.insert(segment->name);
+    owner->tables.insert(segment->table);
+    if (segment->offset) {
+      NameCollector(*owner).walk(segment->offset);
+    }
+    for (auto* item : segment->data) {
+      NameCollector(*owner).walk(item);
+    }
+  });
+
+  // Passive element segments contain expressions (e.g. global.get) in their
+  // data arrays that must be scanned. Since currently functions referring to
+  // segments are forced into the primary module, passive segments always belong
+  // to the primary module.
+  for (auto& segment : primary.elementSegments) {
+    if (segment->isPassive() &&
+        primaryUsed.elementSegments.contains(segment->name)) {
+      for (auto* item : segment->data) {
+        NameCollector(primaryUsed).walk(item);
+      }
+    }
+  }
+
   // Compute the transitive closure of globals referenced in other globals'
   // initializers. WebAssembly validation requires that global initializers only
   // refer to previously defined globals. Therefore, `primary.globals` is
@@ -995,6 +1061,46 @@ void ModuleSplitter::shareImportableItems() {
   for (auto& name : tagsToRemove) {
     primary.removeTag(name);
   }
+
+  // Move segments that are exclusively used in a secondary module. If not, do
+  // nothing. (Segments cannot be imported / exported. They will be handled in
+  // indirectReferencesToSecondaryFunctions.)
+
+  std::vector<Name> dataSegmentsToRemove;
+  for (auto& dataSegment : primary.dataSegments) {
+    auto usingSecondaries =
+      getUsingSecondaries(dataSegment->name, &UsedNames::dataSegments);
+    bool inPrimary = primaryUsed.dataSegments.contains(dataSegment->name);
+
+    if (!inPrimary && usingSecondaries.empty()) {
+      dataSegmentsToRemove.push_back(dataSegment->name);
+    } else if (!inPrimary && usingSecondaries.size() == 1) {
+      auto* secondary = usingSecondaries[0];
+      ModuleUtils::copyDataSegment(dataSegment.get(), *secondary);
+      dataSegmentsToRemove.push_back(dataSegment->name);
+    }
+  }
+  for (auto& name : dataSegmentsToRemove) {
+    primary.removeDataSegment(name);
+  }
+
+  std::vector<Name> elementSegmentsToRemove;
+  for (auto& elementSegment : primary.elementSegments) {
+    auto usingSecondaries =
+      getUsingSecondaries(elementSegment->name, &UsedNames::elementSegments);
+    bool inPrimary = primaryUsed.elementSegments.contains(elementSegment->name);
+
+    if (!inPrimary && usingSecondaries.empty()) {
+      elementSegmentsToRemove.push_back(elementSegment->name);
+    } else if (!inPrimary && usingSecondaries.size() == 1) {
+      auto* secondary = usingSecondaries[0];
+      ModuleUtils::copyElementSegment(elementSegment.get(), *secondary);
+      elementSegmentsToRemove.push_back(elementSegment->name);
+    }
+  }
+  for (auto& name : elementSegmentsToRemove) {
+    primary.removeElementSegment(name);
+  }
 }
 
 void ModuleSplitter::indirectReferencesToSecondaryFunctions() {

test/lit/wasm-split/elem-global-deps.wast

@@ -0,0 +1,55 @@
+;; RUN: wasm-split -all -g %s --keep-funcs=keep -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
+  (global $g i32 (i32.const 42))
+  (table $t 1 1 funcref)
+
+  ;; This elem $e stays in the primary module because table $t is used both in
+  ;; the primary and secondary modules. So, this elem should be marked as "used"
+  ;; in the secondary module, and global $g should not be exported / imported.
+  (elem $e (table $t) (global.get $g) funcref (item (ref.null nofunc)))
+
+  (func $keep
+    (drop
+      (table.get $t
+        (i32.const 0)
+      )
+    )
+  )
+  (func $split
+    (drop
+      (table.get $t
+        (i32.const 0)
+      )
+    )
+  )
+)
+
+;; PRIMARY:      (module
+;; PRIMARY-NEXT:  (type $0 (func))
+;; PRIMARY-NEXT:  (global $g i32 (i32.const 42))
+;; PRIMARY-NEXT:  (table $t 1 1 funcref)
+;; PRIMARY-NEXT:  (elem $e (table $t) (global.get $g) funcref (item (ref.null nofunc)))
+;; PRIMARY-NEXT:  (export "table" (table $t))
+;; PRIMARY-NEXT:  (func $keep
+;; PRIMARY-NEXT:   (drop
+;; PRIMARY-NEXT:    (table.get $t
+;; PRIMARY-NEXT:     (i32.const 0)
+;; PRIMARY-NEXT:    )
+;; PRIMARY-NEXT:   )
+;; PRIMARY-NEXT:  )
+;; PRIMARY-NEXT: )
+
+;; SECONDARY:      (module
+;; SECONDARY-NEXT:  (type $0 (func))
+;; SECONDARY-NEXT:  (import "primary" "table" (table $t 1 1 funcref))
+;; SECONDARY-NEXT:  (func $split
+;; SECONDARY-NEXT:   (drop
+;; SECONDARY-NEXT:    (table.get $t
+;; SECONDARY-NEXT:     (i32.const 0)
+;; SECONDARY-NEXT:    )
+;; SECONDARY-NEXT:   )
+;; SECONDARY-NEXT:  )
+;; SECONDARY-NEXT: )

test/lit/wasm-split/passive-deps.wast

@@ -0,0 +1,38 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: wasm-split %s -all -g -o1 %t.1.wasm -o2 %t.2.wasm --split-funcs=split
+;; RUN: wasm-dis -all %t.1.wasm | filecheck %s --check-prefix PRIMARY
+;; RUN: wasm-dis -all %t.2.wasm | filecheck %s --check-prefix SECONDARY
+
+(module
+ ;; PRIMARY:      (type $0 (func))
+
+ ;; PRIMARY:      (global $g funcref (ref.null nofunc))
+ (global $g funcref (ref.null nofunc))
+
+ ;; We should scan this passive element segment's data and correctly mark $g as
+ ;; used in the primary module.
+ ;; PRIMARY:      (elem $passive-elem funcref (item (global.get $g)))
+ (elem $passive-elem funcref (item (global.get $g)))
+
+ ;; PRIMARY:      (export "global" (global $g))
+
+ ;; PRIMARY:      (func $keep (type $0)
+ ;; PRIMARY-NEXT:  (elem.drop $passive-elem)
+ ;; PRIMARY-NEXT: )
+ (func $keep
+  (elem.drop $passive-elem)
+ )
+
+ ;; SECONDARY:      (type $0 (func))
+
+ ;; SECONDARY:      (import "primary" "global" (global $g funcref))
+
+ ;; SECONDARY:      (func $split (type $0)
+ ;; SECONDARY-NEXT:  (drop
+ ;; SECONDARY-NEXT:   (global.get $g)
+ ;; SECONDARY-NEXT:  )
+ ;; SECONDARY-NEXT: )
+ (func $split
+  (drop (global.get $g))
+ )
+)