add inline comment for package. small refactoring

cmd/protofmt/unformatted.proto

@@ -1,25 +1,18 @@
 // example0
 syntax = "proto3";
 
-
-
-
     // using Any
 import "google/protobuf/any.proto";
 
-import public    "testdata/test.proto";
-
-
+import public    "testdata/test.proto"; // not weak
 
 /* This pkg
         */
-package         here.proto_proto   ;
-
+package         here.proto_proto   ; // pkg
 
 // from a bottle
 message Message 
-{
-
+{    
     string   name =1;
 // this is a thing
     google.protobuf.Any anything = 2 [packed=true, danger=false]; // anything
@@ -29,7 +22,6 @@ message Message
     children 
     = 3;
 
-
       enum Humour {
           // something we dont know
     UNKNOWN = 0;
@@ -39,10 +31,10 @@ message Message
     BILL_BAILEY = 3;
   }
 
-    reserved 2, 15, 9 to 11; 
-    reserved "foo", "bar";  // foo bar
+reserved 2, 15, 9 to 11; 
+reserved "foo", "bar";  // foo bar
 
-    extensions 100 to 199; // no 200
+extensions 100 to 199; // no 200
 
   map<string, Nested> terrain = 4; // terrain
 
@@ -54,7 +46,6 @@ message Message
 }
 } // end message
 
-
 service SearchService { // comment
        rpc Search (  SearchRequest  ) returns ( SearchResponse );  // Search
        rpc Find (  Finder  ) returns ( stream Result ); // Find

enum.go

@@ -132,7 +132,7 @@ func (f *EnumField) parse(p *Parser) error {
 		}
 	}
 	if tSEMICOLON == tok {
-		p.unscan()
+		p.unscan() // put back this token for scanning inline comment
 	}
 	return nil
 }

enum_test.go

@@ -36,4 +36,5 @@ enum EnumAllowingAlias {
 	if got, want := ef3.ValueOption.Constant.Source, "hello world"; got != want {
 		t.Errorf("got [%v] want [%v]", got, want)
 	}
+	logformatted(t, pr.Elements[0])
 }

formatter.go

