fix(instrument): place signal before terminal stmts inside branches

When a segment wraps a terminal statement (RETURN, RAISE EXCEPTION)
inside a control-flow keyword (IF/ELSIF/ELSE … THEN … RETURN …), the
coverage signal was placed after the entire segment — past the RETURN —
making it unreachable. Coverage for these branches was always reported
as 0 regardless of test execution.

Add findTerminalPos() to scan a segment for an embedded terminal
statement and return its byte offset. In emitSegment, use a three-way
dispatch:

  1. Segment starts with terminal → signal before (existing logic)
  2. Segment contains terminal inside → signal before the inner
     terminal (new logic)
  3. No terminal → signal after (existing logic)

This ensures every injected PERFORM pg_notify() is reachable by at
least one execution path.

Add indentOf() helper to extract indentation for the inner terminal
line (may differ from the segment's first-line indent).

Add tests: TestFindTerminalPos (8 cases),
TestInstrumentBody_ReturnInBranches,
TestInstrumentBody_RaiseExceptionInBranch, and
TestInstrumentBody_MixedTerminalNonTerminalBranches.

Author: pashagolub

internal/instrument/instrumenter.go

@@ -165,14 +165,21 @@ func instrumentBody(stmt *parser.Statement, filePath string, skipToBegin bool, n
 			}
 		}
 
-		notifyCall := fmt.Sprintf("%s%s pg_notify('pgcov', '%s');",
-			indent, notifyCmd, strings.ReplaceAll(cp.SignalID, "'", "''"))
-
-		if isTerminalSegment(segText) {
-			// Terminal statements (RETURN, RAISE EXCEPTION) exit the
-			// current scope — the signal must fire before the statement.
+		escapedSignal := strings.ReplaceAll(cp.SignalID, "'", "''")
+
+		if termPos := findTerminalPos(segText); termPos >= 0 {
+			// Segment contains (or starts with) a terminal statement
+			// (RETURN, RAISE EXCEPTION).  Place the signal before the
+			// terminal so it fires before the scope exits.
+			// When termPos == 0 the segment starts with the terminal;
+			// when termPos > 0 it is nested inside a control structure
+			// (e.g. IF … THEN RETURN …).
+			termIndent := indentOf(segText[termPos:])
+			notifyCall := fmt.Sprintf("%s%s pg_notify('pgcov', '%s');",
+				termIndent, notifyCmd, escapedSignal)
+			instrumentedBody.WriteString(segText[:termPos])
 			fmt.Fprintf(&instrumentedBody, "%s\n", notifyCall)
-			instrumentedBody.WriteString(segText)
+			instrumentedBody.WriteString(segText[termPos:])
 			lastWrittenPos = segEnd
 		} else {
 			// Non-terminal statements: emit the signal after the
@@ -186,6 +193,8 @@ func instrumentBody(stmt *parser.Statement, filePath string, skipToBegin bool, n
 			} else {
 				lastWrittenPos = segEnd
 			}
+			notifyCall := fmt.Sprintf("%s%s pg_notify('pgcov', '%s');",
+				indent, notifyCmd, escapedSignal)
 			fmt.Fprintf(&instrumentedBody, "\n%s", notifyCall)
 		}
 	}
@@ -274,49 +283,52 @@ func isExecutableSegment(segmentContent string) bool {
 	return true
 }
 
