[2835313] Ensure correct stack balance when break and continue exceptions are about: the hard cases with potential concealed exception generation.

11 files changed, 628 insertions, 37 deletions

diff --git a/generic/tclAssembly.c b/generic/tclAssembly.c
index 946c729..fc51457 100644
--- a/generic/tclAssembly.c
+++ b/generic/tclAssembly.c
@@ -246,6 +246,8 @@ static void		BBEmitInstInt4(AssemblyEnv* assemEnvPtr, int tblIdx,
 			    int opnd, int count);
 static void		BBEmitInst1or4(AssemblyEnv* assemEnvPtr, int tblIdx,
 			    int param, int count);
+static void		BBEmitInvoke1or4(AssemblyEnv* assemEnvPtr, int tblIdx,
+			    int param, int count);
 static void		BBEmitOpcode(AssemblyEnv* assemEnvPtr, int tblIdx,
 			    int count);
 static int		BuildExceptionRanges(AssemblyEnv* assemEnvPtr);
@@ -679,10 +681,13 @@ BBEmitInstInt4(
 /*
  *-----------------------------------------------------------------------------
  *
- * BBEmitInst1or4 --
+ * BBEmitInst1or4, BBEmitInvoke1or4 --
  *
  *	Emits a 1- or 4-byte operation according to the magnitude of the
- *	operand
+ *	operand. The Invoke variant generates wrapping stack-balance
+ *	management if necessary (which is not normally required in assembled
+ *	code, as loop exception ranges, expansions, breaks and continues can't
+ *	be issued currently).
  *
  *-----------------------------------------------------------------------------
  */
@@ -714,6 +719,29 @@ BBEmitInst1or4(
     TclUpdateAtCmdStart(op, envPtr);
     BBUpdateStackReqs(bbPtr, tblIdx, count);
 }
+
+static void
+BBEmitInvoke1or4(
+    AssemblyEnv* assemEnvPtr,	/* Assembly environment */
+    int tblIdx,			/* Index in TalInstructionTable of op */
+    int param,			/* Variable-length parameter */
+    int count)			/* Arity if variadic */
+{
+    CompileEnv* envPtr = assemEnvPtr->envPtr;
+				/* Compilation environment */
+    BasicBlock* bbPtr = assemEnvPtr->curr_bb;
+				/* Current basic block */
+    int op = TalInstructionTable[tblIdx].tclInstCode;
+
+    if (param <= 0xff) {
+	op >>= 8;
+    } else {
+	op &= 0xff;
+    }
+    TclEmitInvoke(envPtr, op, param);
+    TclUpdateAtCmdStart(op, envPtr);
+    BBUpdateStackReqs(bbPtr, tblIdx, count);
+}
 
 /*
  *-----------------------------------------------------------------------------
@@ -1450,7 +1478,7 @@ AssembleOneLine(
 	    goto cleanup;
 	}
 
-	BBEmitInst1or4(assemEnvPtr, tblIdx, opnd, opnd);
+	BBEmitInvoke1or4(assemEnvPtr, tblIdx, opnd, opnd);
 	break;
 
     case ASSEM_JUMP:
diff --git a/generic/tclCompCmds.c b/generic/tclCompCmds.c
index 7e6b6da..9c43bfe 100644
--- a/generic/tclCompCmds.c
+++ b/generic/tclCompCmds.c
@@ -269,7 +269,7 @@ TclCompileArraySetCmd(
     if (isDataValid && !isDataEven) {
 	PushStringLiteral(envPtr, "list must have an even number of elements");
 	PushStringLiteral(envPtr, "-errorCode {TCL ARGUMENT FORMAT}");
-	TclEmitInstInt4(INST_RETURN_IMM, 1,			envPtr);
+	TclEmitInstInt4(INST_RETURN_IMM, TCL_ERROR,		envPtr);
 	TclEmitInt4(		0,				envPtr);
 	goto done;
     }
@@ -354,7 +354,7 @@ TclCompileArraySetCmd(
 	TclEmitInstInt1(INST_JUMP_FALSE1, 0,			envPtr);
 	PushStringLiteral(envPtr, "list must have an even number of elements");
 	PushStringLiteral(envPtr, "-errorCode {TCL ARGUMENT FORMAT}");
-	TclEmitInstInt4(INST_RETURN_IMM, 1,			envPtr);
+	TclEmitInstInt4(INST_RETURN_IMM, TCL_ERROR,		envPtr);
 	TclEmitInt4(		0,				envPtr);
 	TclAdjustStackDepth(-1, envPtr);
 	fwd = CurrentOffset(envPtr) - offsetFwd;
@@ -620,7 +620,7 @@ TclCompileCatchCmd(
 	TclEmitInstInt4(	INST_BEGIN_CATCH4, range,	envPtr);
 	ExceptionRangeStarts(envPtr, range);
 	TclEmitOpcode(		INST_DUP,			envPtr);
-	TclEmitOpcode(		INST_EVAL_STK,			envPtr);
+	TclEmitInvoke(envPtr,	INST_EVAL_STK);
     }
     /* Stack at this point:
      *    nonsimple:  script <mark> result
@@ -1675,7 +1675,7 @@ TclCompileDictUpdateCmd(
 
     TclEmitInstInt4(	INST_DICT_UPDATE_END, dictIndex,	envPtr);
     TclEmitInt4(		infoIndex,			envPtr);
-    TclEmitOpcode(	INST_RETURN_STK,			envPtr);
+    TclEmitInvoke(envPtr,INST_RETURN_STK);
 
     if (TclFixupForwardJumpToHere(envPtr, &jumpFixup, 127)) {
 	Tcl_Panic("TclCompileDictCmd(update): bad jump distance %d",
@@ -2033,7 +2033,7 @@ TclCompileDictWithCmd(
     } else {
 	TclEmitInstInt4(	INST_DICT_RECOMBINE_IMM, dictVar, envPtr);
     }
-    TclEmitOpcode(		INST_RETURN_STK,		envPtr);
+    TclEmitInvoke(envPtr,	INST_RETURN_STK);
 
     /*
      * Prepare for the start of the next command.
@@ -2135,19 +2135,48 @@ TclCompileErrorCmd(
 {
     /*
      * General syntax: [error message ?errorInfo? ?errorCode?]
-     * However, we only deal with the case where there is just a message.
      */
-    Tcl_Token *messageTokenPtr;
+
+    Tcl_Token *tokenPtr;
     DefineLineInformation;	/* TIP #280 */
 
-    if (parsePtr->numWords != 2) {
+    if (parsePtr->numWords < 2 || parsePtr->numWords > 4) {
 	return TCL_ERROR;
     }
-    messageTokenPtr = TokenAfter(parsePtr->tokenPtr);
 
-    PushStringLiteral(envPtr, "-code error -level 0");
-    CompileWord(envPtr, messageTokenPtr, interp, 1);
-    TclEmitOpcode(INST_RETURN_STK, envPtr);
+    /*
+     * Handle the message.
+     */
+
+    tokenPtr = TokenAfter(parsePtr->tokenPtr);
+    CompileWord(envPtr, tokenPtr, interp, 1);
+
+    /*
+     * Construct the options. Note that -code and -level are not here.
+     */
+
+    if (parsePtr->numWords == 2) {
+	PushStringLiteral(envPtr, "");
+    } else {
+	PushStringLiteral(envPtr, "-errorinfo");
+	tokenPtr = TokenAfter(tokenPtr);
+	CompileWord(envPtr, tokenPtr, interp, 2);
+	if (parsePtr->numWords == 3) {
+	    TclEmitInstInt4(	INST_LIST, 2,			envPtr);
+	} else {
+	    PushStringLiteral(envPtr, "-errorcode");
+	    tokenPtr = TokenAfter(tokenPtr);
+	    CompileWord(envPtr, tokenPtr, interp, 3);
+	    TclEmitInstInt4(	INST_LIST, 4,			envPtr);
+	}
+    }
+
+    /*
+     * Issue the error via 'returnImm error 0'.
+     */
+
+    TclEmitInstInt4(		INST_RETURN_IMM, TCL_ERROR,	envPtr);
+    TclEmitInt4(			0,			envPtr);
     return TCL_OK;
 }
 
diff --git a/generic/tclCompCmdsGR.c b/generic/tclCompCmdsGR.c
index 43ea3d3..d00327d 100644
--- a/generic/tclCompCmdsGR.c
+++ b/generic/tclCompCmdsGR.c
@@ -2367,7 +2367,7 @@ TclCompileReturnCmd(
 
 	CompileWord(envPtr, optsTokenPtr, interp, 2);
 	CompileWord(envPtr, msgTokenPtr,  interp, 3);
-	TclEmitOpcode(INST_RETURN_STK, envPtr);
+	TclEmitInvoke(envPtr, INST_RETURN_STK);
 	return TCL_OK;
     }
 
@@ -2381,6 +2381,10 @@ TclCompileReturnCmd(
      * Scan through the return options. If any are unknown at compile time,
      * there is no value in bytecompiling. Save the option values known in an
      * objv array for merging into a return options dictionary.
+     *
+     * TODO: There is potential for improvement if all option keys are known
+     * at compile time and all option values relating to '-code' and '-level'
+     * are known at compile time.
      */
 
     for (objc = 0; objc < numOptionWords; objc++) {
@@ -2388,7 +2392,7 @@ TclCompileReturnCmd(
 	Tcl_IncrRefCount(objv[objc]);
 	if (!TclWordKnownAtCompileTime(wordTokenPtr, objv[objc])) {
 	    /*
-	     * Non-literal, so punt to run-time.
+	     * Non-literal, so punt to run-time assembly of the dictionary.
 	     */
 
 	    for (; objc>=0 ; objc--) {
@@ -2509,7 +2513,7 @@ TclCompileReturnCmd(
      * Issue the RETURN itself.
      */
 
-    TclEmitOpcode(INST_RETURN_STK, envPtr);
+    TclEmitInvoke(envPtr, INST_RETURN_STK);
     return TCL_OK;
 }
 
@@ -2521,6 +2525,23 @@ CompileReturnInternal(
     int level,
     Tcl_Obj *returnOpts)
 {
+    if (level == 0 && (code == TCL_BREAK || code == TCL_CONTINUE)) {
+	ExceptionRange *rangePtr;
+	ExceptionAux *exceptAux;
+
+	rangePtr = TclGetInnermostExceptionRange(envPtr, code, &exceptAux);
+	if (rangePtr && rangePtr->type == LOOP_EXCEPTION_RANGE) {
+	    TclCleanupStackForBreakContinue(envPtr, exceptAux);
+	    if (code == TCL_BREAK) {
+		TclAddLoopBreakFixup(envPtr, exceptAux);
+	    } else {
+		TclAddLoopContinueFixup(envPtr, exceptAux);
+	    }
+	    Tcl_DecrRefCount(returnOpts);
+	    return;
+	}
+    }
+
     TclEmitPush(TclAddLiteralObj(envPtr, returnOpts, NULL), envPtr);
     TclEmitInstInt4(op, code, envPtr);
     TclEmitInt4(level, envPtr);
diff --git a/generic/tclCompCmdsSZ.c b/generic/tclCompCmdsSZ.c
index 44cb66e..754238f 100644
--- a/generic/tclCompCmdsSZ.c
+++ b/generic/tclCompCmdsSZ.c
@@ -99,6 +99,8 @@ const AuxDataType tclJumptableInfoType = {
     if ((idx)<256) {OP1(LOAD_SCALAR1,(idx));} else {OP4(LOAD_SCALAR4,(idx));}
 #define STORE(idx) \
     if ((idx)<256) {OP1(STORE_SCALAR1,(idx));} else {OP4(STORE_SCALAR4,(idx));}
+#define INVOKE(name) \
+    TclEmitInvoke(envPtr,INST_##name)
 
 /*
  *----------------------------------------------------------------------
@@ -873,7 +875,7 @@ TclSubstCompile(
 	OP(	END_CATCH);
 	OP(	RETURN_CODE_BRANCH);
 
-	/* ERROR -> reraise it */
+	/* ERROR -> reraise it; NB: can't require BREAK/CONTINUE handling */
 	OP(	RETURN_STK);
 	OP(	NOP);
 
@@ -1965,7 +1967,7 @@ TclCompileThrowCmd(
 	OP(			LIST_LENGTH);
 	OP1(			JUMP_FALSE1, 16);
 	OP4(			LIST, 2);
-	OP44(			RETURN_IMM, 1, 0);
+	OP44(			RETURN_IMM, TCL_ERROR, 0);
 	TclAdjustStackDepth(2, envPtr);
 	OP(			POP);
 	OP(			POP);
@@ -1974,7 +1976,7 @@ TclCompileThrowCmd(
 	PUSH(			"type must be non-empty list");
 	PUSH(			"-errorcode {TCL OPERATION THROW BADEXCEPTION}");
     }
-    OP44(			RETURN_IMM, 1, 0);
+    OP44(			RETURN_IMM, TCL_ERROR, 0);
     return TCL_OK;
 }
 
@@ -2396,7 +2398,7 @@ IssueTryClausesInstructions(
 	    TclAdjustStackDepth(-1, envPtr);
 	    FIXJUMP1(		dontChangeOptions);
 	    OP4(			REVERSE, 2);
-	    OP(				RETURN_STK);
+	    INVOKE(			RETURN_STK);
 	}
 
 	JUMP4(				JUMP, addrsToFix[i]);
@@ -2415,7 +2417,7 @@ IssueTryClausesInstructions(
     OP(					POP);
     LOAD(				optionsVar);
     LOAD(				resultVar);
-    OP(					RETURN_STK);
+    INVOKE(				RETURN_STK);
 
     /*
      * Fix all the jumps from taken clauses to here (which is the end of the
@@ -2724,7 +2726,7 @@ IssueTryClausesFinallyInstructions(
     FIXJUMP1(			finalOK);
     LOAD(				optionsVar);
     LOAD(				resultVar);
-    OP(					RETURN_STK);
+    INVOKE(				RETURN_STK);
 
     return TCL_OK;
 }
@@ -2783,7 +2785,7 @@ IssueTryFinallyInstructions(
     OP1(				JUMP1, 7);
     FIXJUMP1(		jumpOK);
     OP4(				REVERSE, 2);
-    OP(					RETURN_STK);
+    INVOKE(				RETURN_STK);
     return TCL_OK;
 }
 
diff --git a/generic/tclCompExpr.c b/generic/tclCompExpr.c
index d8e4d9f..94c1bd6 100644
--- a/generic/tclCompExpr.c
+++ b/generic/tclCompExpr.c
@@ -2335,9 +2335,9 @@ CompileExprTree(
 		 */
 		
 		if (numWords < 255) {
-		    TclEmitInstInt1(INST_INVOKE_STK1, numWords, envPtr);
+		    TclEmitInvoke(envPtr, INST_INVOKE_STK1, numWords);
 		} else {
-		    TclEmitInstInt4(INST_INVOKE_STK4, numWords, envPtr);
+		    TclEmitInvoke(envPtr, INST_INVOKE_STK4, numWords);
 		}
 
 		/*
diff --git a/generic/tclCompile.c b/generic/tclCompile.c
index d15ef3a..ae6e56c 100644
--- a/generic/tclCompile.c
+++ b/generic/tclCompile.c
@@ -1738,9 +1738,9 @@ TclCompileInvocation(
     }
 
     if (wordIdx <= 255) {
-	TclEmitInstInt1(INST_INVOKE_STK1, wordIdx, envPtr);
+	TclEmitInvoke(envPtr, INST_INVOKE_STK1, wordIdx);
     } else {
-	TclEmitInstInt4(INST_INVOKE_STK4, wordIdx, envPtr);
+	TclEmitInvoke(envPtr, INST_INVOKE_STK4, wordIdx);
     }
 }
 
@@ -1802,9 +1802,7 @@ CompileExpanded(
      * stack-neutral in general.
      */
 
-    TclEmitOpcode(INST_INVOKE_EXPANDED, envPtr);
-    envPtr->expandCount--;
-    TclAdjustStackDepth(1 - wordIdx, envPtr);
+    TclEmitInvoke(envPtr, INST_INVOKE_EXPANDED, wordIdx);
 }
 
 static int 
@@ -2458,7 +2456,7 @@ TclCompileCmdWord(
 	 */
 
 	TclCompileTokens(interp, tokenPtr, count, envPtr);
-	TclEmitOpcode(INST_EVAL_STK, envPtr);
+	TclEmitInvoke(envPtr, INST_EVAL_STK);
     }
 }
 
@@ -3904,6 +3902,198 @@ TclFixupForwardJump(
 /*
  *----------------------------------------------------------------------
  *
+ * TclEmitInvoke --
+ *
+ *	Emit one of the invoke-related instructions, wrapping it if necessary
+ *	in code that ensures that any break or continue operation passing
+ *	through it gets the stack unwinding correct, converting it into an
+ *	internal jump if in an appropriate context.
+ *
+ * Results:
+ *	None
+ *
+ * Side effects:
+ *	Issues the jump with all correct stack management. May create another
+ *	loop exception range; pointers to ExceptionRange and ExceptionAux
+ *	structures should not be held across this call.
+ *
+ *----------------------------------------------------------------------
+ */
+
+void
+TclEmitInvoke(
+    CompileEnv *envPtr,
+    int opcode,
+    ...)
+{
+    va_list argList;
+    ExceptionRange *rangePtr;
+    ExceptionAux *auxBreakPtr, *auxContinuePtr;
+    int arg1, arg2, wordCount = 0, expandCount = 0;
+    int loopRange, breakRange, continueRange;
+
+    /*
+     * Parse the arguments.
+     */
+
+    va_start(argList, opcode);
+    switch (opcode) {
+    case INST_INVOKE_STK1:
+	wordCount = arg1 = va_arg(argList, int);
+	arg2 = 0;
+	break;
+    case INST_INVOKE_STK4:
+	wordCount = arg1 = va_arg(argList, int);
+	arg2 = 0;
+	break;
+    case INST_INVOKE_REPLACE:
+	arg1 = va_arg(argList, int);
+	arg2 = va_arg(argList, int);
+	wordCount = arg1 + arg2 - 1;
+	break;
+    default:
+	Tcl_Panic("unexpected opcode");
+    case INST_EVAL_STK:
+	wordCount = 1;
+	arg1 = arg2 = 0;
+	break;
+    case INST_RETURN_STK:
+	wordCount = 2;
+	arg1 = arg2 = 0;
+	break;
+    case INST_INVOKE_EXPANDED:
+	wordCount = arg1 = va_arg(argList, int);
+	arg2 = 0;
+	expandCount = 1;
+	break;
+    }
+    va_end(argList);
+
+    /*
+     * Determine if we need to handle break and continue exceptions with a
+     * special handling exception range (so that we can correctly unwind the
+     * stack).
+     *
+     * These must be done separately; they can be different (especially for
+     * calls from inside a [for] increment clause).
+     */
+
+    rangePtr = TclGetInnermostExceptionRange(envPtr, TCL_BREAK, &auxBreakPtr);
+    if (rangePtr == NULL || rangePtr->type != LOOP_EXCEPTION_RANGE) {
+	auxBreakPtr = NULL;
+    } else if (auxBreakPtr->stackDepth == envPtr->currStackDepth-wordCount
+	    && auxBreakPtr->expandTarget == envPtr->expandCount-expandCount) {
+	auxBreakPtr = NULL;
+    } else {
+	breakRange = auxBreakPtr - envPtr->exceptAuxArrayPtr;
+    }
+
+    rangePtr = TclGetInnermostExceptionRange(envPtr, TCL_CONTINUE,
+	    &auxContinuePtr);
+    if (rangePtr == NULL || rangePtr->type != LOOP_EXCEPTION_RANGE) {
+	auxContinuePtr = NULL;
+    } else if (auxContinuePtr->stackDepth == envPtr->currStackDepth-wordCount
+	    && auxContinuePtr->expandTarget == envPtr->expandCount-expandCount) {
+	auxContinuePtr = NULL;
+    } else {
+	continueRange = auxBreakPtr - envPtr->exceptAuxArrayPtr;
+    }
+
+    if (auxBreakPtr != NULL || auxContinuePtr != NULL) {
+	loopRange = TclCreateExceptRange(LOOP_EXCEPTION_RANGE, envPtr);
+	ExceptionRangeStarts(envPtr, loopRange);
+    }
+
+    /*
+     * Issue the invoke itself.
+     */
+
+    switch (opcode) {
+    case INST_INVOKE_STK1:
+	TclEmitInstInt1(INST_INVOKE_STK1, arg1, envPtr);
+	break;
+    case INST_INVOKE_STK4:
+	TclEmitInstInt4(INST_INVOKE_STK4, arg1, envPtr);
+	break;
+    case INST_INVOKE_EXPANDED:
+	TclEmitOpcode(INST_INVOKE_EXPANDED, envPtr);
+	envPtr->expandCount--;
+	TclAdjustStackDepth(1 - arg1, envPtr);
+	break;
+    case INST_EVAL_STK:
+	TclEmitOpcode(INST_EVAL_STK, envPtr);
+	break;
+    case INST_RETURN_STK:
+	TclEmitOpcode(INST_RETURN_STK, envPtr);
+	break;
+    case INST_INVOKE_REPLACE:
+	TclEmitInstInt4(INST_INVOKE_REPLACE, arg1, envPtr);
+	TclEmitInt1(arg2, envPtr);
+	TclAdjustStackDepth(-1, envPtr); /* Correction to stack depth calcs */
+	break;
+    }
+
+    /*
+     * If we're generating a special wrapper exception range, we need to
+     * finish that up now.
+     */
+
+    if (auxBreakPtr != NULL || auxContinuePtr != NULL) {
+	int savedStackDepth = envPtr->currStackDepth;
+	int savedExpandCount = envPtr->expandCount;
+	JumpFixup nonTrapFixup;
+	ExceptionAux *exceptAux = envPtr->exceptAuxArrayPtr + loopRange;
+
+	if (auxBreakPtr != NULL) {
+	    auxBreakPtr = envPtr->exceptAuxArrayPtr + breakRange;
+	}
+	if (auxContinuePtr != NULL) {
+	    auxContinuePtr = envPtr->exceptAuxArrayPtr + continueRange;
+	}
+
+	ExceptionRangeEnds(envPtr, loopRange);
+	TclEmitForwardJump(envPtr, TCL_UNCONDITIONAL_JUMP, &nonTrapFixup);
+
+	/*
+	 * Careful! When generating these stack unwinding sequences, the depth
+	 * of stack in the cases where they are taken is not the same as if
+	 * the exception is not taken.
+	 */
+
+	if (auxBreakPtr != NULL) {
+	    TclAdjustStackDepth(-1, envPtr);
+	    exceptAux->stackDepth = auxBreakPtr->stackDepth;
+	    exceptAux->expandTarget = auxBreakPtr->expandTarget;
+
+	    ExceptionRangeTarget(envPtr, loopRange, breakOffset);
+	    TclCleanupStackForBreakContinue(envPtr, exceptAux);
+	    TclAddLoopBreakFixup(envPtr, auxBreakPtr);
+
+	    envPtr->currStackDepth = savedStackDepth;
+	    envPtr->expandCount = savedExpandCount;
+	}
+
+	if (auxContinuePtr != NULL) {
+	    TclAdjustStackDepth(-1, envPtr);
+	    exceptAux->stackDepth = auxContinuePtr->stackDepth;
+	    exceptAux->expandTarget = auxContinuePtr->expandTarget;
+
+	    ExceptionRangeTarget(envPtr, loopRange, continueOffset);
+	    TclCleanupStackForBreakContinue(envPtr, exceptAux);
+	    TclAddLoopContinueFixup(envPtr, auxContinuePtr);
+
+	    envPtr->currStackDepth = savedStackDepth;
+	    envPtr->expandCount = savedExpandCount;
+	}
+
+	TclFinalizeLoopExceptionRange(envPtr, loopRange);
+	TclFixupForwardJumpToHere(envPtr, &nonTrapFixup, 127);
+    }
+}
+
+/*
+ *----------------------------------------------------------------------
+ *
  * TclGetInstructionTable --
  *
  *	Returns a pointer to the table describing Tcl bytecode instructions.
diff --git a/generic/tclCompile.h b/generic/tclCompile.h
index 5660055..a39e0f1 100644
--- a/generic/tclCompile.h
+++ b/generic/tclCompile.h
@@ -1021,6 +1021,7 @@ MODULE_SCOPE void	TclDeleteLiteralTable(Tcl_Interp *interp,
 			    LiteralTable *tablePtr);
 MODULE_SCOPE void	TclEmitForwardJump(CompileEnv *envPtr,
 			    TclJumpType jumpType, JumpFixup *jumpFixupPtr);
+MODULE_SCOPE void	TclEmitInvoke(CompileEnv *envPtr, int opcode, ...);
 MODULE_SCOPE ExceptionRange * TclGetExceptionRangeForPc(unsigned char *pc,
 			    int catchOnly, ByteCode *codePtr);
 MODULE_SCOPE void	TclExpandJumpFixupArray(JumpFixupArray *fixupArrayPtr);
diff --git a/generic/tclEnsemble.c b/generic/tclEnsemble.c
index ad11785..9bb7a0c 100644
--- a/generic/tclEnsemble.c
+++ b/generic/tclEnsemble.c
@@ -3179,9 +3179,7 @@ CompileToInvokedCommand(
      * Do the replacing dispatch.
      */
 
-    TclEmitInstInt4(INST_INVOKE_REPLACE, parsePtr->numWords, envPtr);
-    TclEmitInt1(numWords+1, envPtr);
-    TclAdjustStackDepth(-1, envPtr);	/* Correction to stack depth calcs. */
+    TclEmitInvoke(envPtr, INST_INVOKE_REPLACE, parsePtr->numWords,numWords+1);
 }
 
 /*
diff --git a/generic/tclOptimize.c b/generic/tclOptimize.c
index b7f4173..3b16e6e 100644
--- a/generic/tclOptimize.c
+++ b/generic/tclOptimize.c
@@ -344,21 +344,28 @@ AdvanceJumps(
     CompileEnv *envPtr)
 {
     unsigned char *currentInstPtr;
+    Tcl_HashTable jumps;
 
     for (currentInstPtr = envPtr->codeStart ;
 	    currentInstPtr < envPtr->codeNext-1 ;
 	    currentInstPtr += AddrLength(currentInstPtr)) {
-	int offset, delta;
+	int offset, delta, isNew;
 
 	switch (*currentInstPtr) {
 	case INST_JUMP1:
 	case INST_JUMP_TRUE1:
 	case INST_JUMP_FALSE1:
 	    offset = TclGetInt1AtPtr(currentInstPtr + 1);
+	    Tcl_InitHashTable(&jumps, TCL_ONE_WORD_KEYS);
 	    for (delta=0 ; offset+delta != 0 ;) {
 		if (offset + delta < -128 || offset + delta > 127) {
 		    break;
 		}
+		Tcl_CreateHashEntry(&jumps, INT2PTR(offset), &isNew);
+		if (!isNew) {
+		    offset = TclGetInt1AtPtr(currentInstPtr + 1);
+		    break;
+		}
 		offset += delta;
 		switch (*(currentInstPtr + offset)) {
 		case INST_NOP:
@@ -373,13 +380,21 @@ AdvanceJumps(
 		}
 		break;
 	    }
+	    Tcl_DeleteHashTable(&jumps);
 	    TclStoreInt1AtPtr(offset, currentInstPtr + 1);
 	    continue;
 
 	case INST_JUMP4:
 	case INST_JUMP_TRUE4:
 	case INST_JUMP_FALSE4:
+	    Tcl_InitHashTable(&jumps, TCL_ONE_WORD_KEYS);
+	    Tcl_CreateHashEntry(&jumps, INT2PTR(0), &isNew);
 	    for (offset = TclGetInt4AtPtr(currentInstPtr + 1); offset!=0 ;) {
+		Tcl_CreateHashEntry(&jumps, INT2PTR(offset), &isNew);
+		if (!isNew) {
+		    offset = TclGetInt4AtPtr(currentInstPtr + 1);
+		    break;
+		}
 		switch (*(currentInstPtr + offset)) {
 		case INST_NOP:
 		    offset += InstLength(INST_NOP);
@@ -393,6 +408,7 @@ AdvanceJumps(
 		}
 		break;
 	    }
+	    Tcl_DeleteHashTable(&jumps);
 	    TclStoreInt4AtPtr(offset, currentInstPtr + 1);
 	    continue;
 	}
diff --git a/tests/compile.test b/tests/compile.test
index 51db0a2..36e24de 100644
--- a/tests/compile.test
+++ b/tests/compile.test
@@ -713,6 +713,70 @@ test compile-19.0 {Bug 3614102: reset stack housekeeping} -body {
     apply {{} {list [if 1]}}
 } -returnCodes error -match glob -result *
 
+test compile-20.1 {ensure there are no infinite loops in optimizing} {
+    tcl::unsupported::disassemble script {
+	while 1 {
+	    return -code continue -level 0
+	}
+    }
+    return
+} {}
+test compile-20.2 {ensure there are no infinite loops in optimizing} {
+    tcl::unsupported::disassemble script {
+	while 1 {
+	    while 1 {
+		return -code break -level 0
+	    }
+	}
+    }
+    return
+} {}
+
+test compile-21.1 {stack balance management} {
+    apply {{} {
+	set result {}
+	while 1 {
+	    lappend result a
+	    lappend result [list b [break]]
+	    lappend result c
+	}
+	return $result
+    }}
+} a
+test compile-21.2 {stack balance management} {
+    apply {{} {
+	set result {}
+	while {[incr i] <= 10} {
+	    lappend result $i
+	    lappend result [list b [continue] c]
+	    lappend result c
+	}
+	return $result
+    }}
+} {1 2 3 4 5 6 7 8 9 10}
+test compile-21.3 {stack balance management} {
+    apply {args {
+	set result {}
+	while 1 {
+	    lappend result a
+	    lappend result [concat {*}$args [break]]
+	    lappend result c
+	}
+	return $result
+    }} P Q R S T
+} a
+test compile-21.4 {stack balance management} {
+    apply {args {
+	set result {}
+	while {[incr i] <= 10} {
+	    lappend result $i
+	    lappend result [concat {*}$args [continue] c]
+	    lappend result c
+	}
+	return $result
+    }} P Q R S T
+} {1 2 3 4 5 6 7 8 9 10}
+
 # TODO sometime - check that bytecode from tbcload is *not* disassembled.
 
 # cleanup
diff --git a/tests/for.test b/tests/for.test
index 8936682..8abd270 100644
--- a/tests/for.test
+++ b/tests/for.test
@@ -942,6 +942,248 @@ test for-7.8 {Bug 3614226: ensure that continue only cleans up the right amount}
 	expr {$end - $tmp}
     }}
 } 0
+test for-7.9 {Bug 3614226: ensure that break from invoked command cleans up the stack} memory {
+    apply {{} {
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {$x < 5} {incr x} {
+		list a b c [apply {{} {return -code break}}] d e f
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
+test for-7.10 {Bug 3614226: ensure that continue from invoked command cleans up the stack} memory {
+    apply {{} {
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {$x < 5} {incr x} {
+		list a b c [apply {{} {return -code continue}}] d e f
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
+test for-7.11 {Bug 3614226: ensure that break from invoked command cleans up the expansion stack} memory {
+    apply {{} {
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {[incr x]<50} {} {
+		puts {*}[puts a b c {*}[apply {{} {return -code break}}] d e f]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
+test for-7.12 {Bug 3614226: ensure that continue from invoked command cleans up the expansion stack} memory {
+    apply {{} {
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {[incr x]<50} {} {
+		puts {*}[puts a b c {*}[apply {{} {
+		    return -code continue
+		}}] d e f]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
+test for-7.13 {Bug 3614226: ensure that break from invoked command cleans up the combination of main and expansion stack} memory {
+    apply {{} {
+	set l [lrepeat 50 p q r]
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {[incr x]<50} {} {
+		puts [puts {*}$l {*}[puts a b c {*}$l {*}[apply {{} {
+		    return -code break
+		}}] d e f]]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
+test for-7.14 {Bug 3614226: ensure that continue from invoked command cleans up the combination of main and expansion stack} memory {
+    apply {{} {
+	set l [lrepeat 50 p q r]
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {[incr x]<50} {} {
+		puts [puts {*}$l {*}[puts a b c {*}$l {*}[apply {{} {
+		    return -code continue
+		}}] d e f]]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
+test for-7.15 {Bug 3614226: ensure that break from invoked command only cleans up the right amount} memory {
+    apply {{} {
+	set l [lrepeat 50 p q r]
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    unset -nocomplain {*}[for {set x 0} {[incr x]<50} {} {
+		puts [puts {*}$l {*}[puts a b c {*}$l {*}[apply {{} {
+		    return -code break
+		}}] d e f]]
+	    }]
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
+test for-7.16 {Bug 3614226: ensure that continue from invoked command only cleans up the right amount} memory {
+    apply {{} {
+	set l [lrepeat 50 p q r]
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    unset -nocomplain {*}[for {set x 0} {[incr x]<50} {} {
+		puts [puts {*}$l {*}[puts a b c {*}$l {*}[apply {{} {
+		    return -code continue
+		}}] d e f]]
+	    }]
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
+test for-7.17 {Bug 3614226: ensure that break from expanded command cleans up the stack} memory {
+    apply {op {
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {$x < 5} {incr x} {
+		list a b c [{*}$op] d e f
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }} {return -level 0 -code break}
+} 0
+test for-7.18 {Bug 3614226: ensure that continue from expanded command cleans up the stack} memory {
+    apply {op {
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {$x < 5} {incr x} {
+		list a b c [{*}$op] d e f
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }} {return -level 0 -code continue}
+} 0
+test for-7.19 {Bug 3614226: ensure that break from expanded command cleans up the expansion stack} memory {
+    apply {op {
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {[incr x]<50} {} {
+		puts {*}[puts a b c {*}[{*}$op] d e f]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }} {return -level 0 -code break}
+} 0
+test for-7.20 {Bug 3614226: ensure that continue from expanded command cleans up the expansion stack} memory {
+    apply {op {
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {[incr x]<50} {} {
+		puts {*}[puts a b c {*}[{*}$op] d e f]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }} {return -level 0 -code continue}
+} 0
+test for-7.21 {Bug 3614226: ensure that break from expanded command cleans up the combination of main and expansion stack} memory {
+    apply {op {
+	set l [lrepeat 50 p q r]
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {[incr x]<50} {} {
+		puts [puts {*}$l {*}[puts a b c {*}$l {*}[{*}$op] d e f]]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }} {return -level 0 -code break}
+} 0
+test for-7.22 {Bug 3614226: ensure that continue from expanded command cleans up the combination of main and expansion stack} memory {
+    apply {op {
+	set l [lrepeat 50 p q r]
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    for {set x 0} {[incr x]<50} {} {
+		puts [puts {*}$l {*}[puts a b c {*}$l {*}[{*}$op] d e f]]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }} {return -level 0 -code continue}
+} 0
+test for-7.23 {Bug 3614226: ensure that break from expanded command only cleans up the right amount} memory {
+    apply {op {
+	set l [lrepeat 50 p q r]
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    unset -nocomplain {*}[for {set x 0} {[incr x]<50} {} {
+		puts [puts {*}$l {*}[puts a b c {*}$l {*}[{*}$op] d e f]]
+	    }]
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }} {return -level 0 -code break}
+} 0
+test for-7.24 {Bug 3614226: ensure that continue from expanded command only cleans up the right amount} memory {
+    apply {op {
+	set l [lrepeat 50 p q r]
+	# Can't use [memtest]; must be careful when we change stack frames
+	set end [meminfo]
+	for {set i 0} {$i < 5} {incr i} {
+	    unset -nocomplain {*}[for {set x 0} {[incr x]<50} {} {
+		puts [puts {*}$l {*}[puts a b c {*}$l {*}[{*}$op] d e f]]
+	    }]
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }} {return -level 0 -code continue}
+} 0
 
 # cleanup
 ::tcltest::cleanupTests