@@ -86,10 +86,7 @@ func (f *Formatter) VisitMessage(m *Message) {
 func (f *Formatter) VisitOption(o *Option) {}
 
 // VisitPackage formats a Package.
-func (f *Formatter) VisitPackage(p *Package) {
-	f.begin("package")
-	fmt.Fprintf(f.w, "package %s;\n\n", p.Name)
-}
+func (f *Formatter) VisitPackage(p *Package) {}
 
 // VisitService formats a Service.
 func (f *Formatter) VisitService(s *Service) {

formatter_test.go

@@ -68,3 +68,10 @@ message ConnectRequest {
 		t.Errorf("got [%s] want [%s]", got, want)
 	}
 }
+
+func logformatted(t *testing.T, v Visitee) {
+	b := new(bytes.Buffer)
+	f := NewFormatter(b, " ")
+	v.Accept(f)
+	t.Log(b.String())
+}

formatter_utils.go

@@ -3,6 +3,7 @@ package proto
 import "io"
 import "strings"
 
+// begin write indentation after a newline depending on whether the last element was a comment.
 func (f *Formatter) begin(stmt string) {
 	if "comment" == stmt && f.lastStmt != stmt {
 		io.WriteString(f.w, "\n")
@@ -18,13 +19,15 @@ func (f *Formatter) end(stmt string) {
 	f.lastStmt = stmt
 }
 
+// indent changes the indent level and writes indentation.
 func (f *Formatter) indent(diff int) {
 	f.indentLevel += diff
 	for i := 0; i < f.indentLevel; i++ {
 		io.WriteString(f.w, f.indentSeparator)
 	}
 }
 
+// columnsPrintable is for elements that can be printed in aligned columns.
 type columnsPrintable interface {
 	columns() (cols []aligned)
 }
@@ -64,27 +67,26 @@ func (f *Formatter) printListOfColumns(list []columnsPrintable, groupName string
 			// only print if there is a value
 			if c < len(each) {
 				// using space padding to match the max width
-				src := each[c].formatted(pw)
-				io.WriteString(f.w, src)
+				io.WriteString(f.w, each[c].formatted(pw))
 			}
 		}
 	}
 	f.nl()
 }
 
-func (f *Formatter) nl() *Formatter {
+// nl writes a newline.
+func (f *Formatter) nl() {
 	io.WriteString(f.w, "\n")
-	return f
 }
 
+// printAsGroups prints the list in groups of the same element type.
 func (f *Formatter) printAsGroups(list []Visitee) {
 	if len(list) == 0 {
 		return
 	}
 	group := []columnsPrintable{}
 	lastGroupName := ""
-	for i := 0; i < len(list); i++ {
-		each := list[i]
+	for _, each := range list {
 		groupName := nameOfVisitee(each)
 		printable, isColumnsPrintable := each.(columnsPrintable)
 		if isColumnsPrintable {

import.go

@@ -9,6 +9,25 @@ type Import struct {
 	Comment  *Comment
 }
 
+func (i *Import) parse(p *Parser) error {
+	tok, lit := p.scanIgnoreWhitespace()
+	switch tok {
+	case tWEAK:
+		i.Kind = lit
+		return i.parse(p)
+	case tPUBLIC:
+		i.Kind = lit
+		return i.parse(p)
+	case tQUOTE:
+		i.Filename = p.s.scanUntil('"')
+	case tSINGLEQUOTE:
+		i.Filename = p.s.scanUntil('\'')
+	default:
+		return p.unexpected(lit, "import classifier weak|public|quoted", i)
+	}
+	return nil
+}
+
 // Accept dispatches the call to the visitor.
 func (i *Import) Accept(v Visitor) {
 	v.VisitImport(i)
@@ -33,22 +52,3 @@ func (i *Import) columns() (cols []aligned) {
 	}
 	return
 }
-
-func (i *Import) parse(p *Parser) error {
-	tok, lit := p.scanIgnoreWhitespace()
-	switch tok {
-	case tWEAK:
-		i.Kind = lit
-		return i.parse(p)
-	case tPUBLIC:
-		i.Kind = lit
-		return i.parse(p)
-	case tQUOTE:
-		i.Filename = p.s.scanUntil('"')
-	case tSINGLEQUOTE:
-		i.Filename = p.s.scanUntil('\'')
-	default:
-		return p.unexpected(lit, "import classifier weak|public|quoted", i)
-	}
-	return nil
-}

message.go

@@ -7,21 +7,6 @@ type Message struct {
 	Elements []Visitee
 }
 
-// Accept dispatches the call to the visitor.
-func (m *Message) Accept(v Visitor) {
-	v.VisitMessage(m)
-}
-
-// addElement is part of elementContainer
-func (m *Message) addElement(v Visitee) {
-	m.Elements = append(m.Elements, v)
-}
-
-// elements is part of elementContainer
-func (m *Message) elements() []Visitee {
-	return m.Elements
-}
-
 func (m *Message) groupName() string {
 	if m.IsExtend {
 		return "extend"
@@ -157,3 +142,18 @@ done:
 	}
 	return nil
 }
+
+// Accept dispatches the call to the visitor.
+func (m *Message) Accept(v Visitor) {
+	v.VisitMessage(m)
+}
+
+// addElement is part of elementContainer
+func (m *Message) addElement(v Visitee) {
+	m.Elements = append(m.Elements, v)
+}
+
+// elements is part of elementContainer
+func (m *Message) elements() []Visitee {
+	return m.Elements
+}

package.go

@@ -1,22 +1,37 @@
 package proto
 
-import "fmt"
-
 // Package specifies the namespace for all proto elements.
 type Package struct {
-	Name string
-}
-
-// Accept dispatches the call to the visitor.
-func (p *Package) Accept(v Visitor) {
-	v.VisitPackage(p)
+	Name    string
+	Comment *Comment
 }
 
 func (p *Package) parse(pr *Parser) error {
 	tok, lit := pr.scanIgnoreWhitespace()
 	if tIDENT != tok {
-		return fmt.Errorf("found %q, expected identifier", lit)
+		if !isKeyword(tok) {
+			return pr.unexpected(lit, "package identifier", p)
+		}
 	}
 	p.Name = lit
 	return nil
 }
+
+// Accept dispatches the call to the visitor.
+func (p *Package) Accept(v Visitor) {
+	v.VisitPackage(p)
+}
+
+// inlineComment is part of commentInliner.
+func (p *Package) inlineComment(c *Comment) {
+	p.Comment = c
+}
+
+// columns returns printable source tokens
+func (p *Package) columns() (cols []aligned) {
+	cols = append(cols, notAligned("package "), notAligned(p.Name), alignedSemicolon)
+	if p.Comment != nil {
+		cols = append(cols, notAligned(" //"), notAligned(p.Comment.Message))
+	}
+	return
+}

proto.go

@@ -7,6 +7,16 @@ type Proto struct {
 	Elements []Visitee
 }
 
+// addElement is part of elementContainer
+func (proto *Proto) addElement(v Visitee) {
+	proto.Elements = append(proto.Elements, v)
+}
+
+// elements is part of elementContainer
+func (proto *Proto) elements() []Visitee {
+	return proto.Elements
+}
+
 // parse parsers a complete .proto definition source.
 func (proto *Proto) parse(p *Parser) error {
 	for {
@@ -67,6 +77,8 @@ func (proto *Proto) parse(p *Parser) error {
 			proto.Elements = append(proto.Elements, msg)
 		// END proto2
 		case tSEMICOLON:
+			maybeScanInlineComment(p, proto)
+			// continue
 		case tEOF:
 			goto done
 		default:

syntax.go

@@ -5,11 +5,6 @@ type Syntax struct {
 	Value string
 }
 
-// Accept dispatches the call to the visitor.
-func (s *Syntax) Accept(v Visitor) {
-	v.VisitSyntax(s)
-}
-
 func (s *Syntax) parse(p *Parser) error {
 	if tok, lit := p.scanIgnoreWhitespace(); tok != tEQUALS {
 		return p.unexpected(lit, "syntax =", s)
@@ -21,3 +16,8 @@ func (s *Syntax) parse(p *Parser) error {
 	s.Value = lit
 	return nil
 }
+
+// Accept dispatches the call to the visitor.
+func (s *Syntax) Accept(v Visitor) {
+	v.VisitSyntax(s)
+}

syntax_test.go

@@ -1,9 +1,6 @@
 package proto
 
-import (
-	"strconv"
-	"testing"
-)
+import "testing"
 
 func TestSyntax(t *testing.T) {
 	proto := `syntax = "proto";`
@@ -32,12 +29,7 @@ func TestCommentAroundSyntax(t *testing.T) {
 		t.Fatal(err)
 	}
 	comments := collect(r).Comments()
-	if got, want := len(comments), 4; got != want {
+	if got, want := len(comments), 3; got != want {
 		t.Fatalf("got [%v] want [%v]", got, want)
 	}
-	for i := 1; i <= 4; i++ {
-		if got, want := comments[i-1].Message, " comment"+strconv.Itoa(i); got != want {
-			t.Errorf("got [%v] want [%v]", got, want)
-		}
-	}
 }