commentmap.go

     1  // Copyright 2012 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package ast
     6  
     7  import (
     8  	"bytes"
     9  	"cmp"
    10  	"fmt"
    11  	"go/token"
    12  	"slices"
    13  	"strings"
    14  )
    15  
    16  // sortComments sorts the list of comment groups in source order.
    17  func sortComments(list []*CommentGroup) {
    18  	slices.SortFunc(list, func(a, b *CommentGroup) int {
    19  		return cmp.Compare(a.Pos(), b.Pos())
    20  	})
    21  }
    22  
    23  // A CommentMap maps an AST node to a list of comment groups
    24  // associated with it. See [NewCommentMap] for a description of
    25  // the association.
    26  type CommentMap map[Node][]*CommentGroup
    27  
    28  func (cmap CommentMap) addComment(n Node, c *CommentGroup) {
    29  	list := cmap[n]
    30  	if len(list) == 0 {
    31  		list = []*CommentGroup{c}
    32  	} else {
    33  		list = append(list, c)
    34  	}
    35  	cmap[n] = list
    36  }
    37  
    38  // nodeList returns the list of nodes of the AST n in source order.
    39  func nodeList(n Node) []Node {
    40  	var list []Node
    41  	Inspect(n, func(n Node) bool {
    42  		// don't collect comments
    43  		switch n.(type) {
    44  		case nil, *CommentGroup, *Comment:
    45  			return false
    46  		}
    47  		list = append(list, n)
    48  		return true
    49  	})
    50  
    51  	// Note: The current implementation assumes that Inspect traverses the
    52  	//       AST in depth-first and thus _source_ order. If AST traversal
    53  	//       does not follow source order, the sorting call below will be
    54  	//       required.
    55  	// slices.Sort(list, func(a, b Node) int {
    56  	//       r := cmp.Compare(a.Pos(), b.Pos())
    57  	//       if r != 0 {
    58  	//               return r
    59  	//       }
    60  	//       return cmp.Compare(a.End(), b.End())
    61  	// })
    62  
    63  	return list
    64  }
    65  
    66  // A commentListReader helps iterating through a list of comment groups.
    67  type commentListReader struct {
    68  	fset     *token.FileSet
    69  	list     []*CommentGroup
    70  	index    int
    71  	comment  *CommentGroup  // comment group at current index
    72  	pos, end token.Position // source interval of comment group at current index
    73  }
    74  
    75  func (r *commentListReader) eol() bool {
    76  	return r.index >= len(r.list)
    77  }
    78  
    79  func (r *commentListReader) next() {
    80  	if !r.eol() {
    81  		r.comment = r.list[r.index]
    82  		r.pos = r.fset.Position(r.comment.Pos())
    83  		r.end = r.fset.Position(r.comment.End())
    84  		r.index++
    85  	}
    86  }
    87  
    88  // A nodeStack keeps track of nested nodes.
    89  // A node lower on the stack lexically contains the nodes higher on the stack.
    90  type nodeStack []Node
    91  
    92  // push pops all nodes that appear lexically before n
    93  // and then pushes n on the stack.
    94  func (s *nodeStack) push(n Node) {
    95  	s.pop(n.Pos())
    96  	*s = append((*s), n)
    97  }
    98  
    99  // pop pops all nodes that appear lexically before pos
   100  // (i.e., whose lexical extent has ended before or at pos).
   101  // It returns the last node popped.
   102  func (s *nodeStack) pop(pos token.Pos) (top Node) {
   103  	i := len(*s)
   104  	for i > 0 && (*s)[i-1].End() <= pos {
   105  		top = (*s)[i-1]
   106  		i--
   107  	}
   108  	*s = (*s)[0:i]
   109  	return top
   110  }
   111  
   112  // NewCommentMap creates a new comment map by associating comment groups
   113  // of the comments list with the nodes of the AST specified by node.
   114  //
   115  // A comment group g is associated with a node n if:
   116  //
   117  //   - g starts on the same line as n ends
   118  //   - g starts on the line immediately following n, and there is
   119  //     at least one empty line after g and before the next node
   120  //   - g starts before n and is not associated to the node before n
   121  //     via the previous rules
   122  //
   123  // NewCommentMap tries to associate a comment group to the "largest"
   124  // node possible: For instance, if the comment is a line comment
   125  // trailing an assignment, the comment is associated with the entire
   126  // assignment rather than just the last operand in the assignment.
   127  func NewCommentMap(fset *token.FileSet, node Node, comments []*CommentGroup) CommentMap {
   128  	if len(comments) == 0 {
   129  		return nil // no comments to map
   130  	}
   131  
   132  	cmap := make(CommentMap)
   133  
   134  	// set up comment reader r
   135  	tmp := make([]*CommentGroup, len(comments))
   136  	copy(tmp, comments) // don't change incoming comments
   137  	sortComments(tmp)
   138  	r := commentListReader{fset: fset, list: tmp} // !r.eol() because len(comments) > 0
   139  	r.next()
   140  
   141  	// create node list in lexical order
   142  	nodes := nodeList(node)
   143  	nodes = append(nodes, nil) // append sentinel
   144  
   145  	// set up iteration variables
   146  	var (
   147  		p     Node           // previous node
   148  		pend  token.Position // end of p
   149  		pg    Node           // previous node group (enclosing nodes of "importance")
   150  		pgend token.Position // end of pg
   151  		stack nodeStack      // stack of node groups
   152  	)
   153  
   154  	for _, q := range nodes {
   155  		var qpos token.Position
   156  		if q != nil {
   157  			qpos = fset.Position(q.Pos()) // current node position
   158  		} else {
   159  			// set fake sentinel position to infinity so that
   160  			// all comments get processed before the sentinel
   161  			const infinity = 1 << 30
   162  			qpos.Offset = infinity
   163  			qpos.Line = infinity
   164  		}
   165  
   166  		// process comments before current node
   167  		for r.end.Offset <= qpos.Offset {
   168  			// determine recent node group
   169  			if top := stack.pop(r.comment.Pos()); top != nil {
   170  				pg = top
   171  				pgend = fset.Position(pg.End())
   172  			}
   173  			// Try to associate a comment first with a node group
   174  			// (i.e., a node of "importance" such as a declaration);
   175  			// if that fails, try to associate it with the most recent
   176  			// node.
   177  			// TODO(gri) try to simplify the logic below
   178  			var assoc Node
   179  			switch {
   180  			case pg != nil &&
   181  				(pgend.Line == r.pos.Line ||
   182  					pgend.Line+1 == r.pos.Line && r.end.Line+1 < qpos.Line):
   183  				// 1) comment starts on same line as previous node group ends, or
   184  				// 2) comment starts on the line immediately after the
   185  				//    previous node group and there is an empty line before
   186  				//    the current node
   187  				// => associate comment with previous node group
   188  				assoc = pg
   189  			case p != nil &&
   190  				(pend.Line == r.pos.Line ||
   191  					pend.Line+1 == r.pos.Line && r.end.Line+1 < qpos.Line ||
   192  					q == nil):
   193  				// same rules apply as above for p rather than pg,
   194  				// but also associate with p if we are at the end (q == nil)
   195  				assoc = p
   196  			default:
   197  				// otherwise, associate comment with current node
   198  				if q == nil {
   199  					// we can only reach here if there was no p
   200  					// which would imply that there were no nodes
   201  					panic("internal error: no comments should be associated with sentinel")
   202  				}
   203  				assoc = q
   204  			}
   205  			cmap.addComment(assoc, r.comment)
   206  			if r.eol() {
   207  				return cmap
   208  			}
   209  			r.next()
   210  		}
   211  
   212  		// update previous node
   213  		p = q
   214  		pend = fset.Position(p.End())
   215  
   216  		// update previous node group if we see an "important" node
   217  		switch q.(type) {
   218  		case *File, *Field, Decl, Spec, Stmt:
   219  			stack.push(q)
   220  		}
   221  	}
   222  
   223  	return cmap
   224  }
   225  
   226  // Update replaces an old node in the comment map with the new node
   227  // and returns the new node. Comments that were associated with the
   228  // old node are associated with the new node.
   229  func (cmap CommentMap) Update(old, new Node) Node {
   230  	if list := cmap[old]; len(list) > 0 {
   231  		delete(cmap, old)
   232  		cmap[new] = append(cmap[new], list...)
   233  	}
   234  	return new
   235  }
   236  
   237  // Filter returns a new comment map consisting of only those
   238  // entries of cmap for which a corresponding node exists in
   239  // the AST specified by node.
   240  func (cmap CommentMap) Filter(node Node) CommentMap {
   241  	umap := make(CommentMap)
   242  	Inspect(node, func(n Node) bool {
   243  		if g := cmap[n]; len(g) > 0 {
   244  			umap[n] = g
   245  		}
   246  		return true
   247  	})
   248  	return umap
   249  }
   250  
   251  // Comments returns the list of comment groups in the comment map.
   252  // The result is sorted in source order.
   253  func (cmap CommentMap) Comments() []*CommentGroup {
   254  	list := make([]*CommentGroup, 0, len(cmap))
   255  	for _, e := range cmap {
   256  		list = append(list, e...)
   257  	}
   258  	sortComments(list)
   259  	return list
   260  }
   261  
   262  func summary(list []*CommentGroup) string {
   263  	const maxLen = 40
   264  	var buf bytes.Buffer
   265  
   266  	// collect comments text
   267  loop:
   268  	for _, group := range list {
   269  		// Note: CommentGroup.Text() does too much work for what we
   270  		//       need and would only replace this innermost loop.
   271  		//       Just do it explicitly.
   272  		for _, comment := range group.List {
   273  			if buf.Len() >= maxLen {
   274  				break loop
   275  			}
   276  			buf.WriteString(comment.Text)
   277  		}
   278  	}
   279  
   280  	// truncate if too long
   281  	if buf.Len() > maxLen {
   282  		buf.Truncate(maxLen - 3)
   283  		buf.WriteString("...")
   284  	}
   285  
   286  	// replace any invisibles with blanks
   287  	bytes := buf.Bytes()
   288  	for i, b := range bytes {
   289  		switch b {
   290  		case '\t', '\n', '\r':
   291  			bytes[i] = ' '
   292  		}
   293  	}
   294  
   295  	return string(bytes)
   296  }
   297  
   298  func (cmap CommentMap) String() string {
   299  	// print map entries in sorted order
   300  	var nodes []Node
   301  	for node := range cmap {
   302  		nodes = append(nodes, node)
   303  	}
   304  	slices.SortFunc(nodes, func(a, b Node) int {
   305  		r := cmp.Compare(a.Pos(), b.Pos())
   306  		if r != 0 {
   307  			return r
   308  		}
   309  		return cmp.Compare(a.End(), b.End())
   310  	})
   311  
   312  	var buf strings.Builder
   313  	fmt.Fprintln(&buf, "CommentMap {")
   314  	for _, node := range nodes {
   315  		comment := cmap[node]
   316  		// print name of identifiers; print node type for other nodes
   317  		var s string
   318  		if ident, ok := node.(*Ident); ok {
   319  			s = ident.Name
   320  		} else {
   321  			s = fmt.Sprintf("%T", node)
   322  		}
   323  		fmt.Fprintf(&buf, "\t%p  %20s:  %s\n", node, s, summary(comment))
   324  	}
   325  	fmt.Fprintln(&buf, "}")
   326  	return buf.String()
   327  }
   328
View as plain text