--- a/js/src/wasm/WasmBaselineCompile.cpp
+++ b/js/src/wasm/WasmBaselineCompile.cpp
@@ -151,16 +151,17 @@ struct BaseCompilePolicy : OpIterPolicy
 
 typedef OpIter<BaseCompilePolicy> BaseOpIter;
 
 typedef bool IsUnsigned;
 typedef bool IsSigned;
 typedef bool ZeroOnOverflow;
 typedef bool IsKnownNotZero;
 typedef bool HandleNaNSpecially;
+typedef bool PopStack;
 typedef unsigned ByteSize;
 typedef unsigned BitSize;
 
 // UseABI::Wasm implies that the Tls/Heap/Global registers are nonvolatile,
 // except when InterModule::True is also set, when they are volatile.
 //
 // UseABI::System implies that the Tls/Heap/Global registers are volatile.
 // Additionally, the parameter passing mechanism may be slightly different from
@@ -1929,19 +1930,16 @@ class BaseCompiler
     }
 
     void popControl() {
         Control last = ctl_.popCopy();
         if (last.label)
             freeLabel(last.label);
         if (last.otherLabel)
             freeLabel(last.otherLabel);
-
-        if (deadCode_ && !ctl_.empty())
-            popValueStackTo(ctl_.back().stackSize);
     }
 
     Control& controlItem(uint32_t relativeDepth) {
         return ctl_[ctl_.length() - 1 - relativeDepth];
     }
 
     MOZ_MUST_USE PooledLabel* newLabel() {
         // TODO / INVESTIGATE (Bug 1316819): allocate() is fallible, but we can
@@ -2472,18 +2470,19 @@ class BaseCompiler
         }
 #elif defined(JS_CODEGEN_ARM)
         if (call.usesSystemAbi && !call.hardFP)
             masm.ma_vxfer(r0, r1, rv.reg);
 #endif
         return rv;
     }
 
