Stack cleanup works now even in the most evil expansion cases.

6 files changed, 160 insertions, 60 deletions

diff --git a/generic/tclAssembly.c b/generic/tclAssembly.c
index d1af8c6..62641e6 100644
--- a/generic/tclAssembly.c
+++ b/generic/tclAssembly.c
@@ -20,7 +20,7 @@
  *-   break and continue - if exception ranges can be sorted out.
  *-   foreach_start4, foreach_step4
  *-   returnImm, returnStk
- *-   expandStart, expandStkTop, invokeExpanded
+ *-   expandStart, expandStkTop, invokeExpanded, expandDrop
  *-   dictFirst, dictNext, dictDone
  *-   dictUpdateStart, dictUpdateEnd
  *-   jumpTable testing
diff --git a/generic/tclCompCmds.c b/generic/tclCompCmds.c
index 3d6abcf..365e647 100644
--- a/generic/tclCompCmds.c
+++ b/generic/tclCompCmds.c
@@ -493,40 +493,21 @@ TclCompileBreakCmd(
 
     rangePtr = TclGetInnermostExceptionRange(envPtr, TCL_BREAK, &auxPtr);
     if (rangePtr && rangePtr->type == LOOP_EXCEPTION_RANGE) {
-	int toPop;
-
-	/*
-	 * Ditch the extra elements from the auxiliary stack.
-	 */
-
-	toPop = envPtr->expandCount - auxPtr->expandTarget;
-	while (toPop > 0) {
-	    TclEmitOpcode(INST_EXPAND_DROP, envPtr);
-	    toPop--;
-	}
-
-	/*
-	 * Pop off the extra stack frames.
-	 */
-
-	toPop = envPtr->currStackDepth - auxPtr->stackDepth;
-	while (toPop > 0) {
-	    TclEmitOpcode(INST_POP, envPtr);
-	    TclAdjustStackDepth(1, envPtr);
-	    toPop--;
-	}
-
 	/*
 	 * Found the target! No need for a nasty INST_BREAK here.
 	 */
 
+	TclCleanupStackForBreakContinue(envPtr, auxPtr);
 	TclAddLoopBreakFixup(envPtr, auxPtr);
     } else {
 	/*
-	 * Emit a break instruction.
+	 * Emit a real break.
 	 */
 
-	TclEmitOpcode(INST_BREAK, envPtr);
+	PushStringLiteral(envPtr, "");
+	TclEmitOpcode(INST_DUP, envPtr);
+	TclEmitInstInt4(INST_RETURN_IMM, TCL_BREAK, envPtr);
+	TclEmitInt4(0, envPtr);
     }
 
     /*
@@ -846,40 +827,21 @@ TclCompileContinueCmd(
 
     rangePtr = TclGetInnermostExceptionRange(envPtr, TCL_CONTINUE, &auxPtr);
     if (rangePtr && rangePtr->type == LOOP_EXCEPTION_RANGE) {
-	int toPop;
-
-	/*
-	 * Ditch the extra elements from the auxiliary stack.
-	 */
-
-	toPop = envPtr->expandCount - auxPtr->expandTarget;
-	while (toPop > 0) {
-	    TclEmitOpcode(INST_EXPAND_DROP, envPtr);
-	    toPop--;
-	}
-
-	/*
-	 * Pop off the extra stack frames.
-	 */
-
-	toPop = envPtr->currStackDepth - auxPtr->stackDepth;
-	while (toPop > 0) {
-	    TclEmitOpcode(INST_POP, envPtr);
-	    TclAdjustStackDepth(1, envPtr);
-	    toPop--;
-	}
-
 	/*
 	 * Found the target! No need for a nasty INST_CONTINUE here.
 	 */
 
+	TclCleanupStackForBreakContinue(envPtr, auxPtr);
 	TclAddLoopContinueFixup(envPtr, auxPtr);
     } else {
 	/*
-	 * Emit a continue instruction.
+	 * Emit a real continue.
 	 */
 
-	TclEmitOpcode(INST_CONTINUE, envPtr);
+	PushStringLiteral(envPtr, "");
+	TclEmitOpcode(INST_DUP, envPtr);
+	TclEmitInstInt4(INST_RETURN_IMM, TCL_CONTINUE, envPtr);
+	TclEmitInt4(0, envPtr);
     }
 
     /*
diff --git a/generic/tclCompile.c b/generic/tclCompile.c
index f8dd504..69517bc 100644
--- a/generic/tclCompile.c
+++ b/generic/tclCompile.c
@@ -541,7 +541,8 @@ InstructionDesc const tclInstructionTable[] = {
 	 * Stack: ... list1 list2 => ... [lconcat list1 list2] */
 
     {"expandDrop",       1,    0,          0,	{OPERAND_NONE}},
