path: root/formatter.go

package main

import (
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"reflect"
	"regexp"
	"strings"
)

type CommentMode int

const (
	CommentsFollow CommentMode = iota
	CommentsPrecede
	CommentsStandalone
)

type Formatter struct {
	CommentMode CommentMode
}

type lineInfo struct {
	statementType string
	isTopLevel    bool
	isScoped      bool
	isStartLine   bool
}

func (f *Formatter) Format(source []byte) ([]byte, error) {
	formatted, err := format.Source(source)

	if err != nil {
		return nil, err
	}

	fileSet := token.NewFileSet()
	file, err := parser.ParseFile(fileSet, "", formatted, parser.ParseComments)

	if err != nil {
		return nil, err
	}

	lineInfoMap := f.buildLineInfo(fileSet, file)

	return f.rewrite(formatted, lineInfoMap), nil
}

func (f *Formatter) buildLineInfo(fileSet *token.FileSet, file *ast.File) map[int]*lineInfo {
	lineInfoMap := make(map[int]*lineInfo)
	tokenFile := fileSet.File(file.Pos())

	if tokenFile == nil {
		return lineInfoMap
	}

	for _, declaration := range file.Decls {
		startLine := tokenFile.Line(declaration.Pos())
		endLine := tokenFile.Line(declaration.End())
		typeName := reflect.TypeOf(declaration).String()
		lineInfoMap[startLine] = &lineInfo{statementType: typeName, isTopLevel: true, isStartLine: true}
		lineInfoMap[endLine] = &lineInfo{statementType: typeName, isTopLevel: true, isStartLine: false}
	}

	ast.Inspect(file, func(node ast.Node) bool {
		if node == nil {
			return true
		}

		switch typedNode := node.(type) {
		case *ast.BlockStmt:
			f.processBlock(tokenFile, typedNode, lineInfoMap)
		case *ast.CaseClause:
			f.processStatementList(tokenFile, typedNode.Body, lineInfoMap)
		case *ast.CommClause:
			f.processStatementList(tokenFile, typedNode.Body, lineInfoMap)
		}

		return true
	})

	return lineInfoMap
}

func (f *Formatter) processBlock(tokenFile *token.File, block *ast.BlockStmt, lineInfoMap map[int]*lineInfo) {
	if block == nil {
		return
	}

	f.processStatementList(tokenFile, block.List, lineInfoMap)
}

func (f *Formatter) processStatementList(tokenFile *token.File, statements []ast.Stmt, lineInfoMap map[int]*lineInfo) {
	for _, statement := range statements {
		startLine := tokenFile.Line(statement.Pos())
		endLine := tokenFile.Line(statement.End())
		typeName := reflect.TypeOf(statement).String()
		isScoped := false

		switch statement.(type) {
		case *ast.IfStmt, *ast.ForStmt, *ast.RangeStmt, *ast.SwitchStmt,
			*ast.TypeSwitchStmt, *ast.SelectStmt, *ast.BlockStmt:

			isScoped = true
		}

		existingStart := lineInfoMap[startLine]

		if existingStart == nil || !existingStart.isStartLine {
			lineInfoMap[startLine] = &lineInfo{statementType: typeName, isTopLevel: false, isScoped: isScoped, isStartLine: true}
		}

		existingEnd := lineInfoMap[endLine]

		if existingEnd == nil || !existingEnd.isStartLine {
			lineInfoMap[endLine] = &lineInfo{statementType: typeName, isTopLevel: false, isScoped: isScoped, isStartLine: false}
		}

		switch typedStatement := statement.(type) {
		case *ast.IfStmt:
			f.processBlock(tokenFile, typedStatement.Body, lineInfoMap)

			if typedStatement.Else != nil {
				if elseBlock, isBlockStatement := typedStatement.Else.(*ast.BlockStmt); isBlockStatement {
					f.processBlock(tokenFile, elseBlock, lineInfoMap)
				} else if elseIfStatement, isIfStatement := typedStatement.Else.(*ast.IfStmt); isIfStatement {
					f.processIfStatement(tokenFile, elseIfStatement, lineInfoMap)
				}
			}
		case *ast.ForStmt:
			f.processBlock(tokenFile, typedStatement.Body, lineInfoMap)
		case *ast.RangeStmt:
			f.processBlock(tokenFile, typedStatement.Body, lineInfoMap)
		case *ast.SwitchStmt:
			f.processBlock(tokenFile, typedStatement.Body, lineInfoMap)
		case *ast.TypeSwitchStmt:
			f.processBlock(tokenFile, typedStatement.Body, lineInfoMap)
		case *ast.SelectStmt:
			f.processBlock(tokenFile, typedStatement.Body, lineInfoMap)
		case *ast.BlockStmt:
			f.processBlock(tokenFile, typedStatement, lineInfoMap)
		}
	}
}