-    void returnCleanup() {
-        popStackBeforeBranch(ctl_[0].framePushed);
+    void returnCleanup(bool popStack) {
+        if (popStack)
+            popStackBeforeBranch(ctl_[0].framePushed);
         masm.jump(&returnLabel_);
     }
 
     void pop2xI32ForIntMulDiv(RegI32* r0, RegI32* r1) {
 #if defined(JS_CODEGEN_X86) || defined(JS_CODEGEN_X64)
         // srcDest must be eax, and edx will be clobbered.
         need2xI32(specific_eax, specific_edx);
         *r1 = popI32();
@@ -3641,22 +3640,22 @@ class BaseCompiler
 
     template<bool isSetLocal> MOZ_MUST_USE inline bool emitSetOrTeeLocal(uint32_t slot);
     template<bool isSetGlobal> MOZ_MUST_USE inline bool emitSetOrTeeGlobal(uint32_t id);
     template<bool isStore>
     MOZ_MUST_USE inline bool emitStoreOrTeeStore(ValType resultType,
                                                  Scalar::Type viewType,
                                                  LinearMemoryAddress<Nothing> addr);
 
-    void endBlock(ExprType type, bool isFunctionBody);
+    void endBlock(ExprType type);
     void endLoop(ExprType type);
     void endIfThen();
     void endIfThenElse(ExprType type);
 
-    void doReturn(ExprType returnType);
+    void doReturn(ExprType returnType, bool popStack);
     void pushReturned(const FunctionCall& call, ExprType type);
 
     void emitCompareI32(JSOp compareOp, MCompare::CompareType compareType);
     void emitCompareI64(JSOp compareOp, MCompare::CompareType compareType);
     void emitCompareF32(JSOp compareOp, MCompare::CompareType compareType);
     void emitCompareF64(JSOp compareOp, MCompare::CompareType compareType);
 
     void emitAddI32();
@@ -4824,17 +4823,17 @@ BaseCompiler::emitBlock()
 
     if (!deadCode_)
         sync();                    // Simplifies branching out from block
 
     return pushControl(&blockEnd);
 }
 
 void
-BaseCompiler::endBlock(ExprType type, bool isFunctionBody)
+BaseCompiler::endBlock(ExprType type)
 {
     Control& block = controlItem(0);
 
     // Save the value.
     AnyReg r;
     if (!deadCode_)
         r = popJoinRegUnlessVoid(type);
 
@@ -4846,27 +4845,22 @@ BaseCompiler::endBlock(ExprType type, bo
         masm.bind(block.label);
         // No value was provided by the fallthrough but the branch out will
         // have stored one in joinReg, so capture that.
         if (deadCode_)
             r = captureJoinRegUnlessVoid(type);
         deadCode_ = false;
     }
 
-    MOZ_ASSERT(stk_.length() == block.stackSize);
+    popValueStackTo(block.stackSize);
+    popControl();
 
     // Retain the value stored in joinReg by all paths, if there are any.
-    if (!deadCode_) {
+    if (!deadCode_)
         pushJoinRegUnlessVoid(r);
-
-        if (isFunctionBody)
-            doReturn(func_.sig().ret());
-    }
-
-    popControl();
 }
 
 bool
 BaseCompiler::emitLoop()
 {
     if (!iter_.readLoop())
         return false;
 
@@ -4894,18 +4888,17 @@ BaseCompiler::endLoop(ExprType type)
     Control& block = controlItem(0);
 
     AnyReg r;
     if (!deadCode_)
         r = popJoinRegUnlessVoid(type);
 
     popStackOnBlockExit(block.framePushed);
 
-    MOZ_ASSERT(stk_.length() == block.stackSize);
-
+    popValueStackTo(block.stackSize);
     popControl();
 
     // Retain the value stored in joinReg by all paths.
     if (!deadCode_)
         pushJoinRegUnlessVoid(r);
 }
 
 // The bodies of the "then" and "else" arms can be arbitrary sequences
@@ -4964,18 +4957,17 @@ BaseCompiler::endIfThen()
     if (ifThen.otherLabel->used())
         masm.bind(ifThen.otherLabel);
 
     if (ifThen.label->used())
         masm.bind(ifThen.label);
 
     deadCode_ = ifThen.deadOnArrival;
 
-    MOZ_ASSERT(stk_.length() == ifThen.stackSize);
-
+    popValueStackTo(ifThen.stackSize);
     popControl();
 }
 
 bool
 BaseCompiler::emitElse()
 {
     ExprType thenType;
     Nothing unused_thenValue;
@@ -4999,17 +4991,17 @@ BaseCompiler::emitElse()
     if (!deadCode_)
         masm.jump(ifThenElse.label);
 
     if (ifThenElse.otherLabel->used())
         masm.bind(ifThenElse.otherLabel);
 
     // Reset to the "else" branch.
 
-    MOZ_ASSERT(stk_.length() == ifThenElse.stackSize);
+    popValueStackTo(ifThenElse.stackSize);
 
     if (!deadCode_)
         freeJoinRegUnlessVoid(r);
 
     deadCode_ = ifThenElse.deadOnArrival;
 
     return true;
 }
@@ -5041,35 +5033,34 @@ BaseCompiler::endIfThenElse(ExprType typ
     if (joinLive) {
         // No value was provided by the "then" path but capture the one
         // provided by the "else" path.
         if (deadCode_)
             r = captureJoinRegUnlessVoid(type);
         deadCode_ = false;
     }
 
-    MOZ_ASSERT(stk_.length() == ifThenElse.stackSize);
-
+    popValueStackTo(ifThenElse.stackSize);
     popControl();
 
     if (!deadCode_)
         pushJoinRegUnlessVoid(r);
 }
 
 bool
 BaseCompiler::emitEnd()
 {
     LabelKind kind;
     ExprType type;
     Nothing unused_value;
     if (!iter_.readEnd(&kind, &type, &unused_value))
         return false;
 
     switch (kind) {
-      case LabelKind::Block: endBlock(type, iter_.controlStackEmpty()); break;
+      case LabelKind::Block: endBlock(type); break;
       case LabelKind::Loop:  endLoop(type); break;
       case LabelKind::UnreachableThen:
       case LabelKind::Then:  endIfThen(); break;
       case LabelKind::Else:  endIfThenElse(type); break;
     }
 
     return true;
 }