-	/* Drops an element from the auxiliary stack. */
+	/* Drops an element from the auxiliary stack, popping stack elements
+	 * until the matching stack depth is reached. */
 
     {NULL, 0, 0, 0, {OPERAND_NONE}}
 };
@@ -574,6 +575,7 @@ static void		RecordByteCodeStats(ByteCode *codePtr);
 static void		RegisterAuxDataType(const AuxDataType *typePtr);
 static int		SetByteCodeFromAny(Tcl_Interp *interp,
 			    Tcl_Obj *objPtr);
+static void		StartExpanding(CompileEnv *envPtr);
 static int		FormatInstruction(ByteCode *codePtr,
 			    const unsigned char *pc, Tcl_Obj *bufferObj);
 static void		PrintSourceToObj(Tcl_Obj *appendObj,
@@ -2063,8 +2065,7 @@ TclCompileScript(
 	     */
 
 	    if (expand) {
-		TclEmitOpcode(INST_EXPAND_START, envPtr);
-		envPtr->expandCount++;
+		StartExpanding(envPtr);
 	    }
 
 	    /*
@@ -3469,6 +3470,7 @@ TclCreateExceptRange(
     auxPtr->supportsContinue = 1;
     auxPtr->stackDepth = envPtr->currStackDepth;
     auxPtr->expandTarget = envPtr->expandCount;
+    auxPtr->expandTargetDepth = -1;
     auxPtr->numBreakTargets = 0;
     auxPtr->breakTargets = NULL;
     auxPtr->allocBreakTargets = 0;
@@ -3589,6 +3591,103 @@ TclAddLoopContinueFixup(
 /*
  * ---------------------------------------------------------------------
  *
+ * TclCleanupStackForBreakContinue --
+ *
+ *	Ditch the extra elements from the auxiliary stack and the main
+ *	stack. How to do this exactly depends on whether there are any
+ *	elements on the auxiliary stack to pop.
+ *
+ * ---------------------------------------------------------------------
+ */
+
+void
+TclCleanupStackForBreakContinue(
+    CompileEnv *envPtr,
+    ExceptionAux *auxPtr)
+{
+    int toPop = envPtr->expandCount - auxPtr->expandTarget;
+
+    if (toPop > 0) {
+	while (toPop > 0) {
+	    TclEmitOpcode(INST_EXPAND_DROP, envPtr);
+	    toPop--;
+	}
+	toPop = auxPtr->expandTargetDepth - auxPtr->stackDepth;
+	while (toPop > 0) {
+	    TclEmitOpcode(INST_POP, envPtr);
+	    TclAdjustStackDepth(1, envPtr);
+	    toPop--;
+	}
+    } else {
+	toPop = envPtr->currStackDepth - auxPtr->stackDepth;
+	while (toPop > 0) {
+	    TclEmitOpcode(INST_POP, envPtr);
+	    TclAdjustStackDepth(1, envPtr);
+	    toPop--;
+	}
+    }
+}
+
+/*
+ * ---------------------------------------------------------------------
+ *
+ * StartExpanding --
+ *
+ *	Pushes an INST_EXPAND_START and does some additional housekeeping so
+ *	that the [break] and [continue] compilers can use an exception-free
+ *	issue to discard it.
+ *
+ * ---------------------------------------------------------------------
+ */
+
+static void
+StartExpanding(
+    CompileEnv *envPtr)
+{
+    int i;
+
+    TclEmitOpcode(INST_EXPAND_START, envPtr);
+
+    /*
+     * Update inner exception ranges with information about the environment
+     * where this expansion started.
+     */
+
+    for (i=0 ; i<envPtr->exceptArrayNext ; i++) {
+	ExceptionRange *rangePtr = &envPtr->exceptArrayPtr[i];
+	ExceptionAux *auxPtr = &envPtr->exceptAuxArrayPtr[i];
+
+	/*
+	 * Ignore loops unless they're still being built.
+	 */
+
+	if (rangePtr->codeOffset > CurrentOffset(envPtr)) {
+	    continue;
+	}
+	if (rangePtr->numCodeBytes != -1) {
+	    continue;
+	}
+
+	/*
+	 * Adequate condition: further out loops and further in exceptions
+	 * don't actually need this information.
+	 */
+
+	if (auxPtr->expandTarget == envPtr->expandCount) {
+	    auxPtr->expandTargetDepth = envPtr->currStackDepth;
+	}
+    }
+
+    /*
+     * There's now one more expansion being processed on the auxiliary stack.
+     */
+
+    envPtr->expandCount++;
+}
+
+/*
+ * ---------------------------------------------------------------------
+ *
  * TclFinalizeLoopExceptionRange --
  *
  *	Finalizes a loop exception range, binding the registered [break] and
@@ -3629,7 +3728,8 @@ TclFinalizeLoopExceptionRange(
 	    int j;
 
 	    /*
-	     * WTF? Can't bind, so revert to an INST_CONTINUE.
+	     * WTF? Can't bind, so revert to an INST_CONTINUE. Not enough
+	     * space to do anything else.
 	     */
 
 	    *site = INST_CONTINUE;
diff --git a/generic/tclCompile.h b/generic/tclCompile.h
index 75de025..15b5477 100644
--- a/generic/tclCompile.h
+++ b/generic/tclCompile.h
@@ -120,6 +120,11 @@ typedef struct ExceptionAux {
 				 * we can't currently discard them except by
 				 * doing INST_INVOKE_EXPANDED; this is a known
 				 * problem. */
+    int expandTargetDepth;	/* The stack depth expected at the outermost
+				 * expansion within the loop. Not meaningful
+				 * if there have are no open expansions
+				 * between the looping level and the point of
+				 * jump issue. */
     int numBreakTargets;	/* The number of [break]s that want to be
 				 * targeted to the place where this loop
 				 * exception will be bound to. */
@@ -987,6 +992,8 @@ MODULE_SCOPE ByteCode *	TclCompileObj(Tcl_Interp *interp, Tcl_Obj *objPtr,
  */
 
 MODULE_SCOPE void	TclCleanupByteCode(ByteCode *codePtr);
+MODULE_SCOPE void	TclCleanupStackForBreakContinue(CompileEnv *envPtr,
+			    ExceptionAux *auxPtr);
 MODULE_SCOPE void	TclCompileCmdWord(Tcl_Interp *interp,
 			    Tcl_Token *tokenPtr, int count,
 			    CompileEnv *envPtr);
diff --git a/generic/tclExecute.c b/generic/tclExecute.c
index 559df0b..fc50a74 100644
--- a/generic/tclExecute.c
+++ b/generic/tclExecute.c
@@ -2720,14 +2720,15 @@ TEBCresume(
 
     case INST_EXPAND_DROP:
 	/*
-	 * Drops an element of the auxObjList. Does not do any clean up of the
-	 * actual stack.
-	 *
-	 * TODO: POP MAIN STACK BACK TO MARKER
+	 * Drops an element of the auxObjList, popping stack elements to
+	 * restore the stack to the state before the point where the aux
+	 * element was created.
 	 */
 
+	CLANG_ASSERT(auxObjList);
+	objc = CURR_DEPTH - auxObjList->internalRep.ptrAndLongRep.value;
 	POP_TAUX_OBJ();
-	NEXT_INST_F(1, 0, 0);
+	NEXT_INST_V(1, objc, 0);
 
     case INST_EXPAND_STKTOP: {
 	int i;
diff --git a/tests/for.test b/tests/for.test
index cfba1fe..c5803ee 100644
--- a/tests/for.test
+++ b/tests/for.test
@@ -882,6 +882,36 @@ test for-7.4 {Bug 3614226: ensure that continue cleans up the expansion stack} m
 	expr {$end - $tmp}
     }}
 } 0
+test for-7.5 {Bug 3614226: ensure that break cleans up the 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 {*}[break] d e f]]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
+test for-7.6 {Bug 3614226: ensure that continue cleans up the 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 {*}[continue] d e f]]
+	    }
+	    set tmp $end
+	    set end [meminfo]
+	}
+	expr {$end - $tmp}
+    }}
+} 0
 
 # cleanup
 ::tcltest::cleanupTests