diff --git a/src/passes/HeapStoreOptimization.cpp b/src/passes/HeapStoreOptimization.cpp
index 566290c3643..f227dc53631 100644
--- a/src/passes/HeapStoreOptimization.cpp
+++ b/src/passes/HeapStoreOptimization.cpp
@@ -20,6 +20,7 @@
 // TODO: Add dead store elimination / load forwarding here.
 //
 
+#include "cfg/cfg-traversal.h"
 #include "ir/effects.h"
 #include "pass.h"
 #include "wasm-builder.h"
@@ -27,8 +28,17 @@
 
 namespace wasm {
 
+namespace {
+
+// In each basic block we will store the relevant heap store operations and
+// other actions that matter to our analysis.
+struct Info {
+  std::vector<Expression**> actions;
+};
+
 struct HeapStoreOptimization
-  : public WalkerPass<PostWalker<HeapStoreOptimization>> {
+  : public WalkerPass<
+      CFGWalker<HeapStoreOptimization, Visitor<HeapStoreOptimization>, Info>> {
   bool isFunctionParallel() override { return true; }
 
   // Locals are not modified here.
@@ -38,7 +48,39 @@ struct HeapStoreOptimization
     return std::make_unique<HeapStoreOptimization>();
   }
 
-  void visitStructSet(StructSet* curr) {
+  // Branches outside of the function can be ignored, as we only look at local
+  // state in the function. (This may need to change if we do more general dead
+  // store elimination.)
+  bool ignoreBranchesOutsideOfFunc = true;
+
+  // Store struct.sets and blocks, as we can find patterns among those.
+  void addAction() {
+    if (currBasicBlock) {
+      currBasicBlock->contents.actions.push_back(getCurrentPointer());
+    }
+  }
+  void visitStructSet(StructSet* curr) { addAction(); }
+  void visitBlock(Block* curr) { addAction(); }
+
+  void visitFunction(Function* curr) {
+    // Now that the walk is complete and we have a CFG, find things to optimize.
+    for (auto& block : basicBlocks) {
+      for (auto** currp : block->contents.actions) {
+        auto* curr = *currp;
+        if (auto* set = curr->dynCast<StructSet>()) {
+          optimizeStructSet(set, currp);
+        } else if (auto* block = curr->dynCast<Block>()) {
+          optimizeBlock(block);
+        } else {
+          WASM_UNREACHABLE("bad action");
+        }
+      }
+    }
+  }
+
+  // Optimize a struct.set. Receives also a pointer to where it is referred to,
+  // so we can replace it (which we do if we optimize).
+  void optimizeStructSet(StructSet* curr, Expression** currp) {
     // If our reference is a tee of a struct.new, we may be able to fold the
     // stored value into the new itself:
     //
@@ -52,7 +94,7 @@ struct HeapStoreOptimization
           // Success, so we do not need the struct.set any more, and the tee
           // can just be a set instead of us.
           tee->makeSet();
-          replaceCurrent(tee);
+          *currp = tee;
         }
       }
     }
@@ -69,9 +111,9 @@ struct HeapStoreOptimization
   // We also handle other struct.sets immediately after this one. If the
   // instruction following the new is not a struct.set we push the new down if
   // possible.
-  void visitBlock(Block* curr) { optimizeHeapStores(curr->list); }
+  void optimizeBlock(Block* curr) {
+    auto& list = curr->list;
 
-  void optimizeHeapStores(ExpressionList& list) {
     for (Index i = 0; i < list.size(); i++) {
       auto* localSet = list[i]->dynCast<LocalSet>();
       if (!localSet) {
@@ -234,6 +276,8 @@ struct HeapStoreOptimization
   }
 };
 
+} // anonymous namespace
+
 Pass* createHeapStoreOptimizationPass() { return new HeapStoreOptimization(); }
 
 } // namespace wasm