3 files changed, 48 insertions, 25 deletions

diff --git a/src/passes/SafeHeap.cpp b/src/passes/SafeHeap.cpp
index a90696f2c..068c8ef73 100644
--- a/src/passes/SafeHeap.cpp
+++ b/src/passes/SafeHeap.cpp
@@ -22,6 +22,7 @@
 
 #include "asmjs/shared-constants.h"
 #include "ir/bits.h"
+#include "ir/find_all.h"
 #include "ir/import-utils.h"
 #include "ir/load-utils.h"
 #include "pass.h"
@@ -100,11 +101,30 @@ struct AccessInstrumenter : public WalkerPass<PostWalker<AccessInstrumenter>> {
   }
 };
 
-struct FindDirectCallees : public WalkerPass<PostWalker<FindDirectCallees>> {
-public:
-  void visitCall(Call* curr) { callees.insert(curr->target); }
-  std::set<Name> callees;
-};
+static std::set<Name> findCalledFunctions(Module* module, Name startFunc) {
+  std::set<Name> called;
+  std::vector<Name> toVisit;
+
+  auto addFunction = [&](Name name) {
+    if (called.insert(name).second) {
+      toVisit.push_back(name);
+    }
+  };
+
+  if (startFunc.is()) {
+    addFunction(startFunc);
+    while (!toVisit.empty()) {
+      auto next = toVisit.back();
+      toVisit.pop_back();
+      auto* func = module->getFunction(next);
+      for (auto* call : FindAll<Call>(func->body).list) {
+        addFunction(call->target);
+      }
+    }
+  }
+
+  return called;
+}
 
 struct SafeHeap : public Pass {
   PassOptions options;
@@ -114,24 +134,13 @@ struct SafeHeap : public Pass {
     // add imports
     addImports(module);
     // instrument loads and stores
-    // We avoid instrumenting the module start function of any function
-    // that it directly calls.  This is because in some cases the linker
-    // generates `__wasm_init_memory` (either as the start function or
-    // a function directly called from it) and this function is used in shared
-    // memory builds to load the passive memory segments, which in turn means
-    // that value of sbrk() is not available until after it has run.
-    std::set<Name> ignoreFunctions;
-    if (module->start.is()) {
-      // Note that this only finds directly called functions, not transitively
-      // called ones.  That is enough given the current LLVM output as start
-      // will only contain very specific, linker-generated code
-      // (__wasm_init_memory etc. as mentioned above).
-      FindDirectCallees findDirectCallees;
-      findDirectCallees.walkFunctionInModule(module->getFunction(module->start),
-                                             module);
-      ignoreFunctions = findDirectCallees.callees;
-      ignoreFunctions.insert(module->start);
-    }
+    // We avoid instrumenting the module start function of any function that it
+    // directly calls.  This is because in some cases the linker generates
+    // `__wasm_init_memory` (either as the start function or a function directly
+    // called from it) and this function is used in shared memory builds to load
+    // the passive memory segments, which in turn means that value of sbrk() is
+    // not available until after it has run.
+    std::set<Name> ignoreFunctions = findCalledFunctions(module, module->start);
     ignoreFunctions.insert(getSbrkPtr);
     AccessInstrumenter(ignoreFunctions).run(runner, module);
     // add helper checking funcs and imports
diff --git a/test/passes/safe-heap_start-function.txt b/test/passes/safe-heap_start-function.txt
index 58bb77d44..7670530cc 100644
--- a/test/passes/safe-heap_start-function.txt
+++ b/test/passes/safe-heap_start-function.txt
@@ -2,8 +2,8 @@
  (type $i32_i32_=>_i64 (func (param i32 i32) (result i64)))
  (type $i32_i32_i64_=>_none (func (param i32 i32 i64)))
  (type $i32_i32_=>_i32 (func (param i32 i32) (result i32)))
- (type $i32_i32_i32_=>_none (func (param i32 i32 i32)))
  (type $none_=>_none (func))
+ (type $i32_i32_i32_=>_none (func (param i32 i32 i32)))
  (type $i32_i32_=>_f64 (func (param i32 i32) (result f64)))
  (type $i32_i32_f64_=>_none (func (param i32 i32 f64)))
  (type $i32_i32_=>_f32 (func (param i32 i32) (result f32)))
@@ -30,6 +30,15 @@
    )
    (i32.const 5678)
   )
+  (call $foo2)
+ )
+ (func $foo2
+  (i32.store
+   (i32.load
+    (i32.const 98)
+   )
+   (i32.const 99)
+  )
  )
  (func $bar
   (call $SAFE_HEAP_STORE_i32_4_4
diff --git a/test/passes/safe-heap_start-function.wast b/test/passes/safe-heap_start-function.wast
index 7fc2d5201..ae6525c68 100644
--- a/test/passes/safe-heap_start-function.wast
+++ b/test/passes/safe-heap_start-function.wast
@@ -6,8 +6,13 @@
    (call $foo)
   )
   (func $foo
-   ;; should not be modified because its called from the start function
+   ;; should not be modified because its reachable from start function
    (i32.store (i32.load (i32.const 1234)) (i32.const 5678))
+   (call $foo2)
+  )
+  (func $foo2
+   ;; should not be modified because its reachable from start function
+   (i32.store (i32.load (i32.const 98)) (i32.const 99))
   )
   (func $bar
    (i32.store (i32.load (i32.const 1234)) (i32.const 5678))