1package whitespace
2
3import (
4	"go/ast"
5	"go/token"
6)
7
8// Message contains a message
9type Message struct {
10	Pos     token.Position
11	Type    MessageType
12	Message string
13}
14
15// MessageType describes what should happen to fix the warning
16type MessageType uint8
17
18// List of MessageTypes
19const (
20	MessageTypeLeading MessageType = iota + 1
21	MessageTypeTrailing
22	MessageTypeAddAfter
23)
24
25// Settings contains settings for edge-cases
26type Settings struct {
27	MultiIf   bool
28	MultiFunc bool
29}
30
31// Run runs this linter on the provided code
32func Run(file *ast.File, fset *token.FileSet, settings Settings) []Message {
33	var messages []Message
34
35	for _, f := range file.Decls {
36		decl, ok := f.(*ast.FuncDecl)
37		if !ok || decl.Body == nil { // decl.Body can be nil for e.g. cgo
38			continue
39		}
40
41		vis := visitor{file.Comments, fset, nil, make(map[*ast.BlockStmt]bool), settings}
42		ast.Walk(&vis, decl)
43
44		messages = append(messages, vis.messages...)
45	}
46
47	return messages
48}
49
50type visitor struct {
51	comments    []*ast.CommentGroup
52	fset        *token.FileSet
53	messages    []Message
54	wantNewline map[*ast.BlockStmt]bool
55	settings    Settings
56}
57
58func (v *visitor) Visit(node ast.Node) ast.Visitor {
59	if node == nil {
60		return v
61	}
62
63	if stmt, ok := node.(*ast.IfStmt); ok && v.settings.MultiIf {
64		checkMultiLine(v, stmt.Body, stmt.Cond)
65	}
66
67	if stmt, ok := node.(*ast.FuncDecl); ok && v.settings.MultiFunc {
68		checkMultiLine(v, stmt.Body, stmt.Type)
69	}
70
71	if stmt, ok := node.(*ast.BlockStmt); ok {
72		wantNewline := v.wantNewline[stmt]
73
74		comments := v.comments
75		if wantNewline {
76			comments = nil // Comments also count as a newline if we want a newline
77		}
78		first, last := firstAndLast(comments, v.fset, stmt.Pos(), stmt.End(), stmt.List)
79
80		startMsg := checkStart(v.fset, stmt.Lbrace, first)
81
82		if wantNewline && startMsg == nil {
83			v.messages = append(v.messages, Message{v.fset.Position(stmt.Pos()), MessageTypeAddAfter, `multi-line statement should be followed by a newline`})
84		} else if !wantNewline && startMsg != nil {
85			v.messages = append(v.messages, *startMsg)
86		}
87
88		if msg := checkEnd(v.fset, stmt.Rbrace, last); msg != nil {
89			v.messages = append(v.messages, *msg)
90		}
91	}
92
93	return v
94}
95
96func checkMultiLine(v *visitor, body *ast.BlockStmt, stmtStart ast.Node) {
97	start, end := posLine(v.fset, stmtStart.Pos()), posLine(v.fset, stmtStart.End())
98
99	if end > start { // Check only multi line conditions
100		v.wantNewline[body] = true
101	}
102}
103
104func posLine(fset *token.FileSet, pos token.Pos) int {
105	return fset.Position(pos).Line
106}
107
108func firstAndLast(comments []*ast.CommentGroup, fset *token.FileSet, start, end token.Pos, stmts []ast.Stmt) (ast.Node, ast.Node) {
109	if len(stmts) == 0 {
110		return nil, nil
111	}
112
113	first, last := ast.Node(stmts[0]), ast.Node(stmts[len(stmts)-1])
114
115	for _, c := range comments {
116		if posLine(fset, c.Pos()) == posLine(fset, start) || posLine(fset, c.End()) == posLine(fset, end) {
117			continue
118		}
119
120		if c.Pos() < start || c.End() > end {
121			continue
122		}
123		if c.Pos() < first.Pos() {
124			first = c
125		}
126		if c.End() > last.End() {
127			last = c
128		}
129	}
130
131	return first, last
132}
133
134func checkStart(fset *token.FileSet, start token.Pos, first ast.Node) *Message {
135	if first == nil {
136		return nil
137	}
138
139	if posLine(fset, start)+1 < posLine(fset, first.Pos()) {
140		pos := fset.Position(start)
141		return &Message{pos, MessageTypeLeading, `unnecessary leading newline`}
142	}
143
144	return nil
145}
146
147func checkEnd(fset *token.FileSet, end token.Pos, last ast.Node) *Message {
148	if last == nil {
149		return nil
150	}
151
152	if posLine(fset, end)-1 > posLine(fset, last.End()) {
153		pos := fset.Position(end)
154		return &Message{pos, MessageTypeTrailing, `unnecessary trailing newline`}
155	}
156
157	return nil
158}
159