Fix CONTINUE statement in plpgsql FOR ... IN ... SELECT loops and integer FOR loops.

server/plpgsql/json.go

@@ -577,34 +577,42 @@ func (stmt *plpgSQL_stmt_fori) Convert() (block Block, err error) {
 
 	// Build the loop body:
 	//   [0]           InitAssign: varName := lower
-	//   [1]           If(condition, GotoOffset:2) → jumps to [3] (first body stmt) when true
-	//   [2]           ExitGoto → offset=3+bodySize → jumps to ScopeEnd
-	//   [3..3+N-1]    body statements (N = bodySize)
-	//   [3+N]         IncrAssign: varName := varName +/- step
-	//   [3+N+1]       BackGoto → offset=-(3+bodySize) → jumps back to If at [1]
+	//   [1]           SkipIncrGoto → offset=2 → jumps to [3], so the first iteration uses lower unchanged
+	//   [2]           IncrAssign: varName := varName +/- step (the CONTINUE target)
+	//   [3]           If(condition, GotoOffset:2) → jumps to [5] (first body stmt) when true
+	//   [4]           ExitGoto → offset=2+bodySize → jumps to ScopeEnd
+	//   [5..5+N-1]    body statements (N = bodySize)
+	//   [5+N]         BackGoto → offset=-(3+bodySize) → jumps back to IncrAssign at [2]
+	//
+	// The increment sits ahead of the condition, rather than at the end of the body, so that CONTINUE has a
+	// single operation to jump to that both advances the loop and re-tests the condition.
 	//
 	// Because no variables are declared in this block (the loop variable is already
 	// declared by the caller's DECLARE section), ScopeBegin is at M and the
 	// InitAssign is at M+1, so all offsets are consistent.
+	block.ContinueTargetOffset = 2
 	block.Body = []Statement{
 		Assignment{
 			VariableName: varName,
 			Expression:   lowerExpr,
 		},
+		Goto{
+			Offset: 2,
+		},
+		Assignment{
+			VariableName: varName,
+			Expression:   incrExpr,
+		},
 		If{
 			Condition:  condition,
 			GotoOffset: 2,
 		},
 		Goto{
-			Offset: 3 + bodySize,
+			Offset: 2 + bodySize,
 		},
 	}
 	block.Body = append(block.Body, convertedBody...)
 	block.Body = append(block.Body,
-		Assignment{
-			VariableName: varName,
-			Expression:   incrExpr,
-		},
 		Goto{
 			Offset: -(3 + bodySize),
 		},
@@ -649,6 +657,10 @@ func (stmt *plpgSQL_stmt_fors) Convert() (block Block, err error) {
 	//   [1] ForQueryNext  – fetch next row into varName, or jump forward by (bodySize+2) to ScopeEnd
 	//   [2..2+bodySize-1] body statements
 	//   [2+bodySize]      Goto back to ForQueryNext: offset = -(1 + bodySize)
+	//
+	// A CONTINUE must fetch the next row rather than re-run the query, so it targets the ForQueryNext at [1]
+	// rather than the ForQueryInit at [0].
+	block.ContinueTargetOffset = 1
 	block.Body = []Statement{
 		ForQueryInit{CursorName: cursorName, Query: query},
 		ForQueryNext{CursorName: cursorName, RecordVar: varName, GotoOffset: bodySize + 2},

server/plpgsql/reconcile_labels.go

@@ -15,11 +15,18 @@
 package plpgsql
 
 import (
+	"strconv"
+
 	"github.com/cockroachdb/errors"
 
 	"github.com/dolthub/doltgresql/utils"
 )
 
+// continueTargetOption is the key, within a loop's ScopeBegin Options, holding the offset from that ScopeBegin
+// to the loop's next-iteration operation, which is where a CONTINUE for that loop jumps. Block.AppendOperations
+// sets it and reconcileLabels removes it, so it never reaches the persisted operations.
+const continueTargetOption = "continue_target"
+
 // labelStackItem is the stack item used while reconciling labels.
 type labelStackItem struct {
 	label  string
@@ -61,15 +68,29 @@ func reconcileLabels(ops []InterpreterOperation) error {
 				}
 			}
 		case OpCode_ScopeBegin:
+			// A CONTINUE defaults to the operation after this one, else we'll continually increase the scope.
+			// Loops that begin with setup operations instead of their next-iteration step say where it is.
+			start := opIndex + 1
+			if offset, ok := operation.Options[continueTargetOption]; ok {
+				parsedOffset, err := strconv.Atoi(offset)
+				if err != nil {
+					return errors.Wrapf(err, "invalid CONTINUE target for scope at operation %d", opIndex)
+				}
+				start = opIndex + parsedOffset
+			}
 			// We'll push the label and loop status to the stack
 			labels.Push(labelStackItem{
 				label:  operation.PrimaryData,
-				start:  opIndex + 1, // We want to go to the operation after this one, else we'll continually increase the scope
+				start:  start,
 				isLoop: len(operation.Target) > 0,
 			})
-			// We clear the label and loop status since we only set them for reconciliation
+			// We clear the label, loop status, and CONTINUE target since we only set them for reconciliation
 			ops[opIndex].PrimaryData = ""
 			ops[opIndex].Target = ""
+			delete(ops[opIndex].Options, continueTargetOption)
+			if len(ops[opIndex].Options) == 0 {
+				ops[opIndex].Options = nil
+			}
 		case OpCode_ScopeEnd:
 			stackItem := labels.Pop()
 			for gotoIdx, gotoOp := range gotos {

server/plpgsql/statements.go

@@ -16,6 +16,7 @@ package plpgsql
 
 import (
 	"fmt"
+	"strconv"
 
 	"github.com/dolthub/go-mysql-server/sql"
 
@@ -71,6 +72,10 @@ type Block struct {
 	Body       []Statement
 	Label      string
 	IsLoop     bool
+	// ContinueTargetOffset gives the loop's next-iteration operation, where a CONTINUE for this loop jumps, as
+	// an offset from the body's first operation. It applies only when IsLoop is true, and the zero value suits
+	// WHILE and plain LOOP, whose bodies begin with that step rather than with loop setup.
+	ContinueTargetOffset int32
 }
 
 var _ Statement = Block{}
@@ -107,6 +112,7 @@ func (stmt Block) AppendOperations(ops *[]InterpreterOperation, stack *Interpret
 			stmt.Label = stack.GetCurrentLabel()
 		}
 	}
+	scopeBeginIndex := len(*ops)
 	*ops = append(*ops, InterpreterOperation{
 		OpCode:      OpCode_ScopeBegin,
 		PrimaryData: stmt.Label,
@@ -143,6 +149,14 @@ func (stmt Block) AppendOperations(ops *[]InterpreterOperation, stack *Interpret
 			})
 		}
 	}
+	if stmt.IsLoop {
+		// Declarations are already appended, so the body starts at the next operation. reconcileLabels
+		// resolves this loop's CONTINUE statements through this offset.
+		continueTarget := len(*ops) + int(stmt.ContinueTargetOffset)
+		(*ops)[scopeBeginIndex].Options = map[string]string{
+			continueTargetOption: strconv.Itoa(continueTarget - scopeBeginIndex),
+		}
+	}
 	for _, innerStmt := range stmt.Body {
 		if err := innerStmt.AppendOperations(ops, stack); err != nil {
 			return err