fix: make sure coverage point starts with the current token start

Before that coverage point spans from the end of previous token till
the end of the last token in the statement

Author: pashagolub

internal/instrument/instrumenter.go

@@ -201,13 +201,12 @@ type executableSegment struct {
 func findExecutableSegments(bodyContent string, tokens []*pgquery.ScanToken) []executableSegment {
 	var segments []executableSegment
 
-	segmentStart := 0
 	hasExecutableContent := false
+	firstExecutableTokenPos := -1
 
 	for idx, token := range tokens {
 		if token.Token == pgquery.Token_BEGIN_P { // Skip until BEGIN token
 			tokens = tokens[idx+1:]
-			segmentStart = int(token.End)
 			break
 		}
 	}
@@ -220,39 +219,42 @@ func findExecutableSegments(bodyContent string, tokens []*pgquery.ScanToken) []e
 
 		// Check if this is a semicolon (statement separator)
 		if token.Token == pgquery.Token_ASCII_59 { // Token_ASCII_59 - ";"
-			if hasExecutableContent {
+			if hasExecutableContent && firstExecutableTokenPos >= 0 {
 				// Check if this segment represents an executable statement
 				segmentEnd := int(token.Start)
-				segmentContent := bodyContent[segmentStart:segmentEnd]
+				segmentContent := bodyContent[firstExecutableTokenPos:segmentEnd]
 
 				if isExecutableSegment(segmentContent) {
 					segment := executableSegment{
-						startPos:  segmentStart,
+						startPos:  firstExecutableTokenPos,
 						endPos:    segmentEnd,
-						lineStart: convertByteOffsetToLine(bodyContent, segmentStart),
+						lineStart: convertByteOffsetToLine(bodyContent, firstExecutableTokenPos),
 						lineEnd:   convertByteOffsetToLine(bodyContent, segmentEnd),
 					}
 					segments = append(segments, segment)
 				}
 			}
 
 			// Reset for next segment
-			segmentStart = int(token.End)
 			hasExecutableContent = false
+			firstExecutableTokenPos = -1
 		} else {
 			// This is some non-comment token, so we have content
+			if !hasExecutableContent {
+				firstExecutableTokenPos = int(token.Start)
+			}
 			hasExecutableContent = true
 		}
 	}
 
 	// Handle the last segment if there's remaining content
-	if hasExecutableContent && segmentStart < len(bodyContent) {
-		segmentContent := bodyContent[segmentStart:]
+	if hasExecutableContent && firstExecutableTokenPos >= 0 && firstExecutableTokenPos < len(bodyContent) {
+		segmentContent := bodyContent[firstExecutableTokenPos:]
 		if isExecutableSegment(segmentContent) {
 			segment := executableSegment{
-				startPos:  segmentStart,
+				startPos:  firstExecutableTokenPos,
 				endPos:    len(bodyContent),
-				lineStart: convertByteOffsetToLine(bodyContent, segmentStart),
+				lineStart: convertByteOffsetToLine(bodyContent, firstExecutableTokenPos),
 				lineEnd:   convertByteOffsetToLine(bodyContent, len(bodyContent)),
 			}
 			segments = append(segments, segment)

internal/instrument/location.go

@@ -161,18 +161,3 @@ func GetUniqueFiles(locations []CoveragePoint) []string {
 	}
 	return files
 }
-
-// ConvertPositionToLine converts a byte position in source text to a line number (1-indexed)
-func ConvertPositionToLine(sourceText string, startPos int) int {
-	if startPos < 0 || startPos > len(sourceText) {
-		return 1
-	}
-
-	line := 1
-	for i := 0; i < startPos && i < len(sourceText); i++ {
-		if sourceText[i] == '\n' {
-			line++
-		}
-	}
-	return line
-}

internal/instrument/plpgsql_ast_test.go

@@ -62,16 +62,45 @@ $$ LANGUAGE plpgsql;`
 		t.Errorf("Expected 4 coverage points, got %d", len(instrumented.Locations))
 	}
 
-	// Verify coverage points are at the correct lines (by converting positions to lines)
-	expectedLines := []int{7, 9, 11, 14}
+	// Verify coverage points are positioned at the expected executable statements
+	sqlContent, _ := os.ReadFile(tmpFile)
+	sqlText := string(sqlContent)
+	
+	// Expected executable statements that should have coverage points
+	expectedStatements := []string{
+		"discount_rate := 0.20",
+		"discount_rate := 0.10", 
+		"discount_rate := 0.05",
+		"RETURN total_amount * discount_rate",
+	}
+	
+	if len(instrumented.Locations) != len(expectedStatements) {
+		t.Errorf("Expected %d coverage points, got %d", len(expectedStatements), len(instrumented.Locations))
+	}
+	
 	for i, cp := range instrumented.Locations {
-		// Convert position to line number for validation
-		// For this test, we'll use the original SQL content
-		sqlContent, _ := os.ReadFile(tmpFile)
-		actualLine := ConvertPositionToLine(string(sqlContent), cp.StartPos)
-		if actualLine != expectedLines[i] {
-			t.Errorf("Coverage point %d: expected line %d, got %d", i, expectedLines[i], actualLine)
+		if i >= len(expectedStatements) {
+			break
 		}
+		
+		// Extract the actual code at the coverage point position
+		if cp.StartPos < 0 || cp.StartPos >= len(sqlText) {
+			t.Errorf("Coverage point %d: invalid start position %d", i, cp.StartPos)
+			continue
+		}
+		
+		// Get a reasonable chunk of text around the coverage point to verify it's correct
+		endPos := cp.StartPos + cp.Length
+		if endPos > len(sqlText) {
+			endPos = len(sqlText)
+		}
+		actualText := strings.TrimSpace(sqlText[cp.StartPos:endPos])
+		
+		// Check if the coverage point contains the expected statement
+		if !strings.Contains(actualText, expectedStatements[i]) {
+			t.Errorf("Coverage point %d: expected to contain %q, got %q", i, expectedStatements[i], actualText)
+		}
+		
 		if cp.ImplicitCoverage {
 			t.Errorf("Coverage point %d: should not be implicit", i)
 		}