From 45eede45b986cf49c08dc394d22baf438ba16d4d Mon Sep 17 00:00:00 2001
From: Jeff Thompson <jeff@thefirst.org>
Date: Wed, 11 Oct 2023 11:22:05 +0200
Subject: [PATCH] chore: In amino/pkg, add optional WithComments, use them in
 GenerateProto3MessagePartial

Signed-off-by: Jeff Thompson <jeff@thefirst.org>
---
 tm2/pkg/amino/genproto/comments_test.go | 61 ++++++++++++++++++++++++
 tm2/pkg/amino/genproto/genproto.go      | 10 ++++
 tm2/pkg/amino/pkg/pkg.go                | 62 ++++++++++++++++++++++++-
 3 files changed, 131 insertions(+), 2 deletions(-)
 create mode 100644 tm2/pkg/amino/genproto/comments_test.go

diff --git a/tm2/pkg/amino/genproto/comments_test.go b/tm2/pkg/amino/genproto/comments_test.go
new file mode 100644
index 00000000000..2294dfe42de
--- /dev/null
+++ b/tm2/pkg/amino/genproto/comments_test.go
@@ -0,0 +1,61 @@
+package genproto
+
+import (
+	"path"
+	"testing"
+
+	"github.com/gnolang/gno/tm2/pkg/amino"
+	"github.com/jaekwon/testify/assert"
+)
+
+// message comment
+type TestMessageName struct {
+	// field comment 1
+	FieldName1 string
+	// field comment 2
+	FieldName2 []uint64
+}
+
+// message comment 2
+type TestMessageName2 struct {
+	// another field comment
+	FieldName string
+}
+
+func TestComments(t *testing.T) {
+	pkg := amino.RegisterPackage(
+		amino.NewPackage(
+			"github.com/gnolang/gno/tm2/pkg/amino/genproto",
+			"amino_test",
+			amino.GetCallersDirname(),
+		).WithTypes(
+			&TestMessageName{},
+			&TestMessageName2{},
+		// Add comments from this same source file.
+		).WithComments(path.Join(amino.GetCallersDirname(), "comments_test.go")))
+
+	p3c := NewP3Context()
+	p3c.RegisterPackage(pkg)
+	p3c.ValidateBasic()
+	p3doc := p3c.GenerateProto3SchemaForTypes(pkg, pkg.ReflectTypes()...)
+	proto3Schema := p3doc.Print()
+	assert.Equal(t, proto3Schema, `syntax = "proto3";
+package amino_test;
+
+option go_package = "github.com/gnolang/gno/tm2/pkg/amino/genproto/pb";
+
+// messages
+// message comment
+message TestMessageName {
+	// field comment 1
+	string FieldName1 = 1;
+	// field comment 2
+	repeated uint64 FieldName2 = 2;
+}
+
+// message comment 2
+message TestMessageName2 {
+	// another field comment
+	string FieldName = 1;
+}`)
+}
diff --git a/tm2/pkg/amino/genproto/genproto.go b/tm2/pkg/amino/genproto/genproto.go
index 8af73b2690a..a8dc18922ea 100644
--- a/tm2/pkg/amino/genproto/genproto.go
+++ b/tm2/pkg/amino/genproto/genproto.go
@@ -191,6 +191,13 @@ func (p3c *P3Context) GenerateProto3MessagePartial(p3doc *P3Doc, rt reflect.Type
 
 	p3msg.Name = info.Name // not rinfo.
 
+	var fieldComments map[string]string
+	if pkgType, ok := rinfo.Package.GetType(rt); ok {
+		p3msg.Comment = pkgType.Comment
+		// We will check for optional field comments below.
+		fieldComments = pkgType.FieldComments
+	}
+
 	// Append to p3msg.Fields, fields of the struct.
 	for _, field := range rsfields { // rinfo.
 		fp3, fp3IsRepeated, implicit := typeToP3Type(info.Package, field.TypeInfo, field.FieldOptions)
@@ -209,6 +216,9 @@ func (p3c *P3Context) GenerateProto3MessagePartial(p3doc *P3Doc, rt reflect.Type
 			Name:     field.Name,
 			Number:   field.FieldOptions.BinFieldNum,
 		}
+		if fieldComments != nil {
+			p3Field.Comment = fieldComments[field.Name]
+		}
 		p3msg.Fields = append(p3msg.Fields, p3Field)
 	}
 
diff --git a/tm2/pkg/amino/pkg/pkg.go b/tm2/pkg/amino/pkg/pkg.go
index fad46c4cad7..5ce0a0ef7ea 100644
--- a/tm2/pkg/amino/pkg/pkg.go
+++ b/tm2/pkg/amino/pkg/pkg.go
@@ -2,6 +2,9 @@ package pkg
 
 import (
 	"fmt"
+	"go/ast"
+	"go/parser"
+	"go/token"
 	"path/filepath"
 	"reflect"
 	"regexp"
@@ -11,8 +14,10 @@ import (
 
 type Type struct {
 	Type             reflect.Type
-	Name             string // proto3 name (override)
-	PointerPreferred bool   // whether pointer is preferred for decoding interface.
+	Name             string            // proto3 name (override)
+	PointerPreferred bool              // whether pointer is preferred for decoding interface.
+	Comment          string            // optional doc comment for the type
+	FieldComments    map[string]string // If not nil, the optional doc comment for each field name
 }
 
 func (t *Type) FullName(pkg *Package) string {
@@ -196,6 +201,59 @@ func (pkg *Package) WithP3SchemaFile(file string) *Package {
 	return pkg
 }
 
+// Parse the Go code in filename and scan the AST looking for struct doc comments.
+// Find the Type in pkg.Types and set its Comment and FieldComments, which are
+// used by genproto.GenerateProto3MessagePartial to set the Comment in the P3Doc
+// and related P3Field objects.
+func (pkg *Package) WithComments(filename string) *Package {
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
+	if err != nil {
+		panic(err)
+	}
+
+	ast.Inspect(f, func(node ast.Node) bool {
+		if genDecl, ok := node.(*ast.GenDecl); ok {
+			for _, spec := range genDecl.Specs {
+				if typeSpec, ok := spec.(*ast.TypeSpec); ok {
+					if pkgType := pkg.getTypeByName(typeSpec.Name.Name); pkgType != nil {
+						if genDecl.Doc != nil {
+							// Set the type comment.
+							pkgType.Comment = strings.TrimSpace(genDecl.Doc.Text())
+						}
+						if structType, ok := typeSpec.Type.(*ast.StructType); ok {
+							for _, field := range structType.Fields.List {
+								if field.Names != nil && len(field.Names) == 1 && field.Doc != nil {
+									// Set the field comment.
+									if pkgType.FieldComments == nil {
+										pkgType.FieldComments = make(map[string]string)
+									}
+
+									pkgType.FieldComments[field.Names[0].Name] = strings.TrimSpace(field.Doc.Text())
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+		return true
+	})
+
+	return pkg
+}
+
+// Get the Type by name. If not found, return nil.
+func (pkg *Package) getTypeByName(name string) *Type {
+	for _, t := range pkg.Types {
+		if t.Name == name {
+			return t
+		}
+	}
+
+	return nil
+}
+
 // Result cannot be modified.
 func (pkg *Package) GetType(rt reflect.Type) (t Type, ok bool) {
 	if rt.Kind() == reflect.Ptr {