fix: recover RETURN NEXT variable when retvarno is lost by libpg-query

libpg-query's protobuf does not serialize the retvarno field of
PLpgSQL_stmt_return_next, causing RETURN NEXT v_entry to deparse as
bare RETURN NEXT.

Fix:
- Thread enclosingForVarName through FOR loop contexts so the deparser
  knows the loop variable name
- Add inferReturnNextVar() that recovers the variable from:
  1. Enclosing FOR loop variable (most common pattern)
  2. Single user-declared variable (SETOF scalar, no loop)
- Only infer when returnInfo is provided (kind='setof'), preventing
  false inference on OUT-param functions where bare RETURN NEXT is
  correct
- Add extractReturnInfo() to test infrastructure to extract return
  type from CREATE FUNCTION SQL for round-trip tests

Tests:
- 4 new snapshot tests for RETURN NEXT retvarno recovery
- 2 new SQL fixtures (Tests 47-48) in plpgsql_deparser_fixes.sql
- All 236 fixtures round-trip correctly

Author: pyramation

__fixtures__/plpgsql-generated/generated.json

@@ -145,6 +145,8 @@
   "plpgsql_deparser_fixes-44.sql": "-- Test 44: Nested COALESCE expressions\nCREATE FUNCTION test_nested_coalesce(a text, b text, c text) RETURNS text\nLANGUAGE plpgsql AS $$\nBEGIN\n  RETURN COALESCE(a, COALESCE(b, COALESCE(c, 'fallback')));\nEND$$",
   "plpgsql_deparser_fixes-45.sql": "-- Test 45: SELECT INTO with COALESCE and function call in FROM\nCREATE FUNCTION test_select_into_from_srf(v_data jsonb) RETURNS jsonb\nLANGUAGE plpgsql AS $$\nDECLARE\n  v_first jsonb;\nBEGIN\n  SELECT elem INTO v_first\n  FROM jsonb_array_elements(v_data) AS elem\n  LIMIT 1;\n  RETURN v_first;\nEND$$",
   "plpgsql_deparser_fixes-46.sql": "-- Test 46: FOR loop with jsonb_array_elements and WHERE filter\nCREATE FUNCTION test_for_srf_with_filter(v_config jsonb, v_key text) RETURNS jsonb\nLANGUAGE plpgsql AS $$\nDECLARE\n  v_entry jsonb;\nBEGIN\n  FOR v_entry IN\n    SELECT elem FROM jsonb_array_elements(v_config) AS elem\n    WHERE elem->>'key' = v_key\n  LOOP\n    RETURN v_entry;\n  END LOOP;\n  RETURN NULL;\nEND$$",
+  "plpgsql_deparser_fixes-47.sql": "-- Test 47: RETURN NEXT with FOR loop variable (retvarno recovery)\n-- Exercises the case where libpg-query drops retvarno from PLpgSQL_stmt_return_next\nCREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb\nLANGUAGE plpgsql AS $$\nDECLARE\n  v_entry jsonb;\nBEGIN\n  FOR v_entry IN SELECT jsonb_array_elements(v_data)\n  LOOP\n    RETURN NEXT v_entry;\n  END LOOP;\n  RETURN;\nEND$$",
+  "plpgsql_deparser_fixes-48.sql": "-- Test 48: RETURN NEXT with declared variable (no FOR loop, single candidate)\nCREATE FUNCTION test_return_next_declared_var() RETURNS SETOF int\nLANGUAGE plpgsql AS $$\nDECLARE\n  v_val int := 42;\nBEGIN\n  RETURN NEXT v_val;\nEND$$",
   "plpgsql_control-1.sql": "--\n-- Tests for PL/pgSQL control structures\n--\n\n-- integer FOR loop\n\ndo $$\nbegin\n  -- basic case\n  for i in 1..3 loop\n    raise notice '1..3: i = %', i;\n  end loop;\n  -- with BY, end matches exactly\n  for i in 1..10 by 3 loop\n    raise notice '1..10 by 3: i = %', i;\n  end loop;\n  -- with BY, end does not match\n  for i in 1..11 by 3 loop\n    raise notice '1..11 by 3: i = %', i;\n  end loop;\n  -- zero iterations\n  for i in 1..0 by 3 loop\n    raise notice '1..0 by 3: i = %', i;\n  end loop;\n  -- REVERSE\n  for i in reverse 10..0 by 3 loop\n    raise notice 'reverse 10..0 by 3: i = %', i;\n  end loop;\n  -- potential overflow\n  for i in 2147483620..2147483647 by 10 loop\n    raise notice '2147483620..2147483647 by 10: i = %', i;\n  end loop;\n  -- potential overflow, reverse direction\n  for i in reverse -2147483620..-2147483647 by 10 loop\n    raise notice 'reverse -2147483620..-2147483647 by 10: i = %', i;\n  end loop;\nend$$",
   "plpgsql_control-2.sql": "-- BY can't be zero or negative\ndo $$\nbegin\n  for i in 1..3 by 0 loop\n    raise notice '1..3 by 0: i = %', i;\n  end loop;\nend$$",
   "plpgsql_control-3.sql": "do $$\nbegin\n  for i in 1..3 by -1 loop\n    raise notice '1..3 by -1: i = %', i;\n  end loop;\nend$$",

