From 88e115ff19272359c7e381182e18dbe44edb20e7 Mon Sep 17 00:00:00 2001
From: Vyacheslav Pryimak <vyacheslavpryimak@gmail.com>
Date: Tue, 4 May 2021 15:43:16 +0300
Subject: [PATCH] formatting for struct (#52)

---
 reviser/reviser.go      | 38 +++++++++++++++++++++++---------------
 reviser/reviser_test.go | 14 ++++++++++++++
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/reviser/reviser.go b/reviser/reviser.go
index 3ad5463..00b1502 100644
--- a/reviser/reviser.go
+++ b/reviser/reviser.go
@@ -122,25 +122,33 @@ func formatDecls(f *ast.File, options Options) {
 	}
 
 	for _, decl := range f.Decls {
-		if dd, ok := decl.(*ast.FuncDecl); ok {
-			var formattedComments []*ast.Comment
-			if dd.Doc != nil {
-				formattedComments = make([]*ast.Comment, len(dd.Doc.List))
-			}
+		switch dd := decl.(type) {
+		case *ast.GenDecl:
+			dd.Doc = fixCommentGroup(dd.Doc)
+		case *ast.FuncDecl:
+			dd.Doc = fixCommentGroup(dd.Doc)
+		}
+	}
+}
 
-			formattedDoc := &ast.CommentGroup{
-				List: formattedComments,
-			}
+func fixCommentGroup(commentGroup *ast.CommentGroup) *ast.CommentGroup {
+	if commentGroup == nil {
+		formattedDoc := &ast.CommentGroup{
+			List: []*ast.Comment{},
+		}
 
-			if dd.Doc != nil {
-				for i, comment := range dd.Doc.List {
-					formattedDoc.List[i] = comment
-				}
-			}
+		return formattedDoc
+	}
 
-			dd.Doc = formattedDoc
-		}
+	formattedDoc := &ast.CommentGroup{
+		List: make([]*ast.Comment, len(commentGroup.List)),
 	}
+
+	for i, comment := range commentGroup.List {
+		formattedDoc.List[i] = comment
+	}
+
+	return formattedDoc
 }
 
 func groupImports(
diff --git a/reviser/reviser_test.go b/reviser/reviser_test.go
index bf752eb..86b9455 100644
--- a/reviser/reviser_test.go
+++ b/reviser/reviser_test.go
@@ -954,12 +954,26 @@ func TestExecute_WithFormat(t *testing.T) {
 				projectName: "github.com/incu6us/goimports-reviser",
 				filePath:    "./testdata/example.go",
 				fileContent: `package testdata
+type SomeStruct struct{}
+type SomeStruct1 struct{}
+// SomeStruct2 comments
+type SomeStruct2 struct{}
+func (s *SomeStruct2) test() {}
 func test(){}
 func test1(){}
 `,
 			},
 			want: `package testdata
 
+type SomeStruct struct{}
+
+type SomeStruct1 struct{}
+
+// SomeStruct2 comments
+type SomeStruct2 struct{}
+
+func (s *SomeStruct2) test() {}
+
 func test() {}
 
 func test1() {}