From fdd6a2b909336861bc8c834b766257fdfe99a967 Mon Sep 17 00:00:00 2001
From: Gavin Shriver <g@stately.cloud>
Date: Mon, 16 Dec 2024 12:12:06 -0800
Subject: [PATCH] blah

---
 gencommon/comments.go            | 24 +++++++++++++++++++++++-
 gencommon/interface_test.go      |  4 ++++
 gencommon/internal/test_types.go |  4 ++++
 gogenproto/internal/test.pb.go   |  2 +-
 4 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/gencommon/comments.go b/gencommon/comments.go
index 27093f7..a2dd87f 100644
--- a/gencommon/comments.go
+++ b/gencommon/comments.go
@@ -54,9 +54,31 @@ func CommentsFromObj(pkg *packages.Package, typeName string) Comments {
 		}
 
 		if decl, ok := obj.Decl.(*ast.TypeSpec); ok {
-			return FromCommentGroup(decl.Doc)
+			// structs/interfaces won't have comments built-in unless
+			// they're defined a particular way... e.g.
+			// type {
+			//    // comment on MyType
+			//    MyType struct {
+			//    ...
+			//    }
+			// }
+			// So we'll check that first, but then build a comment map and try to get the
+			// correct comment that way.
+			if decl.Doc != nil && len(decl.Doc.List) > 0 {
+				return FromCommentGroup(decl.Doc)
+			} else if decl.Comment != nil && len(decl.Comment.List) > 0 {
+				return FromCommentGroup(decl.Comment)
+			}
+			cmap := ast.NewCommentMap(pkg.Fset, stax, stax.Comments)
+			for tp, commentCommentGroups := range cmap {
+				if v, ok := tp.(*ast.GenDecl); ok && len(v.Specs) == 1 && v.Specs[0] == decl && len(commentCommentGroups) > 0 {
+					// take the LAST comment.... theoretically the closest comment to the actual struct we're looking at.
+					return FromCommentGroup(commentCommentGroups[len(commentCommentGroups)-1])
+				}
+			}
 		}
 	}
+
 	return nil
 }
 
diff --git a/gencommon/interface_test.go b/gencommon/interface_test.go
index 7de1dac..cc7c277 100644
--- a/gencommon/interface_test.go
+++ b/gencommon/interface_test.go
@@ -67,6 +67,10 @@ func TestFindInterface(t *testing.T) {
 			iface, err := gencommon.FindInterface(imports, pkgs, pkg.PkgPath, "TypeToGenerate", test.options)
 			require.NoError(t, err)
 
+			assert.Equal(t,
+				"// TypeToGenerate has a comment.\n// SecondLine of expected comment.",
+				iface.Comments.String())
+
 			expected := set.Make(test.expectedMethods...)
 			for _, m := range iface.Methods {
 				assert.Truef(t, expected.Remove(m.Name), "found duplicate OR unexpected method %q", m.Name)
diff --git a/gencommon/internal/test_types.go b/gencommon/internal/test_types.go
index b66b6ee..655483e 100644
--- a/gencommon/internal/test_types.go
+++ b/gencommon/internal/test_types.go
@@ -5,6 +5,10 @@ import (
 	"github.com/drshriveer/gtools/gencommon/internal/nestedpkg"
 )
 
+// Comment to throw us off the trail.
+
+// TypeToGenerate has a comment.
+// SecondLine of expected comment.
 type TypeToGenerate struct {
 	EmbeddedA
 	EmbeddedB
diff --git a/gogenproto/internal/test.pb.go b/gogenproto/internal/test.pb.go
index aa9475c..439bdd3 100644
--- a/gogenproto/internal/test.pb.go
+++ b/gogenproto/internal/test.pb.go
@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.31.0
-// 	protoc        v4.25.2
+// 	protoc        v4.25.1
 // source: test.proto
 
 package internal