__fixtures__/plpgsql/plpgsql_deparser_fixes.sql

@@ -592,3 +592,26 @@ BEGIN
   END LOOP;
   RETURN NULL;
 END$$;
+
+-- Test 47: RETURN NEXT with FOR loop variable (retvarno recovery)
+-- Exercises the case where libpg-query drops retvarno from PLpgSQL_stmt_return_next
+CREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb
+LANGUAGE plpgsql AS $$
+DECLARE
+  v_entry jsonb;
+BEGIN
+  FOR v_entry IN SELECT jsonb_array_elements(v_data)
+  LOOP
+    RETURN NEXT v_entry;
+  END LOOP;
+  RETURN;
+END$$;
+
+-- Test 48: RETURN NEXT with declared variable (no FOR loop, single candidate)
+CREATE FUNCTION test_return_next_declared_var() RETURNS SETOF int
+LANGUAGE plpgsql AS $$
+DECLARE
+  v_val int := 42;
+BEGIN
+  RETURN NEXT v_val;
+END$$;

packages/plpgsql-deparser/__tests__/__snapshots__/deparser-fixes.test.ts.snap

@@ -183,6 +183,48 @@ exports[`plpgsql-deparser bug fixes PERFORM SELECT fix should strip SELECT keywo
 END"
 `;
 
+exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should leave RETURN NEXT bare for OUT-param functions 1`] = `
+"BEGIN
+  FOR i IN 1..5 LOOP
+      x := i;
+      y := 'item_' || i::text;
+      RETURN NEXT;
+  END LOOP;
+  RETURN;
+END"
+`;
+
+exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should leave RETURN NEXT bare when returnInfo is not provided 1`] = `
+"DECLARE
+  v_entry jsonb;
+BEGIN
+  FOR v_entry IN SELECT jsonb_array_elements(v_data) LOOP
+      RETURN NEXT;
+  END LOOP;
+  RETURN;
+END"
+`;
+
+exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should recover FOR loop variable in RETURN NEXT for SETOF functions 1`] = `
+"DECLARE
+  v_entry jsonb;
+BEGIN
+  FOR v_entry IN SELECT jsonb_array_elements(v_data) LOOP
+      RETURN NEXT v_entry;
+  END LOOP;
+  RETURN;
+END"
+`;
+
+exports[`plpgsql-deparser bug fixes RETURN NEXT retvarno recovery should recover single declared variable in RETURN NEXT for SETOF functions 1`] = `
+"DECLARE
+  v_val int := 42;
+BEGIN
+  RETURN NEXT v_val;
+  RETURN;
+END"
+`;
+
 exports[`plpgsql-deparser bug fixes Record field qualification (recfield) should handle OLD and NEW record references 1`] = `
 "BEGIN
   IF OLD.status <> NEW.status THEN

