perf: Reduce allocations and syscalls in formatting pipeline

6 files changed, 25 insertions, 20 deletions

diff --git a/adapter_go.go b/adapter_go.go
index 59096bb..1b39e0b 100644
--- a/adapter_go.go
+++ b/adapter_go.go
@@ -1,11 +1,11 @@
 package main
 
 import (
+	"bytes"
 	"github.com/Fuwn/iku/engine"
 	"go/format"
 	"go/parser"
 	"go/token"
-	"strings"
 )
 
 type GoAdapter struct{}
@@ -26,11 +26,12 @@ func (a *GoAdapter) Analyze(source []byte) ([]byte, []engine.LineEvent, error) {
 
 	formatter := &Formatter{}
 	lineInformationMap := formatter.buildLineInfo(tokenFileSet, parsedFile)
-	sourceLines := strings.Split(string(formattedSource), "\n")
-	events := make([]engine.LineEvent, len(sourceLines))
+	sourceByteLines := bytes.Split(formattedSource, []byte("\n"))
+	events := make([]engine.LineEvent, len(sourceByteLines))
 	insideRawString := false
 
-	for lineIndex, currentLine := range sourceLines {
+	for lineIndex, currentLineBytes := range sourceByteLines {
+		currentLine := string(currentLineBytes)
 		backtickCount := countRawStringDelimiters(currentLine)
 		wasInsideRawString := insideRawString
 
diff --git a/engine/engine.go b/engine/engine.go
index 7a5399c..15ee568 100644
--- a/engine/engine.go
+++ b/engine/engine.go
@@ -104,13 +104,16 @@ func (e *Engine) Format(events []LineEvent) []string {
 
 func (e *Engine) FormatToString(events []LineEvent) string {
 	lines := e.Format(events)
-	output := strings.Join(lines, "\n")
+	lines = append(lines, "")
 
-	if !strings.HasSuffix(output, "\n") {
-		output += "\n"
-	}
+	return strings.Join(lines, "\n")
+}
+
+func (e *Engine) FormatToBytes(events []LineEvent) []byte {
+	lines := e.Format(events)
+	lines = append(lines, "")
 
-	return output
+	return []byte(strings.Join(lines, "\n"))
 }
 
 func (e *Engine) findNextNonComment(events []LineEvent, startIndex int) int {
diff --git a/formatter.go b/formatter.go
index 3d4e873..0256694 100644
--- a/formatter.go
+++ b/formatter.go
@@ -33,7 +33,7 @@ func (f *Formatter) Format(source []byte, filename string) ([]byte, error) {
 
 	formattingEngine := &engine.Engine{CommentMode: MapCommentMode(f.CommentMode)}
 
-	return []byte(formattingEngine.FormatToString(events)), nil
+	return formattingEngine.FormatToBytes(events), nil
 }
 
 func analyzeSource(source []byte, filename string) ([]byte, []engine.LineEvent, error) {
diff --git a/formatter_test.go b/formatter_test.go
index 6ce7edd..64b6f6d 100644
--- a/formatter_test.go
+++ b/formatter_test.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"strings"
 	"testing"
 )
@@ -448,8 +449,7 @@ func BenchmarkFormatLarge(b *testing.B) {
 	sourceBuilder.WriteString("package main\n\n")
 
 	for functionIndex := range 100 {
-		sourceBuilder.WriteString("func foo")
-		sourceBuilder.WriteString(string(rune('A' + functionIndex%26)))
+		fmt.Fprintf(&sourceBuilder, "func foo%d", functionIndex)
 		sourceBuilder.WriteString("() {\n")
 		sourceBuilder.WriteString("\tx := 1\n")
 		sourceBuilder.WriteString("\tif x > 0 {\n")
diff --git a/inspect.go b/inspect.go
index a2e39c9..0c05657 100644
--- a/inspect.go
+++ b/inspect.go
@@ -1,9 +1,9 @@
 package main
 
 import (
+	"fmt"
 	"go/ast"
 	"go/token"
-	"reflect"
 )
 
 func isGeneralDeclarationScoped(generalDeclaration *ast.GenDecl) bool {
@@ -20,7 +20,7 @@ func isGeneralDeclarationScoped(generalDeclaration *ast.GenDecl) bool {
 }
 
 func (f *Formatter) buildLineInfo(tokenFileSet *token.FileSet, parsedFile *ast.File) map[int]*lineInformation {
-	lineInformationMap := make(map[int]*lineInformation)
+	lineInformationMap := make(map[int]*lineInformation, 2*len(parsedFile.Decls))
 	tokenFile := tokenFileSet.File(parsedFile.Pos())
 
 	if tokenFile == nil {
@@ -41,7 +41,7 @@ func (f *Formatter) buildLineInfo(tokenFileSet *token.FileSet, parsedFile *ast.F
 			statementType = "func"
 			isScoped = true
 		default:
-			statementType = reflect.TypeOf(declaration).String()
+			statementType = fmt.Sprintf("%T", declaration)
 		}
 
 		lineInformationMap[startLine] = &lineInformation{statementType: statementType, isTopLevel: true, isScoped: isScoped, isStartLine: true}
@@ -88,14 +88,14 @@ func (f *Formatter) processStatementList(tokenFile *token.File, statements []ast
 			if generalDeclaration, isGeneralDeclaration := typedStatement.Decl.(*ast.GenDecl); isGeneralDeclaration {
 				statementType = generalDeclaration.Tok.String()
 			} else {
-				statementType = reflect.TypeOf(statement).String()
+				statementType = fmt.Sprintf("%T", statement)
 			}
 		case *ast.IfStmt, *ast.ForStmt, *ast.RangeStmt, *ast.SwitchStmt,
 			*ast.TypeSwitchStmt, *ast.SelectStmt, *ast.BlockStmt:
-			statementType = reflect.TypeOf(statement).String()
+			statementType = fmt.Sprintf("%T", statement)
 			isScoped = true
 		default:
-			statementType = reflect.TypeOf(statement).String()
+			statementType = fmt.Sprintf("%T", statement)
 		}
 
 		existingStart := lineInformationMap[startLine]
diff --git a/main.go b/main.go
index 649752b..c41bfe2 100644
--- a/main.go
+++ b/main.go
@@ -5,6 +5,7 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"io/fs"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -107,12 +108,12 @@ var supportedFileExtensions = map[string]bool{
 func processDirectory(formatter *Formatter, directoryPath string, exitCode *int) error {
 	var sourceFilePaths []string
 
-	err := filepath.Walk(directoryPath, func(currentPath string, fileInfo os.FileInfo, err error) error {
+	err := filepath.WalkDir(directoryPath, func(currentPath string, dirEntry fs.DirEntry, err error) error {
 		if err != nil {
 			return err
 		}
 
-		if !fileInfo.IsDir() && supportedFileExtensions[filepath.Ext(currentPath)] {
+		if !dirEntry.IsDir() && supportedFileExtensions[filepath.Ext(currentPath)] {
 			sourceFilePaths = append(sourceFilePaths, currentPath)
 		}