Merge pull request emicklei#55 from emicklei/issue26

Issue26

comment.go

@@ -73,6 +73,13 @@ func (c *Comment) Merge(other *Comment) {
 	c.Cstyle = c.Cstyle || other.Cstyle
 }
 
+func (c Comment) hasTextOnLine(line int) bool {
+	if len(c.Lines) == 0 {
+		return false
+	}
+	return c.Position.Line <= line && line <= c.Position.Line+len(c.Lines)-1
+}
+
 // Message returns the first line or empty if no lines.
 func (c Comment) Message() string {
 	if len(c.Lines) == 0 {
@@ -105,12 +112,12 @@ func maybeScanInlineComment(p *Parser, c elementContainer) {
 	}
 }
 
-// takeLastComment removes and returns the last element of the list if it is a Comment.
-func takeLastComment(list []Visitee) (*Comment, []Visitee) {
+// takeLastCommentIfEndsOnLine removes and returns the last element of the list if it is a Comment
+func takeLastCommentIfEndsOnLine(list []Visitee, line int) (*Comment, []Visitee) {
 	if len(list) == 0 {
 		return nil, list
 	}
-	if last, ok := list[len(list)-1].(*Comment); ok {
+	if last, ok := list[len(list)-1].(*Comment); ok && last.hasTextOnLine(line) {
 		return last, list[:len(list)-1]
 	}
 	return nil, list
@@ -119,17 +126,24 @@ func takeLastComment(list []Visitee) (*Comment, []Visitee) {
 // mergeOrReturnComment creates a new comment and tries to merge it with the last element (if is a comment and is on the next line).
 func mergeOrReturnComment(elements []Visitee, lit string, pos scanner.Position) *Comment {
 	com := newComment(pos, lit)
-	// last element must be a comment to merge +
-	// do not merge c-style comments +
-	// last comment line was on previous line
-	if esize := len(elements); esize > 0 {
-		if last, ok := elements[esize-1].(*Comment); ok &&
-			!last.Cstyle &&
-			pos.Line <= last.Position.Line+len(last.Lines) { // less than because last line of file could be inline comment
-			last.Merge(com)
-			// mark as merged
-			com = nil
-		}
+	esize := len(elements)
+	if esize == 0 {
+		return com
+	}
+	// last element must be a comment to merge
+	last, ok := elements[esize-1].(*Comment)
+	if !ok {
+		return com
+	}
+	// do not merge c-style comments
+	if last.Cstyle {
+		return com
+	}
+	// last comment has text on previous line
+	// TODO handle last line of file could be inline comment
+	if !last.hasTextOnLine(pos.Line - 1) {
+		return com
 	}
-	return com
+	last.Merge(com)
+	return nil
 }

comment_test.go

@@ -54,7 +54,7 @@ world`)
 func TestTakeLastComment(t *testing.T) {
 	c0 := newComment(startPosition, "hi")
 	c1 := newComment(startPosition, "there")
-	_, l := takeLastComment([]Visitee{c0, c1})
+	_, l := takeLastCommentIfEndsOnLine([]Visitee{c0, c1}, 1)
 	if got, want := len(l), 1; got != want {
 		t.Fatalf("got [%v] want [%v]", got, want)
 	}
@@ -235,3 +235,46 @@ func TestParseCommentWithTripleSlash(t *testing.T) {
 		t.Fatalf("got [%d] want [%d]", got, want)
 	}
 }
+
+func TestCommentAssociation(t *testing.T) {
+	src := `
+	// foo1
+	// foo2
+
+	// bar
+	
+	syntax = "proto3";
+	
+	// baz
+	
+	// bat1
+	// bat2
+	package bat;
+	
+	// Oneway is the return type to use for an rpc method if
+	// the method should be generated as oneway.
+	message Oneway {
+	  bool ack = 1;
+	}`
+	p := newParserOn(src)
+	def, err := p.Parse()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got, want := len(def.Elements), 6; got != want {
+		t.Fatalf("got [%v] want [%v]", got, want)
+	}
+	pkg := def.Elements[4].(*Package)
+	if got, want := pkg.Comment.Message(), " bat1"; got != want {
+		t.Fatalf("got [%v] want [%v]", got, want)
+	}
+	if got, want := len(pkg.Comment.Lines), 2; got != want {
+		t.Fatalf("got [%v] want [%v]", got, want)
+	}
+	if got, want := pkg.Comment.Lines[1], " bat2"; got != want {
+		t.Fatalf("got [%v] want [%v]", got, want)
+	}
+	if got, want := len(def.Elements[5].(*Message).Comment.Lines), 2; got != want {
+		t.Fatalf("got [%v] want [%v]", got, want)
+	}
+}

enum.go

@@ -57,8 +57,8 @@ func (e *Enum) elements() []Visitee {
 
 // takeLastComment is part of elementContainer
 // removes and returns the last element of the list if it is a Comment.
-func (e *Enum) takeLastComment() (last *Comment) {
-	last, e.Elements = takeLastComment(e.Elements)
+func (e *Enum) takeLastComment(expectedOnLine int) (last *Comment) {
+	last, e.Elements = takeLastCommentIfEndsOnLine(e.Elements, expectedOnLine)
 	return
 }
 
@@ -84,7 +84,7 @@ func (e *Enum) parse(p *Parser) error {
 		case tOPTION:
 			v := new(Option)
 			v.Position = pos
-			v.Comment = e.takeLastComment()
+			v.Comment = e.takeLastComment(pos.Line)
 			err := v.parse(p)
 			if err != nil {
 				return err
@@ -98,7 +98,7 @@ func (e *Enum) parse(p *Parser) error {
 			p.nextPut(pos, tok, lit)
 			f := new(EnumField)
 			f.Position = pos
-			f.Comment = e.takeLastComment()
+			f.Comment = e.takeLastComment(pos.Line - 1)
 			err := f.parse(p)
 			if err != nil {
 				return err

group.go

@@ -62,8 +62,8 @@ func (g *Group) Doc() *Comment {
 
 // takeLastComment is part of elementContainer
 // removes and returns the last element of the list if it is a Comment.
-func (g *Group) takeLastComment() (last *Comment) {
-	last, g.Elements = takeLastComment(g.Elements)
+func (g *Group) takeLastComment(expectedOnLine int) (last *Comment) {
+	last, g.Elements = takeLastCommentIfEndsOnLine(g.Elements, expectedOnLine)
 	return
 }
 

message.go

@@ -76,47 +76,47 @@ func parseMessageBody(p *Parser, c elementContainer) error {
 		case tENUM == tok:
 			e := new(Enum)
 			e.Position = pos
-			e.Comment = c.takeLastComment()
+			e.Comment = c.takeLastComment(pos.Line - 1)
 			if err := e.parse(p); err != nil {
 				return err
 			}
 			c.addElement(e)
 		case tMESSAGE == tok:
 			msg := new(Message)
 			msg.Position = pos
-			msg.Comment = c.takeLastComment()
+			msg.Comment = c.takeLastComment(pos.Line - 1)
 			if err := msg.parse(p); err != nil {
 				return err
 			}
 			c.addElement(msg)
 		case tOPTION == tok:
 			o := new(Option)
 			o.Position = pos
-			o.Comment = c.takeLastComment()
+			o.Comment = c.takeLastComment(pos.Line - 1)
 			if err := o.parse(p); err != nil {
 				return err
 			}
 			c.addElement(o)
 		case tONEOF == tok:
 			o := new(Oneof)
 			o.Position = pos
-			o.Comment = c.takeLastComment()
+			o.Comment = c.takeLastComment(pos.Line - 1)
 			if err := o.parse(p); err != nil {
 				return err
 			}
 			c.addElement(o)
 		case tMAP == tok:
 			f := newMapField()
 			f.Position = pos
-			f.Comment = c.takeLastComment()
+			f.Comment = c.takeLastComment(pos.Line - 1)
 			if err := f.parse(p); err != nil {
 				return err
 			}
 			c.addElement(f)
 		case tRESERVED == tok:
 			r := new(Reserved)
 			r.Position = pos
-			r.Comment = c.takeLastComment()
+			r.Comment = c.takeLastComment(pos.Line - 1)
 			if err := r.parse(p); err != nil {
 				return err
 			}
@@ -129,7 +129,7 @@ func parseMessageBody(p *Parser, c elementContainer) error {
 			if tGROUP == tok {
 				g := new(Group)
 				g.Position = pos
-				g.Comment = c.takeLastComment()
+				g.Comment = c.takeLastComment(pos.Line - 1)
 				g.Optional = prevTok == tOPTIONAL
 				g.Repeated = prevTok == tREPEATED
 				g.Required = prevTok == tREQUIRED
@@ -143,7 +143,7 @@ func parseMessageBody(p *Parser, c elementContainer) error {
 				f := newNormalField()
 				f.Type = lit
 				f.Position = pos
-				f.Comment = c.takeLastComment()
+				f.Comment = c.takeLastComment(pos.Line - 1)
 				f.Optional = prevTok == tOPTIONAL
 				f.Repeated = prevTok == tREPEATED
 				f.Required = prevTok == tREQUIRED
@@ -155,23 +155,23 @@ func parseMessageBody(p *Parser, c elementContainer) error {
 		case tGROUP == tok:
 			g := new(Group)
 			g.Position = pos
-			g.Comment = c.takeLastComment()
+			g.Comment = c.takeLastComment(pos.Line - 1)
 			if err := g.parse(p); err != nil {
 				return err
 			}
 			c.addElement(g)
 		case tEXTENSIONS == tok:
 			e := new(Extensions)
 			e.Position = pos
-			e.Comment = c.takeLastComment()
+			e.Comment = c.takeLastComment(pos.Line - 1)
 			if err := e.parse(p); err != nil {
 				return err
 			}
 			c.addElement(e)
 		case tEXTEND == tok:
 			e := new(Message)
 			e.Position = pos
-			e.Comment = c.takeLastComment()
+			e.Comment = c.takeLastComment(pos.Line - 1)
 			e.IsExtend = true
 			if err := e.parse(p); err != nil {
 				return err
@@ -188,7 +188,7 @@ func parseMessageBody(p *Parser, c elementContainer) error {
 			p.nextPut(pos, tok, lit)
 			f := newNormalField()
 			f.Position = pos
-			f.Comment = c.takeLastComment()
+			f.Comment = c.takeLastComment(pos.Line - 1)
 			if err := f.parse(p); err != nil {
 				return err
 			}
@@ -217,8 +217,8 @@ func (m *Message) elements() []Visitee {
 	return m.Elements
 }
 
-func (m *Message) takeLastComment() (last *Comment) {
-	last, m.Elements = takeLastComment(m.Elements)
+func (m *Message) takeLastComment(expectedOnLine int) (last *Comment) {
+	last, m.Elements = takeLastCommentIfEndsOnLine(m.Elements, expectedOnLine)
 	return
 }
 

oneof.go

@@ -47,8 +47,8 @@ func (o *Oneof) elements() []Visitee {
 
 // takeLastComment is part of elementContainer
 // removes and returns the last element of the list if it is a Comment.
-func (o *Oneof) takeLastComment() (last *Comment) {
-	last, o.Elements = takeLastComment(o.Elements)
+func (o *Oneof) takeLastComment(expectedOnLine int) (last *Comment) {
+	last, o.Elements = takeLastCommentIfEndsOnLine(o.Elements, expectedOnLine)
 	return last
 }
 
@@ -76,7 +76,7 @@ func (o *Oneof) parse(p *Parser) error {
 		case tIDENT:
 			f := newOneOfField()
 			f.Position = pos
-			f.Comment, o.Elements = takeLastComment(o.elements())
+			f.Comment, o.Elements = takeLastCommentIfEndsOnLine(o.elements(), pos.Line-1) // TODO call takeLastComment instead?
 			f.Type = lit
 			if err := parseFieldAfterType(f.Field, p); err != nil {
 				return err
@@ -85,7 +85,7 @@ func (o *Oneof) parse(p *Parser) error {
 		case tGROUP:
 			g := new(Group)
 			g.Position = pos
-			g.Comment, o.Elements = takeLastComment(o.elements())
+			g.Comment, o.Elements = takeLastCommentIfEndsOnLine(o.elements(), pos.Line-1)
 			if err := g.parse(p); err != nil {
 				return err
 			}

parser_test.go

@@ -49,7 +49,7 @@ func TestParseComment(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	if got, want := len(collect(pr).Comments()), 2; got != want {
+	if got, want := len(collect(pr).Comments()), 3; got != want {
 		t.Errorf("got [%v] want [%v]", got, want)
 	}
 }