-// isTerminalSegment checks whether a segment starts with a terminal statement
-// (RETURN or RAISE EXCEPTION / bare RAISE) that exits the current scope.
-// NOTIFY calls placed after such statements would be unreachable.
-// Non-fatal RAISE levels (NOTICE, WARNING, INFO, LOG, DEBUG) are not terminal.
-func isTerminalSegment(segmentContent string) bool {
+// findTerminalPos scans segmentContent for a terminal statement (RETURN or
+// fatal RAISE) and returns its byte position within the string.  Returns -1
+// if no terminal statement is found.  This is used for segments that wrap a
+// terminal inside a control-flow keyword (e.g. IF/ELSIF/ELSE … RETURN …).
+func findTerminalPos(segmentContent string) int {
 	sc := pglex.NewScanner(segmentContent)
-
-	// Find the first non-comment token.
-	var first pglex.Token
 	for {
-		first = sc.Scan()
-		if first.Type == pglex.EOF {
-			return false
+		tok := sc.Scan()
+		if tok.Type == pglex.EOF {
+			return -1
 		}
-		if first.Type != pglex.Comment {
-			break
+		if tok.Type == pglex.KReturn {
+			return tok.Pos
 		}
-	}
-
-	if first.Type == pglex.KReturn {
-		return true
-	}
-
-	if first.Type == pglex.KRaise {
-		// RAISE is terminal unless followed by a non-fatal level.
-		for {
-			tok := sc.Scan()
-			if tok.Type == pglex.EOF {
-				return true // bare RAISE; — re-raise in exception handler
-			}
-			if tok.Type == pglex.Comment {
-				continue
-			}
-			switch tok.Type {
-			case pglex.KNotice, pglex.KWarning, pglex.KInfo, pglex.KLog, pglex.KDebug:
-				return false
-			default:
-				return true // RAISE EXCEPTION, RAISE 'message', etc.
+		if tok.Type == pglex.KRaise {
+			pos := tok.Pos
+			// Peek at the next non-comment token to decide fatality.
+			for {
+				next := sc.Scan()
+				if next.Type == pglex.EOF {
+					return pos // bare RAISE — re-raise
+				}
+				if next.Type == pglex.Comment {
+					continue
+				}
+				switch next.Type {
+				case pglex.KNotice, pglex.KWarning, pglex.KInfo, pglex.KLog, pglex.KDebug:
+					// Non-fatal RAISE; continue scanning for a later terminal.
+					break
+				default:
+					return pos // RAISE EXCEPTION, RAISE 'msg', etc.
+				}
+				break
 			}
 		}
 	}
+}
 
-	return false
+// indentOf returns the leading whitespace of the first non-empty line in s.
+func indentOf(s string) string {
+	for line := range strings.SplitSeq(s, "\n") {
+		if strings.TrimSpace(line) != "" {
+			return getIndentation(line)
+		}
+	}
+	return ""
 }
 
 // getIndentation returns the leading whitespace of a line.

internal/instrument/plpgsql_ast_test.go

@@ -238,11 +238,13 @@ $$ LANGUAGE plpgsql;`
 	t.Logf("Coverage points: %d (may be 0 for malformed SQL)", len(instrumented.Locations))
 }
 
-func TestIsTerminalSegment(t *testing.T) {
+func TestFindTerminalPos_StartingTerminals(t *testing.T) {
+	// Segments that start with a terminal should return pos 0.
+	// Segments without any terminal should return -1.
 	tests := []struct {
-		name     string
-		segment  string
-		terminal bool
+		name      string
+		segment   string
+		wantFound bool
 	}{
 		{"RETURN value", "RETURN a + b", true},
 		{"RETURN bare", "RETURN", true},
@@ -265,9 +267,11 @@ func TestIsTerminalSegment(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := isTerminalSegment(tt.segment)
-			if got != tt.terminal {
-				t.Errorf("isTerminalSegment(%q) = %v, want %v", tt.segment, got, tt.terminal)
+			got := findTerminalPos(tt.segment)
+			if tt.wantFound && got < 0 {
+				t.Errorf("findTerminalPos(%q) = -1, want >= 0", tt.segment)
+			} else if !tt.wantFound && got >= 0 {
+				t.Errorf("findTerminalPos(%q) = %d, want -1", tt.segment, got)
 			}
 		})
 	}
@@ -367,3 +371,204 @@ $$ LANGUAGE plpgsql;`
 
 	t.Log(instrumentedSQL)
 }
+
+func TestFindTerminalPos(t *testing.T) {
+	tests := []struct {
+		name      string
+		segment   string
+		wantFound bool   // whether a terminal is expected
+		wantText  string // expected text at the found position (prefix check)
+	}{
+		{"bare RETURN", "RETURN x", true, "RETURN"},
+		{"IF with RETURN", "IF x > 0 THEN\n        RETURN 1", true, "RETURN"},
+		{"ELSIF with RETURN", "\n    ELSIF x > 5 THEN\n        RETURN 2", true, "RETURN"},
+		{"ELSE with RETURN", "\n    ELSE\n        RETURN 3", true, "RETURN"},
+		{"no terminal", "x := x + 1", false, ""},
+		{"RAISE NOTICE", "RAISE NOTICE 'hello'", false, ""},
+		{"IF with RAISE EXCEPTION", "IF x < 0 THEN\n        RAISE EXCEPTION 'bad'", true, "RAISE"},
+		{"IF with non-terminal", "IF x > 0 THEN\n        x := 1", false, ""},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := findTerminalPos(tt.segment)
+			if tt.wantFound {
+				if got < 0 {
+					t.Fatalf("findTerminalPos(%q) = -1, want >= 0", tt.segment)
+				}
+				if !strings.HasPrefix(tt.segment[got:], tt.wantText) {
+					t.Errorf("findTerminalPos(%q) at %d: got %q, want prefix %q",
+						tt.segment, got, tt.segment[got:], tt.wantText)
+				}
+			} else {
+				if got >= 0 {
+					t.Errorf("findTerminalPos(%q) = %d, want -1", tt.segment, got)
+				}
+			}
+		})
+	}
+}
+
+func TestInstrumentBody_ReturnInBranches(t *testing.T) {
+	// B2 scenario: IF/ELSIF/ELSE with RETURN in each branch.
+	// All signals must be reachable (placed before the RETURN inside each branch).
+	sql := `CREATE OR REPLACE FUNCTION check_stock(v_stock INT)
+RETURNS TEXT AS $$
+BEGIN
+    IF v_stock = 0 THEN
+        RETURN 'out_of_stock';
+    ELSIF v_stock <= 10 THEN
+        RETURN 'low_stock';
+    ELSE
+        RETURN 'in_stock';
+    END IF;
+END;
+$$ LANGUAGE plpgsql;`
+
+	stmts := parser.ParseStatements(sql)
+	if len(stmts) == 0 {
+		t.Fatal("ParseStatements() returned no statements")
+	}
+
+	instrumentedSQL, coveragePoints := instrumentBody(stmts[0], "test.sql", true, "PERFORM")
+	if len(coveragePoints) != 3 {
+		t.Fatalf("expected 3 coverage points, got %d", len(coveragePoints))
+	}
+
+	// For each coverage point, verify the NOTIFY comes BEFORE the RETURN
+	// inside the branch—not after it (which would be unreachable).
+	returns := []string{"RETURN 'out_of_stock'", "RETURN 'low_stock'", "RETURN 'in_stock'"}
+	for i, cp := range coveragePoints {
+		notify := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", cp.SignalID)
+		notifyIdx := strings.Index(instrumentedSQL, notify)
+		returnIdx := strings.Index(instrumentedSQL, returns[i])
+		if notifyIdx < 0 || returnIdx < 0 {
+			t.Fatalf("cp %d: could not find notify or %s", i, returns[i])
+		}
+		if notifyIdx > returnIdx {
+			t.Errorf("cp %d (%s): NOTIFY at %d after RETURN at %d — signal is unreachable",
+				i, returns[i], notifyIdx, returnIdx)
+		}
+	}
+
+	// Also verify the instrumented SQL does NOT have a PERFORM between
+	// a RETURN and the next ELSIF/ELSE/END (that would be unreachable code).
+	for _, kw := range []string{"ELSIF", "ELSE", "END IF"} {
+		kwIdx := strings.Index(instrumentedSQL, kw)
+		if kwIdx < 0 {
+			continue
+		}
+		// Check a narrow window before the keyword for a rogue PERFORM.
+		before := instrumentedSQL[max(0, kwIdx-80):kwIdx]
+		// There should be a RETURN between the PERFORM and the keyword boundary.
+		lastPerform := strings.LastIndex(before, "PERFORM pg_notify")
+		lastReturn := strings.LastIndex(before, "RETURN")
+		if lastPerform >= 0 && lastReturn >= 0 && lastPerform > lastReturn {
+			t.Errorf("unreachable PERFORM found between RETURN and %s", kw)
+		}
+	}
+
+	t.Log(instrumentedSQL)
+}
+
+func TestInstrumentBody_RaiseExceptionInBranch(t *testing.T) {
+	// Segment: IF ... THEN RAISE EXCEPTION ... — terminal inside control structure.
+	sql := `CREATE OR REPLACE FUNCTION validate(x INT)
+RETURNS VOID AS $$
+BEGIN
+    IF x < 0 THEN
+        RAISE EXCEPTION 'negative: %', x;
+    ELSIF x = 0 THEN
+        RAISE EXCEPTION 'zero';
+    END IF;
+    RAISE NOTICE 'ok: %', x;
+END;
+$$ LANGUAGE plpgsql;`
+
+	stmts := parser.ParseStatements(sql)
+	if len(stmts) == 0 {
+		t.Fatal("ParseStatements() returned no statements")
+	}
+
+	instrumentedSQL, coveragePoints := instrumentBody(stmts[0], "test.sql", true, "PERFORM")
+	if len(coveragePoints) != 3 {
+		t.Fatalf("expected 3 coverage points, got %d", len(coveragePoints))
+	}
+
+	// First two are in IF/ELSIF branches with RAISE EXCEPTION (terminal).
+	// Signals must appear before the RAISE EXCEPTION.
+	for i, target := range []string{"RAISE EXCEPTION 'negative", "RAISE EXCEPTION 'zero"} {
+		notify := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", coveragePoints[i].SignalID)
+		notifyIdx := strings.Index(instrumentedSQL, notify)
+		stmtIdx := strings.Index(instrumentedSQL, target)
+		if notifyIdx < 0 || stmtIdx < 0 {
+			t.Fatalf("cp %d: could not find notify or %q", i, target)
+		}
+		if notifyIdx > stmtIdx {
+			t.Errorf("cp %d: NOTIFY at %d after %q at %d — unreachable", i, notifyIdx, target, stmtIdx)
+		}
+	}
+
+	// Third is RAISE NOTICE (non-terminal, standalone). Signal should come after.
+	notify2 := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", coveragePoints[2].SignalID)
+	notifyIdx2 := strings.Index(instrumentedSQL, notify2)
+	noticeIdx := strings.Index(instrumentedSQL, "RAISE NOTICE")
+	if notifyIdx2 < 0 || noticeIdx < 0 {
+		t.Fatal("could not find RAISE NOTICE or its notify")
+	}
+	if notifyIdx2 <= noticeIdx {
+		t.Errorf("RAISE NOTICE: NOTIFY at %d should come after statement at %d", notifyIdx2, noticeIdx)
+	}
+
+	t.Log(instrumentedSQL)
+}
+
+func TestInstrumentBody_MixedTerminalNonTerminalBranches(t *testing.T) {
+	// Branch with RETURN vs branch with assignment — signal placement differs.
+	sql := `CREATE OR REPLACE FUNCTION classify(x INT)
+RETURNS TEXT AS $$
+DECLARE
+    result TEXT;
+BEGIN
+    IF x > 0 THEN
+        result := 'positive';
+    ELSE
+        RETURN 'non-positive';
+    END IF;
+    RETURN result;
+END;
+$$ LANGUAGE plpgsql;`
+
+	stmts := parser.ParseStatements(sql)
+	if len(stmts) == 0 {
+		t.Fatal("ParseStatements() returned no statements")
+	}
+
+	instrumentedSQL, coveragePoints := instrumentBody(stmts[0], "test.sql", true, "PERFORM")
+	if len(coveragePoints) != 3 {
+		t.Fatalf("expected 3 coverage points, got %d", len(coveragePoints))
+	}
+
+	// cp0: IF ... result := 'positive' — no terminal, signal after.
+	assign := "result := 'positive'"
+	assignNotify := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", coveragePoints[0].SignalID)
+	if strings.Index(instrumentedSQL, assignNotify) < strings.Index(instrumentedSQL, assign) {
+		t.Error("assignment branch: NOTIFY should come AFTER the assignment")
+	}
+
+	// cp1: ELSE RETURN 'non-positive' — terminal inside branch, signal before.
+	ret1 := "RETURN 'non-positive'"
+	retNotify1 := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", coveragePoints[1].SignalID)
+	if strings.Index(instrumentedSQL, retNotify1) > strings.Index(instrumentedSQL, ret1) {
+		t.Error("ELSE RETURN branch: NOTIFY should come BEFORE the RETURN")
+	}
+
+	// cp2: standalone RETURN result — terminal at start, signal before.
+	ret2 := "RETURN result"
+	retNotify2 := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", coveragePoints[2].SignalID)
+	if strings.Index(instrumentedSQL, retNotify2) > strings.Index(instrumentedSQL, ret2) {
+		t.Error("standalone RETURN: NOTIFY should come BEFORE the RETURN")
+	}
+
+	t.Log(instrumentedSQL)
+}