From 6bbfd7b94caa227dc6e4d7c72a967f7a92e04c44 Mon Sep 17 00:00:00 2001
From: Hau Nguyen <hauvipapro@gmail.com>
Date: Tue, 17 Jan 2023 17:20:05 +0700
Subject: [PATCH] fix: update 2 times with dstFile

---
 internal/imports/formatter.go | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/internal/imports/formatter.go b/internal/imports/formatter.go
index bb442b9..ec2d229 100644
--- a/internal/imports/formatter.go
+++ b/internal/imports/formatter.go
@@ -234,7 +234,7 @@ func (ft *Formatter) formatImports(
 	if err != nil {
 		return nil, fmt.Errorf("decorator: failed to parse file [%s]: %w", path, err)
 	}
-	if len(dstFile.Imports) == 0 {
+	if len(dstFile.Imports) == 0 || len(dstFile.Decls) == 0 {
 		return nil, ErrEmptyImport
 	}
 	ft.logDSTImportSpecs("formatImports: dstImportSpecs", dstFile.Imports)
@@ -250,8 +250,28 @@ func (ft *Formatter) formatImports(
 	}
 	ft.logDSTImportSpecs("formatImports: formattedDSTImportSpecs: ", formattedDSTImportSpecs)
 
+	// First update
 	dstFile.Imports = formattedDSTImportSpecs
 
+	genSpecs := dstFile.Decls[0].(*dst.GenDecl).Specs
+	formattedGenSpecs := make([]dst.Spec, 0, len(genSpecs))
+
+	// Append all imports first
+	for _, importSpec := range formattedDSTImportSpecs {
+		formattedGenSpecs = append(formattedGenSpecs, importSpec)
+	}
+
+	// Append all non imports later
+	for _, genSpec := range genSpecs {
+		if _, ok := genSpec.(*dst.ImportSpec); !ok {
+			formattedGenSpecs = append(formattedGenSpecs, genSpec)
+			continue
+		}
+	}
+
+	// Second update
+	dstFile.Decls[0].(*dst.GenDecl).Specs = formattedGenSpecs
+
 	var buf bytes.Buffer
 	if err := decorator.Fprint(&buf, dstFile); err != nil {
 		return nil, fmt.Errorf("decorator: failed to fprint [%s]: %w", path, err)