fix + more tests for comment association

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 takeLastCommentIfOnLine(list []Visitee, line int) (*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 && last.Position.Line == line {
+	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 takeLastCommentIfOnLine(list []Visitee, line int) (*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 := takeLastCommentIfOnLine([]Visitee{c0, c1}, 1)
+	_, l := takeLastCommentIfEndsOnLine([]Visitee{c0, c1}, 1)
 	if got, want := len(l), 1; got != want {
 		t.Fatalf("got [%v] want [%v]", got, want)
 	}
@@ -247,6 +247,8 @@ func TestCommentAssociation(t *testing.T) {
 	
 	// baz
 	
+	// bat1
+	// bat2
 	package bat;`
 	p := newParserOn(src)
 	def, err := p.Parse()
@@ -256,4 +258,14 @@ func TestCommentAssociation(t *testing.T) {
 	if got, want := len(def.Elements), 5; 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)
+	}
 }

enum.go

@@ -58,7 +58,7 @@ 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(expectedOnLine int) (last *Comment) {
-	last, e.Elements = takeLastCommentIfOnLine(e.Elements, expectedOnLine)
+	last, e.Elements = takeLastCommentIfEndsOnLine(e.Elements, expectedOnLine)
 	return
 }
 

group.go

@@ -63,7 +63,7 @@ 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(expectedOnLine int) (last *Comment) {
-	last, g.Elements = takeLastCommentIfOnLine(g.Elements, expectedOnLine)
+	last, g.Elements = takeLastCommentIfEndsOnLine(g.Elements, expectedOnLine)
 	return
 }
 

message.go

@@ -218,7 +218,7 @@ func (m *Message) elements() []Visitee {
 }
 
 func (m *Message) takeLastComment(expectedOnLine int) (last *Comment) {
-	last, m.Elements = takeLastCommentIfOnLine(m.Elements, expectedOnLine)
+	last, m.Elements = takeLastCommentIfEndsOnLine(m.Elements, expectedOnLine)
 	return
 }
 

oneof.go

@@ -48,7 +48,7 @@ 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(expectedOnLine int) (last *Comment) {
-	last, o.Elements = takeLastCommentIfOnLine(o.Elements, expectedOnLine)
+	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 = takeLastCommentIfOnLine(o.elements(), pos.Line-1) // TODO call takeLastComment instead?
+			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 = takeLastCommentIfOnLine(o.elements(), pos.Line-1)
+			g.Comment, o.Elements = takeLastCommentIfEndsOnLine(o.elements(), pos.Line-1)
 			if err := g.parse(p); err != nil {
 				return err
 			}

proto.go

@@ -42,7 +42,7 @@ func (proto *Proto) elements() []Visitee {
 // takeLastComment is part of elementContainer
 // removes and returns the last element of the list if it is a Comment.
 func (proto *Proto) takeLastComment(expectedOnLine int) (last *Comment) {
-	last, proto.Elements = takeLastCommentIfOnLine(proto.Elements, expectedOnLine)
+	last, proto.Elements = takeLastCommentIfEndsOnLine(proto.Elements, expectedOnLine)
 	return
 }
 
@@ -58,39 +58,39 @@ func (proto *Proto) parse(p *Parser) error {
 		case tOPTION == tok:
 			o := new(Option)
 			o.Position = pos
-			o.Comment, proto.Elements = takeLastCommentIfOnLine(proto.Elements, pos.Line-1)
+			o.Comment, proto.Elements = takeLastCommentIfEndsOnLine(proto.Elements, pos.Line-1)
 			if err := o.parse(p); err != nil {
 				return err
 			}
 			proto.Elements = append(proto.Elements, o)
 		case tSYNTAX == tok:
 			s := new(Syntax)
 			s.Position = pos
-			s.Comment, proto.Elements = takeLastCommentIfOnLine(proto.Elements, pos.Line-1)
+			s.Comment, proto.Elements = takeLastCommentIfEndsOnLine(proto.Elements, pos.Line-1)
 			if err := s.parse(p); err != nil {
 				return err
 			}
 			proto.Elements = append(proto.Elements, s)
 		case tIMPORT == tok:
 			im := new(Import)
 			im.Position = pos
-			im.Comment, proto.Elements = takeLastCommentIfOnLine(proto.Elements, pos.Line-1)
+			im.Comment, proto.Elements = takeLastCommentIfEndsOnLine(proto.Elements, pos.Line-1)
 			if err := im.parse(p); err != nil {
 				return err
 			}
 			proto.Elements = append(proto.Elements, im)
 		case tENUM == tok:
 			enum := new(Enum)
 			enum.Position = pos
-			enum.Comment, proto.Elements = takeLastCommentIfOnLine(proto.Elements, pos.Line-1)
+			enum.Comment, proto.Elements = takeLastCommentIfEndsOnLine(proto.Elements, pos.Line-1)
 			if err := enum.parse(p); err != nil {
 				return err
 			}
 			proto.Elements = append(proto.Elements, enum)
 		case tSERVICE == tok:
 			service := new(Service)
 			service.Position = pos
-			service.Comment, proto.Elements = takeLastCommentIfOnLine(proto.Elements, pos.Line-1)
+			service.Comment, proto.Elements = takeLastCommentIfEndsOnLine(proto.Elements, pos.Line-1)
 			err := service.parse(p)
 			if err != nil {
 				return err
@@ -99,15 +99,15 @@ func (proto *Proto) parse(p *Parser) error {
 		case tPACKAGE == tok:
 			pkg := new(Package)
 			pkg.Position = pos
-			pkg.Comment, proto.Elements = takeLastCommentIfOnLine(proto.Elements, pos.Line-1)
+			pkg.Comment, proto.Elements = takeLastCommentIfEndsOnLine(proto.Elements, pos.Line-1)
 			if err := pkg.parse(p); err != nil {
 				return err
 			}
 			proto.Elements = append(proto.Elements, pkg)
 		case tMESSAGE == tok:
 			msg := new(Message)
 			msg.Position = pos
-			msg.Comment, proto.Elements = takeLastCommentIfOnLine(proto.Elements, pos.Line-1)
+			msg.Comment, proto.Elements = takeLastCommentIfEndsOnLine(proto.Elements, pos.Line-1)
 			if err := msg.parse(p); err != nil {
 				return err
 			}
@@ -116,7 +116,7 @@ func (proto *Proto) parse(p *Parser) error {
 		case tEXTEND == tok:
 			msg := new(Message)
 			msg.Position = pos
-			msg.Comment, proto.Elements = takeLastCommentIfOnLine(proto.Elements, pos.Line-1)
+			msg.Comment, proto.Elements = takeLastCommentIfEndsOnLine(proto.Elements, pos.Line-1)
 			msg.IsExtend = true
 			if err := msg.parse(p); err != nil {
 				return err

service.go

@@ -58,7 +58,7 @@ func (s *Service) elements() []Visitee {
 // takeLastComment is part of elementContainer
 // removes and returns the last elements of the list if it is a Comment.
 func (s *Service) takeLastComment(expectedOnLine int) (last *Comment) {
-	last, s.Elements = takeLastCommentIfOnLine(s.Elements, expectedOnLine)
+	last, s.Elements = takeLastCommentIfEndsOnLine(s.Elements, expectedOnLine)
 	return
 }
 
@@ -85,7 +85,7 @@ func (s *Service) parse(p *Parser) error {
 		case tRPC:
 			rpc := new(RPC)
 			rpc.Position = pos
-			rpc.Comment, s.Elements = takeLastCommentIfOnLine(s.Elements, pos.Line-1)
+			rpc.Comment, s.Elements = takeLastCommentIfEndsOnLine(s.Elements, pos.Line-1)
 			err := rpc.parse(p)
 			if err != nil {
 				return err
@@ -222,7 +222,7 @@ func (r *RPC) elements() []Visitee {
 	return r.Elements
 }
 
-func (r *RPC) takeLastComment() (last *Comment) {
-	last, r.Elements = takeLastComment(r.Elements)
+func (r *RPC) takeLastComment(expectedOnLine int) (last *Comment) {
+	last, r.Elements = takeLastCommentIfEndsOnLine(r.Elements, expectedOnLine)
 	return
 }