github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/codegen/loaderx/comment_scanner.go (about)

     1  package loaderx
     2  
     3  import (
     4  	"go/ast"
     5  	"go/token"
     6  	"sort"
     7  	"strings"
     8  )
     9  
    10  func CommentsOf(fileSet *token.FileSet, targetNode ast.Node, files ...*ast.File) string {
    11  	file := FileOf(targetNode, files...)
    12  	if file == nil {
    13  		return ""
    14  	}
    15  	commentScanner := NewCommentScanner(fileSet, file)
    16  	doc := commentScanner.CommentsOf(targetNode)
    17  	if doc != "" {
    18  		return doc
    19  	}
    20  	return doc
    21  }
    22  
    23  func NewCommentScanner(fileSet *token.FileSet, file *ast.File) *CommentScanner {
    24  	commentMap := ast.NewCommentMap(fileSet, file, file.Comments)
    25  
    26  	return &CommentScanner{
    27  		file:       file,
    28  		CommentMap: commentMap,
    29  	}
    30  }
    31  
    32  type CommentScanner struct {
    33  	file       *ast.File
    34  	CommentMap ast.CommentMap
    35  }
    36  
    37  func (scanner *CommentScanner) CommentsOf(targetNode ast.Node) string {
    38  	commentGroupList := scanner.CommentGroupListOf(targetNode)
    39  	return StringifyCommentGroup(commentGroupList...)
    40  }
    41  
    42  func (scanner *CommentScanner) CommentGroupListOf(targetNode ast.Node) (commentGroupList []*ast.CommentGroup) {
    43  	if targetNode == nil {
    44  		return
    45  	}
    46  
    47  	switch targetNode.(type) {
    48  	case *ast.File, *ast.Field, ast.Stmt, ast.Decl:
    49  		if comments, ok := scanner.CommentMap[targetNode]; ok {
    50  			commentGroupList = comments
    51  		}
    52  	case ast.Spec:
    53  		// Spec should merge with comments of its parent gen decl when empty
    54  		if comments, ok := scanner.CommentMap[targetNode]; ok {
    55  			commentGroupList = append(commentGroupList, comments...)
    56  		}
    57  
    58  		if len(commentGroupList) == 0 {
    59  			for node, comments := range scanner.CommentMap {
    60  				if genDecl, ok := node.(*ast.GenDecl); ok {
    61  					for _, spec := range genDecl.Specs {
    62  						if targetNode == spec {
    63  							commentGroupList = append(commentGroupList, comments...)
    64  						}
    65  					}
    66  				}
    67  			}
    68  		}
    69  	default:
    70  		// find nearest parent node which have comments
    71  		{
    72  			var deltaPos token.Pos
    73  			var parentNode ast.Node
    74  
    75  			deltaPos = -1
    76  
    77  			ast.Inspect(scanner.file, func(node ast.Node) bool {
    78  				switch node.(type) {
    79  				case *ast.Field, ast.Decl, ast.Spec, ast.Stmt:
    80  					if targetNode.Pos() >= node.Pos() && targetNode.End() <= node.End() {
    81  						nextDelta := targetNode.Pos() - node.Pos()
    82  						if deltaPos == -1 || (nextDelta <= deltaPos) {
    83  							deltaPos = nextDelta
    84  							parentNode = node
    85  						}
    86  					}
    87  				}
    88  				return true
    89  			})
    90  
    91  			if parentNode != nil {
    92  				commentGroupList = scanner.CommentGroupListOf(parentNode)
    93  			}
    94  		}
    95  	}
    96  
    97  	sort.Sort(ByCommentPos(commentGroupList))
    98  	return
    99  }
   100  
   101  type ByCommentPos []*ast.CommentGroup
   102  
   103  func (a ByCommentPos) Len() int {
   104  	return len(a)
   105  }
   106  
   107  func (a ByCommentPos) Swap(i, j int) {
   108  	a[i], a[j] = a[j], a[i]
   109  }
   110  
   111  func (a ByCommentPos) Less(i, j int) bool {
   112  	return a[i].Pos() < a[j].Pos()
   113  }
   114  
   115  func StringifyCommentGroup(commentGroupList ...*ast.CommentGroup) (comments string) {
   116  	if len(commentGroupList) == 0 {
   117  		return ""
   118  	}
   119  	for _, commentGroup := range commentGroupList {
   120  		for _, line := range strings.Split(commentGroup.Text(), "\n") {
   121  			if strings.HasPrefix(line, "go:generate") {
   122  				continue
   123  			}
   124  			comments = comments + "\n" + line
   125  		}
   126  	}
   127  	return strings.TrimSpace(comments)
   128  }