Fix handling of 'invokeExpanded' and start to do 'returnStk'.

3 files changed, 138 insertions, 9 deletions

diff --git a/generic/tclCompCmdsGR.c b/generic/tclCompCmdsGR.c
index 5513b01..fbd370b 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;
     }
 
@@ -2509,7 +2509,7 @@ TclCompileReturnCmd(
      * Issue the RETURN itself.
      */
 
-    TclEmitOpcode(INST_RETURN_STK, envPtr);
+    TclEmitInvoke(envPtr, INST_RETURN_STK);
     return TCL_OK;
 }
 
diff --git a/generic/tclCompile.c b/generic/tclCompile.c
index 89b9011..ae6e56c 100644
--- a/generic/tclCompile.c
+++ b/generic/tclCompile.c
@@ -1802,9 +1802,7 @@ CompileExpanded(
      * stack-neutral in general.
      */
 
-    TclEmitInvoke(envPtr, INST_INVOKE_EXPANDED);
-    envPtr->expandCount--;
-    TclAdjustStackDepth(1 - wordIdx, envPtr);
+    TclEmitInvoke(envPtr, INST_INVOKE_EXPANDED, wordIdx);
 }
 
 static int 
@@ -3931,7 +3929,8 @@ TclEmitInvoke(
     va_list argList;
     ExceptionRange *rangePtr;
     ExceptionAux *auxBreakPtr, *auxContinuePtr;
-    int arg1, arg2, wordCount = 0, loopRange, breakRange, continueRange;
+    int arg1, arg2, wordCount = 0, expandCount = 0;
+    int loopRange, breakRange, continueRange;
 
     /*
      * Parse the arguments.
@@ -3955,8 +3954,17 @@ TclEmitInvoke(
     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 = arg2 = 0;
+	wordCount = arg1 = va_arg(argList, int);
+	arg2 = 0;
+	expandCount = 1;
 	break;
     }
     va_end(argList);
@@ -3974,7 +3982,7 @@ TclEmitInvoke(
     if (rangePtr == NULL || rangePtr->type != LOOP_EXCEPTION_RANGE) {
 	auxBreakPtr = NULL;
     } else if (auxBreakPtr->stackDepth == envPtr->currStackDepth-wordCount
-	    && auxBreakPtr->expandTarget == envPtr->expandCount) {
+	    && auxBreakPtr->expandTarget == envPtr->expandCount-expandCount) {
 	auxBreakPtr = NULL;
     } else {
 	breakRange = auxBreakPtr - envPtr->exceptAuxArrayPtr;
@@ -3985,7 +3993,7 @@ TclEmitInvoke(
     if (rangePtr == NULL || rangePtr->type != LOOP_EXCEPTION_RANGE) {
 	auxContinuePtr = NULL;
     } else if (auxContinuePtr->stackDepth == envPtr->currStackDepth-wordCount
-	    && auxContinuePtr->expandTarget == envPtr->expandCount) {
+	    && auxContinuePtr->expandTarget == envPtr->expandCount-expandCount) {
 	auxContinuePtr = NULL;
     } else {
 	continueRange = auxBreakPtr - envPtr->exceptAuxArrayPtr;
@@ -4009,10 +4017,15 @@ TclEmitInvoke(
 	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);
diff --git a/tests/for.test b/tests/for.test
index 6a18abe..8abd270 100644
--- a/tests/for.test
+++ b/tests/for.test
@@ -1068,6 +1068,122 @@ test for-7.16 {Bug 3614226: ensure that continue from invoked command only clean
 	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