gas-stack/pkg/codegen/modelgenerate/ast_helpers.go

package modelgenerate

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"sort"
)

// These just need to be unique (i.e., something that will never be a real function in actual code).
const (
	commentMarker   = "__comment__jafkjewfkajwlekfjawlejf"
	blankLineMarker = "__blank_line__awefjakwlefjkwalkefj"
)

// TrailingComments is a side map for attaching end-of-line comments to AST nodes.
// Generators populate this map; FprintWithComments consumes and clears it.
// For a regular statement, the comment appears at the end of the statement's line.
// For a BlockStmt, the comment appears after the closing brace.
//
// This is a terrible implementation (global variable), but Node is an interface (pointer type),
// so there's no risk of cross-contamination really.  Still bad for concurrent use, but hopefully
// we don't have to do that.
var TrailingComments = map[ast.Node]string{}

// mustCall wraps a call expression in Must(...), producing AST for Must(inner).
func mustCall(inner ast.Expr) *ast.CallExpr {
	return &ast.CallExpr{
		Fun:  ast.NewIdent("Must"),
		Args: []ast.Expr{inner},
	}
}

// Comment creates a marker statement that will be converted to a real Go comment
// in the generated output by FprintWithComments.
func Comment(text string) ast.Stmt {
	return &ast.ExprStmt{X: &ast.CallExpr{
		Fun:  ast.NewIdent(commentMarker),
		Args: []ast.Expr{&ast.BasicLit{Kind: token.STRING, Value: `"` + text + `"`}},
	}}
}

// BlankLine creates a marker statement that will be converted to a blank line
// in the generated output.  The marker occupies a line in the first-pass output;
// when removed, the position gap causes the printer to insert a blank line.
func BlankLine() *ast.ExprStmt {
	return &ast.ExprStmt{X: &ast.CallExpr{
		Fun: ast.NewIdent(blankLineMarker),
	}}
}

// FprintWithComments prints an ast.File, converting Comment/BlankLine markers
// into real Go comments and blank lines, preserving Doc comments, and applying
// trailing comments from the TrailingComments side map.
//
// It does a round-trip (print -> parse -> modify -> print) to obtain real token
// positions, which Go's comment system requires.
func FprintWithComments(w io.Writer, file *ast.File) error {
	// First pass: print to buffer
	// Doc comments on FuncDecl/GenDecl are NOT added to file.Comments here;
	// the printer handles them via setComment(d.Doc) when visiting each decl.
	var buf bytes.Buffer
	fset := token.NewFileSet()
	if err := printer.Fprint(&buf, fset, file); err != nil {
		return fmt.Errorf("initial pretty-printing to get positioning: %w", err)
	}

	// Re-parse to get real positions (ParseComments preserves doc comments)
	fset = token.NewFileSet()
	parsed, err := parser.ParseFile(fset, "", buf.Bytes(), parser.ParseComments)
	if err != nil {
		return fmt.Errorf("re-parsing pretty-print: %w", err)
	}

	// Parallel walk: apply TrailingComments from the side map.
	// Both trees have identical structure (the reparse is just a positioned copy),
	// so ast.Inspect visits nodes in the same order.  We skip comment nodes to
	// avoid mismatches from Doc fields.
	if len(TrailingComments) > 0 {
		// Helper: convert the tree-of-nodes into a slice-of-nodes
		collectNodes := func(node ast.Node) []ast.Node {
			var nodes []ast.Node
			ast.Inspect(node, func(n ast.Node) bool {
				// Filter out comments, as they only appear in the
				switch n.(type) {
				case *ast.CommentGroup, *ast.Comment:
					return false
				}
				nodes = append(nodes, n)
				return true
			})
			return nodes
		}

		origNodes := collectNodes(file)
		reparsedNodes := collectNodes(parsed)
		if len(origNodes) != len(reparsedNodes) {
			panic(fmt.Sprintf(
				"origNodes: %d; reparsedNodes: %d.  The AST generator is likely generating an invalid AST",
				len(origNodes), len(reparsedNodes),
			))
		}

		for i, orig := range origNodes {
			text, isOk := TrailingComments[orig]
			if !isOk {
				continue
			}
			reparsed := reparsedNodes[i]
			parsed.Comments = append(parsed.Comments, &ast.CommentGroup{
				List: []*ast.Comment{{Slash: reparsed.End(), Text: "// " + text}},
			})
		}
	}

	extractCommentMarker := func(stmt ast.Stmt) (string, bool) {
		expr, isOk := stmt.(*ast.ExprStmt)
		if !isOk {
			return "", false
		}
		call, isOk := expr.X.(*ast.CallExpr)
		if !isOk {
			return "", false
		}
		ident, isOk := call.Fun.(*ast.Ident)
		if !isOk || ident.Name != commentMarker {
			return "", false
		}
		lit, isOk := call.Args[0].(*ast.BasicLit)
		if !isOk {
			return "", false
		}
		return lit.Value[1 : len(lit.Value)-1], true
	}

	isBlankLineMarker := func(stmt ast.Stmt) bool {
		expr, isOk := stmt.(*ast.ExprStmt)
		if !isOk {
			return false
		}
		call, isOk := expr.X.(*ast.CallExpr)
		if !isOk {
			return false
		}
		ident, isOk := call.Fun.(*ast.Ident)
		return isOk && ident.Name == blankLineMarker
	}

	// Convert comment and blank-line markers
	ast.Inspect(parsed, func(n ast.Node) bool {
		// We only care about Block nodes
		block, isOk := n.(*ast.BlockStmt)
		if !isOk {
			return true
		}

		// Check the statements in this block , replacing our artificial Comment(...) and BlankLine(...) nodes
		// in the block statement's body with actual ones
		filtered := block.List[:0]
		for _, stmt := range block.List {
			if text, isOk := extractCommentMarker(stmt); isOk {
				// If it's a comment, add it to the fileset's list of Comments
				parsed.Comments = append(parsed.Comments, &ast.CommentGroup{
					List: []*ast.Comment{{Slash: stmt.Pos(), Text: "// " + text}},
				})
			} else if isBlankLineMarker(stmt) {
				// If it's a blank line, just remove it; the position gap creates a blank line
			} else {
				// Otherwise: it's a normal statement, so keep it.
				filtered = append(filtered, stmt)
			}
		}
		block.List = filtered
		return true
	})

	sort.Slice(parsed.Comments, func(i, j int) bool {
		return parsed.Comments[i].Pos() < parsed.Comments[j].Pos()
	})

	// Clear side map for next invocation
	for k := range TrailingComments {
		delete(TrailingComments, k)
	}

	err = printer.Fprint(w, fset, parsed)
	if err != nil {
		return fmt.Errorf("re-pretty-printing: %w", err)
	}
	return nil
}