func (f *Formatter) processIfStatement(tokenFile *token.File, ifStatement *ast.IfStmt, lineInfoMap map[int]*lineInfo) {
	startLine := tokenFile.Line(ifStatement.Pos())
	endLine := tokenFile.Line(ifStatement.End())
	existingStart := lineInfoMap[startLine]

	if existingStart == nil || !existingStart.isStartLine {
		lineInfoMap[startLine] = &lineInfo{statementType: "*ast.IfStmt", isTopLevel: false, isScoped: true, isStartLine: true}
	}

	existingEnd := lineInfoMap[endLine]

	if existingEnd == nil || !existingEnd.isStartLine {
		lineInfoMap[endLine] = &lineInfo{statementType: "*ast.IfStmt", isTopLevel: false, isScoped: true, isStartLine: false}
	}

	f.processBlock(tokenFile, ifStatement.Body, lineInfoMap)

	if ifStatement.Else != nil {
		if elseBlock, isBlockStatement := ifStatement.Else.(*ast.BlockStmt); isBlockStatement {
			f.processBlock(tokenFile, elseBlock, lineInfoMap)
		} else if elseIfStatement, isIfStatement := ifStatement.Else.(*ast.IfStmt); isIfStatement {
			f.processIfStatement(tokenFile, elseIfStatement, lineInfoMap)
		}
	}
}

func (f *Formatter) rewrite(source []byte, lineInfoMap map[int]*lineInfo) []byte {
	lines := strings.Split(string(source), "\n")

	var result []string

	closingBracePattern := regexp.MustCompile(`^\s*[\}\)]`)
	openingBracePattern := regexp.MustCompile(`[\{\(]\s*$`)
	caseLabelPattern := regexp.MustCompile(`^\s*(case\s+.*|default\s*):\s*$`)
	commentOnlyPattern := regexp.MustCompile(`^\s*//`)
	packageLinePattern := regexp.MustCompile(`^package\s+`)
	previousWasOpenBrace := false
	previousType := ""
	previousWasComment := false
	previousWasTopLevel := false
	insideRawString := false

	for index, line := range lines {
		backtickCount := strings.Count(line, "`")
		wasInsideRawString := insideRawString

		if backtickCount%2 == 1 {
			insideRawString = !insideRawString
		}

		if wasInsideRawString {
			result = append(result, line)
			continue
		}

		lineNumber := index + 1
		trimmed := strings.TrimSpace(line)
		isBlank := trimmed == ""
		isClosingBrace := closingBracePattern.MatchString(line)
		isOpeningBrace := openingBracePattern.MatchString(line)
		isCaseLabel := caseLabelPattern.MatchString(line)
		isCommentOnlyLine := commentOnlyPattern.MatchString(line)
		isPackageLine := packageLinePattern.MatchString(trimmed)

		if isBlank {
			continue
		}

		info := lineInfoMap[lineNumber]
		currentType := ""

		if info != nil {
			currentType = info.statementType
		}

		if isPackageLine {
			currentType = "package"
		}

		needsBlank := false
		currentIsTopLevel := info != nil && info.isTopLevel
		currentIsScoped := info != nil && info.isScoped

		if len(result) > 0 && !previousWasOpenBrace && !isClosingBrace && !isCaseLabel {
			if currentIsTopLevel && previousWasTopLevel {
				if f.CommentMode == CommentsFollow && previousWasComment {
				} else {
					needsBlank = true
				}
			} else if currentIsScoped {
				if f.CommentMode == CommentsFollow && previousWasComment {
				} else {
					needsBlank = true
				}
			} else if currentType != "" && previousType != "" && currentType != previousType {
				if f.CommentMode == CommentsFollow && previousWasComment {
				} else {
					needsBlank = true
				}
			}

			if f.CommentMode == CommentsFollow && isCommentOnlyLine && !previousWasComment {
				nextLineNumber := f.findNextNonCommentLine(lines, index+1)

				if nextLineNumber > 0 {
					nextInfo := lineInfoMap[nextLineNumber]

					if nextInfo != nil {
						nextIsTopLevel := nextInfo.isTopLevel
						nextIsScoped := nextInfo.isScoped

						if nextIsTopLevel && previousWasTopLevel {
							needsBlank = true
						} else if nextIsScoped {
							needsBlank = true
						} else if nextInfo.statementType != "" && previousType != "" && nextInfo.statementType != previousType {
							needsBlank = true
						}
					}
				}
			}
		}

		if needsBlank {
			result = append(result, "")
		}

		result = append(result, line)
		previousWasOpenBrace = isOpeningBrace || isCaseLabel
		previousWasComment = isCommentOnlyLine

		if info != nil {
			previousType = info.statementType
			previousWasTopLevel = info.isTopLevel
		} else if currentType != "" {
			previousType = currentType
			previousWasTopLevel = false
		}
	}

	output := strings.Join(result, "\n")

	if !strings.HasSuffix(output, "\n") {
		output += "\n"
	}

	return []byte(output)
}

func (f *Formatter) findNextNonCommentLine(lines []string, startIndex int) int {
	commentOnlyPattern := regexp.MustCompile(`^\s*//`)

	for index := startIndex; index < len(lines); index++ {
		trimmed := strings.TrimSpace(lines[index])

		if trimmed == "" {
			continue
		}

		if commentOnlyPattern.MatchString(lines[index]) {
			continue
		}

		return index + 1
	}

	return 0
}