packages/plpgsql-deparser/__tests__/deparser-fixes.test.ts

@@ -1021,4 +1021,82 @@ END$$`;
       expect(deparsed).toContain('now()');
     });
   });
+
+  describe('RETURN NEXT retvarno recovery', () => {
+    it('should recover FOR loop variable in RETURN NEXT for SETOF functions', async () => {
+      const sql = `CREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb
+LANGUAGE plpgsql AS $$
+DECLARE
+  v_entry jsonb;
+BEGIN
+  FOR v_entry IN SELECT jsonb_array_elements(v_data)
+  LOOP
+    RETURN NEXT v_entry;
+  END LOOP;
+  RETURN;
+END$$`;
+
+      await testUtils.expectAstMatch('RETURN NEXT FOR loop var', sql);
+      const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
+      const deparsed = deparseSync(parsed, undefined, { kind: 'setof' });
+      expect(deparsed).toMatchSnapshot();
+      expect(deparsed).toContain('RETURN NEXT v_entry');
+    });
+
+    it('should recover single declared variable in RETURN NEXT for SETOF functions', async () => {
+      const sql = `CREATE FUNCTION test_return_next_declared_var() RETURNS SETOF int
+LANGUAGE plpgsql AS $$
+DECLARE
+  v_val int := 42;
+BEGIN
+  RETURN NEXT v_val;
+END$$`;
+
+      await testUtils.expectAstMatch('RETURN NEXT declared var', sql);
+      const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
+      const deparsed = deparseSync(parsed, undefined, { kind: 'setof' });
+      expect(deparsed).toMatchSnapshot();
+      expect(deparsed).toContain('RETURN NEXT v_val');
+    });
+
+    it('should leave RETURN NEXT bare for OUT-param functions', async () => {
+      const sql = `CREATE FUNCTION test_return_next_out(OUT x integer, OUT y text) RETURNS SETOF record
+LANGUAGE plpgsql AS $$
+BEGIN
+  FOR i IN 1..5 LOOP
+    x := i;
+    y := 'item_' || i::text;
+    RETURN NEXT;
+  END LOOP;
+  RETURN;
+END$$`;
+
+      await testUtils.expectAstMatch('RETURN NEXT OUT params', sql);
+      const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
+      const deparsed = deparseSync(parsed, undefined, { kind: 'out_params' });
+      expect(deparsed).toMatchSnapshot();
+      expect(deparsed).toContain('RETURN NEXT;');
+      expect(deparsed).not.toMatch(/RETURN NEXT\s+\w/);
+    });
+
+    it('should leave RETURN NEXT bare when returnInfo is not provided', async () => {
+      const sql = `CREATE FUNCTION test_return_next_for_var(v_data jsonb) RETURNS SETOF jsonb
+LANGUAGE plpgsql AS $$
+DECLARE
+  v_entry jsonb;
+BEGIN
+  FOR v_entry IN SELECT jsonb_array_elements(v_data)
+  LOOP
+    RETURN NEXT v_entry;
+  END LOOP;
+  RETURN;
+END$$`;
+
+      const parsed = parsePlPgSQLSync(sql) as unknown as PLpgSQLParseResult;
+      // Without returnInfo, inference should NOT happen (safety)
+      const deparsed = deparseSync(parsed);
+      expect(deparsed).toMatchSnapshot();
+      expect(deparsed).toContain('RETURN NEXT;');
+    });
+  });
 });

packages/plpgsql-deparser/src/plpgsql-deparser.ts

@@ -86,6 +86,8 @@ export interface PLpgSQLDeparserContext {
   loopVarLinenos?: Set<number>;
   /** Map of block lineno to the set of datum indices that belong to that block */
   blockDatumMap?: Map<number, Set<number>>;
+  /** Name of the enclosing FOR loop variable (used to recover RETURN NEXT <var> when retvarno is missing) */
+  enclosingForVarName?: string;
 }
 
 /**
@@ -1179,7 +1181,7 @@ export class PLpgSQLDeparser {
     }
 
     if (fori.body) {
-      const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
+      const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
       for (const stmt of fori.body) {
         const stmtStr = this.deparseStmt(stmt, bodyContext);
         parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));
@@ -1210,7 +1212,7 @@ export class PLpgSQLDeparser {
     parts.push(`${kw('FOR')} ${varName} ${kw('IN')} ${query} ${kw('LOOP')}`);
 
     if (fors.body) {
-      const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
+      const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
       for (const stmt of fors.body) {
         const stmtStr = this.deparseStmt(stmt, bodyContext);
         parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));
@@ -1246,7 +1248,7 @@ export class PLpgSQLDeparser {
     parts.push(`${forClause} ${kw('LOOP')}`);
 
     if (forc.body) {
-      const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
+      const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
       for (const stmt of forc.body) {
         const stmtStr = this.deparseStmt(stmt, bodyContext);
         parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));
@@ -1283,7 +1285,7 @@ export class PLpgSQLDeparser {
     parts.push(`${kw('FOREACH')} ${varName}${sliceClause} ${kw('IN ARRAY')} ${expr} ${kw('LOOP')}`);
 
     if (foreach.body) {
-      const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
+      const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
       for (const stmt of foreach.body) {
         const stmtStr = this.deparseStmt(stmt, bodyContext);
         parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));
@@ -1363,22 +1365,83 @@ export class PLpgSQLDeparser {
 
   /**
    * Deparse a RETURN NEXT statement
+   *
+   * libpg-query does not serialize `retvarno` for PLpgSQL_stmt_return_next,
+   * so when a function has `RETURN NEXT variable`, the variable reference is
+   * lost.  We recover it by:
+   *   1. Using the enclosing FOR loop variable (most common pattern)
+   *   2. Scanning datums for the sole user-declared variable (non-loop case)
+   * Bare `RETURN NEXT` remains valid for TABLE / OUT-param functions.
    */
   private deparseReturnNext(ret: PLpgSQL_stmt_return_next, context: PLpgSQLDeparserContext): string {
     const kw = this.keyword;
-    
+
     if (ret.expr) {
       return `${kw('RETURN NEXT')} ${this.deparseExpr(ret.expr)}`;
     }
-    
+
     if (ret.retvarno !== undefined && ret.retvarno >= 0) {
       const varName = this.getVarName(ret.retvarno, context);
       return `${kw('RETURN NEXT')} ${varName}`;
     }
-    
+
+    // retvarno is missing (libpg-query serialization gap).  Try to recover.
+    const inferred = this.inferReturnNextVar(context);
+    if (inferred) {
+      return `${kw('RETURN NEXT')} ${inferred}`;
+    }
+
     return kw('RETURN NEXT');
   }
 
+  /**
+   * Attempt to infer the variable for a bare RETURN NEXT when retvarno is
+   * not available.  Returns the variable name or undefined.
+   *
+   * We only infer when `returnInfo` tells us the function returns SETOF
+   * scalar.  Without returnInfo we cannot distinguish a legitimate bare
+   * `RETURN NEXT` (OUT-param / TABLE function) from a dropped retvarno,
+   * so we leave the statement bare to avoid incorrect output.
+   */
+  private inferReturnNextVar(context: PLpgSQLDeparserContext): string | undefined {
+    // Without return-type context we cannot safely infer
+    if (!context.returnInfo) {
+      return undefined;
+    }
+
+    // TABLE / OUT-param functions legitimately use bare RETURN NEXT
+    if (context.returnInfo.kind === 'out_params') {
+      return undefined;
+    }
+
+    // Only attempt recovery for SETOF scalar functions
+    if (context.returnInfo.kind !== 'setof') {
+      return undefined;
+    }
+
+    // 1. Inside a FOR loop → use the loop variable
+    if (context.enclosingForVarName) {
+      return context.enclosingForVarName;
+    }
+
+    // 2. Outside a FOR loop → look for a single user-declared var
+    if (context.datums) {
+      const candidates = context.datums.filter((d) => {
+        if ('PLpgSQL_var' in d) {
+          const v = d.PLpgSQL_var;
+          // Skip implicit 'found' and function parameters (no lineno)
+          return v.refname !== 'found' && v.lineno !== undefined;
+        }
+        return false;
+      });
+      if (candidates.length === 1 && 'PLpgSQL_var' in candidates[0]) {
+        return candidates[0].PLpgSQL_var.refname;
+      }
+    }
+
+    return undefined;
+  }
+
   /**
    * Deparse a RETURN QUERY statement
    */
@@ -1718,7 +1781,7 @@ export class PLpgSQLDeparser {
     parts.push(`${forClause} ${kw('LOOP')}`);
 
     if (fors.body) {
-      const bodyContext = { ...context, indentLevel: context.indentLevel + 1 };
+      const bodyContext = { ...context, indentLevel: context.indentLevel + 1, enclosingForVarName: varName };
       for (const stmt of fors.body) {
         const stmtStr = this.deparseStmt(stmt, bodyContext);
         parts.push(this.indent(stmtStr + ';', bodyContext.indentLevel));

packages/plpgsql-deparser/test-utils/index.ts

@@ -1,5 +1,5 @@
 import { parsePlPgSQL, parsePlPgSQLSync } from '@libpg-query/parser';
-import { deparseSync, PLpgSQLParseResult } from '../src';
+import { deparseSync, PLpgSQLParseResult, ReturnInfo } from '../src';
 import { readFileSync, readdirSync, existsSync } from 'fs';
 import * as path from 'path';
 import { diff } from 'jest-diff';
@@ -328,6 +328,43 @@ function reconstructSql(originalSql: string, newBody: string): string {
   return parts.prefix + newBody + parts.suffix;
 }
 
+/**
+ * Extract ReturnInfo from a CREATE FUNCTION SQL string by inspecting the
+ * RETURNS clause and OUT/INOUT parameters.  Used by the round-trip test to
+ * give the deparser enough context to recover dropped retvarno fields.
+ *
+ * We intentionally only extract for SETOF and OUT-param functions — these
+ * are the cases where RETURN NEXT inference needs context.  For other
+ * return types (scalar, void, trigger) we return undefined to avoid
+ * changing existing RETURN statement behaviour (e.g. bare RETURN in scalar
+ * functions would become RETURN NULL if we provided returnInfo).
+ */
+export function extractReturnInfo(sql: string): ReturnInfo | undefined {
+  // Normalise to a single line for simpler matching
+  const norm = sql.replace(/\s+/g, ' ').trim();
+
+  // Strip the body ($$...$$) so we only inspect the signature
+  const sig = norm.replace(/\$\$.*\$\$/s, '');
+
+  // OUT / INOUT parameters → out_params (prevents false RETURN NEXT inference)
+  if (/\b(OUT|INOUT)\s+/i.test(sig)) {
+    return { kind: 'out_params' };
+  }
+
+  // RETURNS TABLE(...) → out_params
+  if (/RETURNS\s+TABLE\s*\(/i.test(sig)) {
+    return { kind: 'out_params' };
+  }
+
+  // RETURNS SETOF <type> → needed for RETURN NEXT variable recovery
+  if (/RETURNS\s+SETOF\b/i.test(sig)) {
+    return { kind: 'setof' };
+  }
+
+  // All other types: don't provide returnInfo to preserve existing behaviour
+  return undefined;
+}
+
 export class PLpgSQLTestUtils {
   protected printErrorMessage(sql: string, position: number) {
     const lineNumber = sql.slice(0, position).match(/\n/g)?.length || 0;
@@ -368,7 +405,8 @@ export class PLpgSQLTestUtils {
         throw createParseError('PARSE_FAILED', testName, sql);
       }
 
-      const deparsedBody = deparseSync(originalAst);
+      const returnInfo = extractReturnInfo(sql);
+      const deparsedBody = deparseSync(originalAst, undefined, returnInfo);
 
       if (!deparsedBody || deparsedBody.trim().length === 0) {
         throw createParseError('DEPARSE_FAILED', testName, sql, deparsedBody);