@@ -5098,18 +5089,16 @@ BaseCompiler::emitBr()
 
     // The register holding the join value is free for the remainder
     // of this block.
 
     freeJoinRegUnlessVoid(r);
 
     deadCode_ = true;
 
-    popValueStackTo(ctl_.back().stackSize);
-
     return true;
 }
 
 bool
 BaseCompiler::emitBrIf()
 {
     uint32_t relativeDepth;
     ExprType type;
@@ -5238,18 +5227,16 @@ BaseCompiler::emitBrTable()
     // Clean up.
 
     freeI32(rc);
     freeJoinRegUnlessVoid(r);
 
     for (uint32_t i = 0; i < tableLength; i++)
         freeLabel(stubs[i]);
 
-    popValueStackTo(ctl_.back().stackSize);
-
     return true;
 }
 
 bool
 BaseCompiler::emitDrop()
 {
     if (!iter_.readDrop())
         return false;
@@ -5258,44 +5245,44 @@ BaseCompiler::emitDrop()
         return true;
 
     popStackIfMemory();
     popValueStackBy(1);
     return true;
 }
 
 void
-BaseCompiler::doReturn(ExprType type)
+BaseCompiler::doReturn(ExprType type, bool popStack)
 {
     switch (type) {
       case ExprType::Void: {
-        returnCleanup();
+        returnCleanup(popStack);
         break;
       }
       case ExprType::I32: {
         RegI32 rv = popI32(RegI32(ReturnReg));
-        returnCleanup();
+        returnCleanup(popStack);
         freeI32(rv);
         break;
       }
       case ExprType::I64: {
         RegI64 rv = popI64(RegI64(ReturnReg64));
-        returnCleanup();
+        returnCleanup(popStack);
         freeI64(rv);
         break;
       }
       case ExprType::F64: {
         RegF64 rv = popF64(RegF64(ReturnDoubleReg));
-        returnCleanup();
+        returnCleanup(popStack);
         freeF64(rv);
         break;
       }
       case ExprType::F32: {
         RegF32 rv = popF32(RegF32(ReturnFloat32Reg));
-        returnCleanup();
+        returnCleanup(popStack);
         freeF32(rv);
         break;
       }
       default: {
         MOZ_CRASH("Function return type");
       }
     }
 }
@@ -5305,21 +5292,19 @@ BaseCompiler::emitReturn()
 {
     Nothing unused_value;
     if (!iter_.readReturn(&unused_value))
         return false;
 
     if (deadCode_)
         return true;
 
-    doReturn(func_.sig().ret());
+    doReturn(func_.sig().ret(), PopStack(true));
     deadCode_ = true;
 
-    popValueStackTo(ctl_.back().stackSize);
-
     return true;
 }
 
 bool
 BaseCompiler::emitCallArgs(const ValTypeVector& args, FunctionCall& baselineCall)
 {
     MOZ_ASSERT(!deadCode_);
 
@@ -6529,17 +6514,16 @@ BaseCompiler::emitBody()
             CHECK_NEXT(emitBrTable());
           case uint16_t(Op::Return):
             CHECK_NEXT(emitReturn());
           case uint16_t(Op::Unreachable):
             CHECK(iter_.readUnreachable());
             if (!deadCode_) {
                 unreachableTrap();
                 deadCode_ = true;
-                popValueStackTo(ctl_.back().stackSize);
             }
             NEXT();
 
           // Calls
           case uint16_t(Op::Call):
             CHECK_NEXT(emitCall());
           case uint16_t(Op::CallIndirect):
             CHECK_NEXT(emitCallIndirect(/* oldStyle = */ false));
@@ -7125,22 +7109,26 @@ BaseCompiler::emitFunction()
         return false;
 
     beginFunction();
 
     UniquePooledLabel functionEnd(newLabel());
     if (!functionEnd)
         return false;
 
+    // For the function's block; it will be popped by endBlock().
     if (!pushControl(&functionEnd))
         return false;
 
     if (!emitBody())
         return false;
 
+    if (!deadCode_)
+        doReturn(sig.ret(), PopStack(false));
+
     if (!iter_.readFunctionEnd())
         return false;
 
     if (!endFunction())
         return false;
 
     return true;
 }