feat(#275): comment preservation in AST and formatter (#311)

Co-authored-by: Ajit Pratap Singh <[email protected]>

Author: ajitpratap0

pkg/formatter/comment_preservation_test.go

@@ -0,0 +1,103 @@
+package formatter
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestFormat_LineCommentPreservation(t *testing.T) {
+	f := New(Options{Compact: true})
+
+	t.Run("leading line comment", func(t *testing.T) {
+		input := "-- header comment\nSELECT col FROM t"
+		result, err := f.Format(input)
+		if err != nil {
+			t.Fatalf("Format() error = %v", err)
+		}
+		if !strings.Contains(result, "-- header comment") {
+			t.Errorf("Expected leading comment preserved, got: %s", result)
+		}
+		if !strings.Contains(result, "SELECT") || !strings.Contains(result, "col") {
+			t.Errorf("Expected SQL preserved, got: %s", result)
+		}
+	})
+
+	t.Run("trailing line comment", func(t *testing.T) {
+		input := "SELECT col FROM t -- trailing"
+		result, err := f.Format(input)
+		if err != nil {
+			t.Fatalf("Format() error = %v", err)
+		}
+		if !strings.Contains(result, "-- trailing") {
+			t.Errorf("Expected trailing comment preserved, got: %s", result)
+		}
+	})
+}
+
+func TestFormat_BlockCommentPreservation(t *testing.T) {
+	f := New(Options{Compact: true})
+
+	t.Run("leading block comment", func(t *testing.T) {
+		input := "/* header */\nSELECT col FROM t"
+		result, err := f.Format(input)
+		if err != nil {
+			t.Fatalf("Format() error = %v", err)
+		}
+		if !strings.Contains(result, "/* header */") {
+			t.Errorf("Expected block comment preserved, got: %s", result)
+		}
+	})
+
+	t.Run("inline block comment", func(t *testing.T) {
+		input := "SELECT /* inline */ col FROM t"
+		result, err := f.Format(input)
+		if err != nil {
+			t.Fatalf("Format() error = %v", err)
+		}
+		if !strings.Contains(result, "/* inline */") {
+			t.Errorf("Expected inline block comment preserved, got: %s", result)
+		}
+	})
+}
+
+func TestFormat_CommentRoundTrip(t *testing.T) {
+	f := New(Options{Compact: true})
+
+	inputs := []string{
+		"-- header\nSELECT col FROM t",
+		"SELECT col FROM t -- trailing",
+		"/* block */ SELECT col FROM t",
+	}
+
+	for _, input := range inputs {
+		t.Run(input, func(t *testing.T) {
+			first, err := f.Format(input)
+			if err != nil {
+				t.Fatalf("First format error: %v", err)
+			}
+			second, err := f.Format(first)
+			if err != nil {
+				t.Fatalf("Second format error: %v", err)
+			}
+			if first != second {
+				t.Errorf("Round-trip mismatch:\n  first:  %q\n  second: %q", first, second)
+			}
+		})
+	}
+}
+
+func TestFormat_MultipleComments(t *testing.T) {
+	f := New(Options{Compact: true})
+
+	input := "-- first comment\n-- second comment\nSELECT col FROM t"
+	result, err := f.Format(input)
+	if err != nil {
+		t.Fatalf("Format() error = %v", err)
+	}
+	if !strings.Contains(result, "-- first comment") {
+		t.Errorf("Expected first comment, got: %s", result)
+	}
+	if !strings.Contains(result, "-- second comment") {
+		t.Errorf("Expected second comment, got: %s", result)
+	}
+}

pkg/formatter/formatter.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/ajitpratap0/GoSQLX/pkg/models"
 	"github.com/ajitpratap0/GoSQLX/pkg/sql/ast"
 	"github.com/ajitpratap0/GoSQLX/pkg/sql/parser"
 	"github.com/ajitpratap0/GoSQLX/pkg/sql/tokenizer"
@@ -53,13 +54,22 @@ func (f *Formatter) Format(sql string) (string, error) {
 		return "", nil
 	}
 
+	// Capture comments from tokenizer before parsing
+	comments := tkz.Comments
+
 	p := parser.NewParser()
 	parsedAST, err := p.ParseFromModelTokens(tokens)
 	if err != nil {
 		return "", fmt.Errorf("parsing failed: %w", err)
 	}
 	defer ast.ReleaseAST(parsedAST)
 
+	// Attach captured comments to AST
+	if len(comments) > 0 {
+		parsedAST.Comments = make([]models.Comment, len(comments))
+		copy(parsedAST.Comments, comments)
+	}
+
 	// Use AST's built-in Format method
 	style := ast.ReadableStyle()
 	if f.opts.Compact {

pkg/models/comment.go

@@ -0,0 +1,21 @@
+// Package models provides the Comment type for SQL comment preservation.
+package models
+
+// CommentStyle indicates the type of SQL comment.
+type CommentStyle int
+
+const (
+	// LineComment represents a -- single-line comment.
+	LineComment CommentStyle = iota
+	// BlockComment represents a /* multi-line */ comment.
+	BlockComment
+)
+
+// Comment represents a SQL comment captured during tokenization.
+type Comment struct {
+	Text   string       // The comment text including delimiters (e.g., "-- foo" or "/* bar */")
+	Style  CommentStyle // Line or block comment
+	Start  Location     // Start position in source
+	End    Location     // End position in source
+	Inline bool         // True if the comment is on the same line as code (trailing)
+}

pkg/sql/ast/ast.go

@@ -32,7 +32,11 @@
 // FILTER clause, RETURNING clause, JSON/JSONB operators, and FETCH FIRST/NEXT.
 package ast
 
-import "fmt"
+import (
+	"fmt"
+
+	"github.com/ajitpratap0/GoSQLX/pkg/models"
+)
 
 // Node represents any node in the Abstract Syntax Tree.
 //
@@ -1684,6 +1688,7 @@ func (p PartitionDefinition) Children() []Node {
 // AST represents the root of the Abstract Syntax Tree
 type AST struct {
 	Statements []Statement
+	Comments   []models.Comment // Comments captured during tokenization, preserved during formatting
 }
 
 func (a AST) TokenLiteral() string { return "" }

pkg/sql/ast/format.go

@@ -138,7 +138,34 @@ func (a AST) Format(opts FormatOptions) string {
 		// Each statement already gets semicolons from its own Format
 		sep = "\n"
 	}
-	return strings.Join(parts, sep)
+	result := strings.Join(parts, sep)
+
+	// Emit preserved comments around the formatted SQL
+	if len(a.Comments) > 0 {
+		var leading, trailing []string
+		for _, c := range a.Comments {
+			if c.Inline {
+				// Inline comments (on same line as code) → trailing
+				trailing = append(trailing, c.Text)
+			} else {
+				// Comments on their own line → leading
+				leading = append(leading, c.Text)
+			}
+		}
+		var sb strings.Builder
+		for _, lc := range leading {
+			sb.WriteString(lc)
+			sb.WriteString("\n")
+		}
+		sb.WriteString(result)
+		for _, tc := range trailing {
+			sb.WriteString(" ")
+			sb.WriteString(tc)
+		}
+		result = sb.String()
+	}
+
+	return result
 }
 
 // Format returns formatted SQL for a SelectStatement.

pkg/sql/ast/pool.go

@@ -353,6 +353,11 @@ func ReleaseAST(ast *AST) {
 	// Reset slice but keep capacity
 	ast.Statements = ast.Statements[:0]
 
+	// Reset comments but keep capacity
+	if cap(ast.Comments) > 0 {
+		ast.Comments = ast.Comments[:0]
+	}
+
 	// Return to pool
 	astPool.Put(ast)
 }

pkg/sql/tokenizer/pool.go

@@ -161,4 +161,9 @@ func (t *Tokenizer) Reset() {
 
 	// Don't reset keywords as they're constant
 	t.logger = nil
+
+	// Preserve Comments slice capacity but reset length
+	if cap(t.Comments) > 0 {
+		t.Comments = t.Comments[:0]
+	}
 }

pkg/sql/tokenizer/tokenizer.go

@@ -245,6 +245,7 @@ type Tokenizer struct {
 	line       int                // Current line number (1-based)
 	keywords   *keywords.Keywords // Keyword classifier for token type determination
 	logger     *slog.Logger       // Optional structured logger for verbose tracing
+	Comments   []models.Comment   // Comments captured during tokenization
 }
 
 // New creates a new Tokenizer with default configuration and keyword support.
@@ -1233,6 +1234,8 @@ func (t *Tokenizer) readPunctuation() (models.Token, error) {
 			}
 			// Check for line comment: --
 			if nxtR == '-' {
+				commentStartIdx := t.pos.Index - size // back to first '-'
+				commentStartPos := t.toSQLPosition(Position{Index: commentStartIdx})
 				t.pos.AdvanceRune(nxtR, nxtSize)
 				// Skip until end of line or EOF
 				for t.pos.Index < len(t.input) {
@@ -1243,6 +1246,19 @@ func (t *Tokenizer) readPunctuation() (models.Token, error) {
 					}
 					t.pos.AdvanceRune(cr, csize)
 				}
+				commentEndIdx := t.pos.Index
+				// Trim trailing newline from comment text
+				textEnd := commentEndIdx
+				if textEnd > 0 && t.input[textEnd-1] == '\n' {
+					textEnd--
+				}
+				t.Comments = append(t.Comments, models.Comment{
+					Text:   string(t.input[commentStartIdx:textEnd]),
+					Style:  models.LineComment,
+					Start:  commentStartPos,
+					End:    t.toSQLPosition(t.pos),
+					Inline: t.hasCodeBeforeOnLine(commentStartIdx),
+				})
 				// Return the next token (skip the comment)
 				t.skipWhitespace()
 				return t.nextToken()
@@ -1258,6 +1274,8 @@ func (t *Tokenizer) readPunctuation() (models.Token, error) {
 			nxtR, nxtSize := utf8.DecodeRune(t.input[t.pos.Index:])
 			// Check for block comment: /*
 			if nxtR == '*' {
+				commentStartIdx := t.pos.Index - size // back to '/'
+				commentStartPos := t.toSQLPosition(Position{Index: commentStartIdx})
 				t.pos.AdvanceRune(nxtR, nxtSize)
 				// Skip until */ or EOF
 				for t.pos.Index < len(t.input) {
@@ -1275,6 +1293,13 @@ func (t *Tokenizer) readPunctuation() (models.Token, error) {
 						t.pos.AdvanceRune(cr, csize)
 					}
 				}
+				t.Comments = append(t.Comments, models.Comment{
+					Text:   string(t.input[commentStartIdx:t.pos.Index]),
+					Style:  models.BlockComment,
+					Start:  commentStartPos,
+					End:    t.toSQLPosition(t.pos),
+					Inline: t.hasCodeBeforeOnLine(commentStartIdx),
+				})
 				// Return the next token (skip the comment)
 				t.skipWhitespace()
 				return t.nextToken()
@@ -1636,3 +1661,23 @@ func (t *Tokenizer) getLocation(pos int) models.Location {
 func isIdentifierChar(r rune) bool {
 	return isUnicodeIdentifierPart(r)
 }
+
+// hasCodeBeforeOnLine checks if there are non-whitespace characters on the same
+// line before the given byte index. Used to determine if a comment is inline.
+func (t *Tokenizer) hasCodeBeforeOnLine(idx int) bool {
+	// Find the start of the line containing idx
+	lineStart := 0
+	for i := len(t.lineStarts) - 1; i >= 0; i-- {
+		if t.lineStarts[i] <= idx {
+			lineStart = t.lineStarts[i]
+			break
+		}
+	}
+	// Check for non-whitespace between lineStart and idx
+	for i := lineStart; i < idx && i < len(t.input); i++ {
+		if t.input[i] != ' ' && t.input[i] != '\t' && t.input[i] != '\r' {
+			return true
+		}
+	}
+	return false
+}