2 files changed, 42 insertions, 19 deletions

diff --git a/src/passes/LogExecution.cpp b/src/passes/LogExecution.cpp
index 65a802a2f..abdaa8d23 100644
--- a/src/passes/LogExecution.cpp
+++ b/src/passes/LogExecution.cpp
@@ -16,13 +16,17 @@
 
 //
 // Instruments the build with code to log execution at each function
-// entry and loop header. This can be useful in debugging, to log out
+// entry, loop header, and return. This can be useful in debugging, to log out
 // a trace, and diff it to another (running in another browser, to
 // check for bugs, for example).
 //
 // The logging is performed by calling an ffi with an id for each
 // call site. You need to provide that import on the JS side.
 //
+// This pass is more effective on flat IR (--flatten) since when it
+// instruments say a return, there will be no code run in the return's
+// value.
+//
 
 #include <wasm.h>
 #include <wasm-builder.h>
@@ -41,10 +45,19 @@ struct LogExecution : public WalkerPass<PostWalker<LogExecution>> {
     curr->body = makeLogCall(curr->body);
   }
 
+  void visitReturn(Return* curr) {
+    replaceCurrent(makeLogCall(curr));
+  }
+
   void visitFunction(Function* curr) {
     if (curr->imported()) {
       return;
     }
+    if (auto* block = curr->body->dynCast<Block>()) {
+      if (!block->list.empty()) {
+        block->list.back() = makeLogCall(block->list.back());
+      }
+    }
     curr->body = makeLogCall(curr->body);
   }
 
diff --git a/test/passes/log-execution.txt b/test/passes/log-execution.txt
index 3f842a6f1..a7e83f1e8 100644
--- a/test/passes/log-execution.txt
+++ b/test/passes/log-execution.txt
@@ -18,26 +18,31 @@
  )
  (func $workk (; 4 ;) (type $FUNCSIG$v)
   (call $log_execution
-   (i32.const 2)
+   (i32.const 3)
   )
   (block
    (if
     (i32.const 0)
     (nop)
    )
-   (drop
-    (i32.const 1)
+   (block
+    (call $log_execution
+     (i32.const 2)
+    )
+    (drop
+     (i32.const 1)
+    )
    )
   )
  )
  (func $loops (; 5 ;) (type $FUNCSIG$v)
   (call $log_execution
-   (i32.const 6)
+   (i32.const 8)
   )
   (block
    (loop $x
     (call $log_execution
-     (i32.const 3)
+     (i32.const 4)
     )
     (block
      (call $loops)
@@ -48,24 +53,29 @@
     (call $intt)
     (loop $y
      (call $log_execution
-      (i32.const 4)
+      (i32.const 5)
      )
      (call $loops)
     )
    )
-   (loop $loop-in
+   (block
     (call $log_execution
-     (i32.const 5)
+     (i32.const 7)
     )
-    (block
-     (drop
-      (i32.const 10)
-     )
-     (drop
-      (i32.const 20)
+    (loop $loop-in
+     (call $log_execution
+      (i32.const 6)
      )
-     (drop
-      (i32.const 30)
+     (block
+      (drop
+       (i32.const 10)
+      )
+      (drop
+       (i32.const 20)
+      )
+      (drop
+       (i32.const 30)
+      )
      )
     )
    )
@@ -73,11 +83,11 @@
  )
  (func $loops-similar (; 6 ;) (type $FUNCSIG$v)
   (call $log_execution
-   (i32.const 8)
+   (i32.const 10)
   )
   (loop $x
    (call $log_execution
-    (i32.const 7)
+    (i32.const 9)
    )
    (block
     (call $loops)