diff --git a/src/ir/lubs.cpp b/src/ir/lubs.cpp
index f9c04cb36e4..9d9beb06cff 100644
--- a/src/ir/lubs.cpp
+++ b/src/ir/lubs.cpp
@@ -15,7 +15,6 @@
  */
 
 #include "ir/lubs.h"
-#include "ir/find_all.h"
 #include "ir/utils.h"
 #include "wasm-type.h"
 #include "wasm.h"
@@ -51,51 +50,36 @@ LUBFinder getResultsLUB(Function* func, Module& wasm) {
     return lub;
   }
 
-  // Scan the body and look at the returns. First, return expressions.
-  for (auto* ret : FindAll<Return>(func->body).list) {
-    lub.note(ret->value->type);
-    if (lub.getLUB() == originalType) {
-      return lub;
-    }
-  }
+  // Scan the body and look at the returns.
+  struct Finder : public PostWalker<Finder> {
+    Module& wasm;
+    LUBFinder& lub;
 
-  // Process return_calls and call_refs. Unlike return expressions which we
-  // just handled, these only get a type to update, not a value.
-  auto processReturnType = [&](Type type) {
-    // Return whether we still look ok to do the optimization. If this is
-    // false then we can stop here.
-    lub.note(type);
-    return lub.getLUB() != originalType;
-  };
+    Finder(Module& wasm, LUBFinder& lub) : wasm(wasm), lub(lub) {}
 
-  for (auto* call : FindAll<Call>(func->body).list) {
-    if (call->isReturn &&
-        !processReturnType(wasm.getFunction(call->target)->getResults())) {
-      return lub;
-    }
-  }
-  for (auto* call : FindAll<CallIndirect>(func->body).list) {
-    if (call->isReturn &&
-        !processReturnType(call->heapType.getSignature().results)) {
-      return lub;
-    }
-  }
-  for (auto* call : FindAll<CallRef>(func->body).list) {
-    if (call->isReturn) {
-      auto targetType = call->target->type;
-      // We can skip unreachable code and calls to bottom types, as both trap.
-      if (targetType == Type::unreachable) {
-        continue;
+    void visitReturn(Return* curr) { lub.note(curr->value->type); }
+    void visitCall(Call* curr) {
+      if (curr->isReturn) {
+        lub.note(wasm.getFunction(curr->target)->getResults());
       }
-      auto targetHeapType = targetType.getHeapType();
-      if (targetHeapType.isBottom()) {
-        continue;
+    }
+    void visitCallIndirect(CallIndirect* curr) {
+      if (curr->isReturn) {
+        lub.note(curr->heapType.getSignature().results);
       }
-      if (!processReturnType(targetHeapType.getSignature().results)) {
-        return lub;
+    }
+    void visitCallRef(CallRef* curr) {
+      if (curr->isReturn) {
+        auto targetType = curr->target->type;
+        // We can skip unreachable code and calls to bottom types, as both trap.
+        if (!targetType.isSignature()) {
+          return;
+        }
+        lub.note(targetType.getHeapType().getSignature().results);
       }
     }
-  }
+  } finder(wasm, lub);
+  finder.walk(func->body);
 
   return lub;
 }