/
util.go
316 lines (288 loc) · 8.01 KB
/
util.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
// Copyright 2018 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.
package pretty
// This file contains utility and other non-standard functions not listed
// in the paper (see the package documentation). Utility functions are those
// that are generally useful or need to create certain structures (like union)
// in a correct way.
// Join joins Docs d with string s and Line. There is no space between each
// item in d and the subsequent instance of s.
func Join(s string, d ...Doc) Doc {
return JoinDoc(Concat(Text(s), Line), d...)
}
// JoinDoc joins Docs d with Doc s.
func JoinDoc(s Doc, d ...Doc) Doc {
switch len(d) {
case 0:
return Nil
case 1:
return d[0]
default:
return Fold(Concat, d[0], s, JoinDoc(s, d[1:]...))
}
}
// JoinNestedRight nests nested with string s.
// Every item after the first is indented.
// For example:
// aaaa
// <sep> bbb
// bbb
// <sep> ccc
// ccc
func JoinNestedRight(sep Doc, nested ...Doc) Doc {
switch len(nested) {
case 0:
return Nil
case 1:
return nested[0]
default:
return Concat(
nested[0],
FoldMap(Concat,
func(a Doc) Doc { return Concat(Line, ConcatSpace(sep, NestT(Group(a)))) },
nested[1:]...))
}
}
// ConcatDoc concatenates two Docs with between.
func ConcatDoc(a, b, between Doc) Doc {
return simplifyNil(a, b, func(a, b Doc) Doc {
return Concat(
a,
Concat(
between,
b,
),
)
})
}
// ConcatLine concatenates two Docs with a Line.
func ConcatLine(a, b Doc) Doc {
return ConcatDoc(a, b, Line)
}
// ConcatSpace concatenates two Docs with a space.
func ConcatSpace(a, b Doc) Doc {
return ConcatDoc(a, b, textSpace)
}
// Stack concats Docs with a Line between each.
func Stack(d ...Doc) Doc {
return Fold(ConcatLine, d...)
}
// Fillwords fills lines with as many docs as will fit joined with a space or line.
func Fillwords(d ...Doc) Doc {
u := &union{textSpace, line{}}
fill := func(a, b Doc) Doc {
return ConcatDoc(a, b, u)
}
return Fold(fill, d...)
}
// JoinGroupAligned nests joined d with divider under head.
func JoinGroupAligned(head, divider string, d ...Doc) Doc {
return AlignUnder(Text(head), Join(divider, d...))
}
// NestUnder nests nested under head.
func NestUnder(head, nested Doc) Doc {
return Group(Concat(
head,
NestT(Concat(
Line,
Group(nested),
)),
))
}
// AlignUnder aligns nested to the right of head, and, if
// this does not fit on the line, nests nested under head.
func AlignUnder(head, nested Doc) Doc {
g := Group(nested)
a := ConcatSpace(head, Align(g))
b := Concat(head, NestT(Concat(Line, g)))
return Group(&union{a, b})
}
// Fold applies f recursively to all Docs in d.
func Fold(f func(a, b Doc) Doc, d ...Doc) Doc {
switch len(d) {
case 0:
return Nil
case 1:
return d[0]
default:
return f(d[0], Fold(f, d[1:]...))
}
}
// FoldMap applies f recursively to all Docs in d processed through g.
func FoldMap(f func(a, b Doc) Doc, g func(Doc) Doc, d ...Doc) Doc {
switch len(d) {
case 0:
return Nil
case 1:
return g(d[0])
default:
return f(g(d[0]), FoldMap(f, g, d[1:]...))
}
}
// BracketDoc is like Bracket except it accepts Docs instead of strings.
func BracketDoc(l, x, r Doc) Doc {
return Group(Fold(Concat,
l,
NestT(Concat(SoftBreak, x)),
SoftBreak,
r,
))
}
// simplifyNil returns fn(a, b). If either Doc is Nil, the other Doc
// is returned without invoking fn.
func simplifyNil(a, b Doc, fn func(Doc, Doc) Doc) Doc {
if a == Nil {
return b
}
if b == Nil {
return a
}
return fn(a, b)
}
// JoinNestedOuter attempts to de-indent the items after the first so
// that the operator appears in the right margin. This replacement is
// only done if there are enough "simple spaces" (as per previous uses
// of Align) to de-indent. No attempt is made to de-indent hard tabs,
// otherwise alignment may break on output devices with a different
// physical tab width. docFn should be set to Text or Keyword and will be
// used when converting lbl to a Doc.
func JoinNestedOuter(lbl string, docFn func(string) Doc, d ...Doc) Doc {
sep := docFn(lbl)
return &snesting{
f: func(k int16) Doc {
if k < int16(len(lbl)+1) {
// If there is not enough space, don't even try. Just use a
// regular nested section.
return JoinNestedRight(sep, d...)
}
// If there is enough space, on the other hand, we can de-indent
// every line after the first.
return Concat(d[0],
NestS(int16(-len(lbl)-1),
FoldMap(
Concat,
func(x Doc) Doc {
return Concat(
Line,
ConcatSpace(sep, Align(Group(x))))
},
d[1:]...)))
},
}
}
// TableRow is the data for one row of a RLTable (see below).
type TableRow struct {
Label string
Doc Doc
}
// TableAlignment should be used as first argument to Table().
type TableAlignment int
const (
// TableNoAlign does not use alignment and instead uses NestUnder.
TableNoAlign TableAlignment = iota
// TableRightAlignFirstColumn right-aligns (left-pads) the first column.
TableRightAlignFirstColumn
// TableLeftAlignFirstColumn left-aligns (right-pads) the first column.
TableLeftAlignFirstColumn
)
// Table defines a document that formats a list of pairs of items either:
// - as a 2-column table, with the two columns aligned for example:
// SELECT aaa
// bbb
// FROM ccc
// - as sections, for example:
// SELECT
// aaa
// bbb
// FROM
// ccc
//
// We restrict the left value in each list item to be a one-line string
// to make the width computation efficient.
//
// For convenience, the function also skips over rows with a nil
// pointer as doc.
//
// docFn should be set to Text or Keyword and will be used when converting
// TableRow label's to Docs.
func Table(alignment TableAlignment, docFn func(string) Doc, rows ...TableRow) Doc {
// Compute the nested formatting in "sections". It's simple.
// Note that we do not use NestUnder() because we are not grouping
// at this level (the group is done for the final form below).
items := makeTableNestedSections(docFn, rows)
nestedSections := Stack(items...)
finalDoc := nestedSections
if alignment != TableNoAlign {
items = makeAlignedTableItems(alignment == TableRightAlignFirstColumn /* leftPad */, docFn, rows, items)
alignedTable := Stack(items...)
finalDoc = &union{alignedTable, nestedSections}
}
return Group(finalDoc)
}
func computeLeftColumnWidth(rows []TableRow) int {
leftwidth := 0
for _, r := range rows {
if r.Doc == nil {
continue
}
if leftwidth < len(r.Label) {
leftwidth = len(r.Label)
}
}
return leftwidth
}
func makeAlignedTableItems(
leftPad bool, docFn func(string) Doc, rows []TableRow, items []Doc,
) []Doc {
// We'll need the left column width below.
leftWidth := computeLeftColumnWidth(rows)
// Now convert the rows.
items = items[:0]
for _, r := range rows {
if r.Doc == nil || (r.Label == "" && r.Doc == Nil) {
continue
}
if r.Label != "" {
lbl := docFn(r.Label)
var pleft Doc
if leftPad {
pleft = Concat(pad{int16(leftWidth - len(r.Label))}, lbl)
} else {
pleft = Concat(lbl, pad{int16(leftWidth - len(r.Label))})
}
d := simplifyNil(pleft, r.Doc,
func(a, b Doc) Doc {
return ConcatSpace(a, Align(Group(b)))
})
items = append(items, d)
} else {
items = append(items, Concat(pad{int16(leftWidth + 1)}, Align(Group(r.Doc))))
}
}
return items
}
func makeTableNestedSections(docFn func(string) Doc, rows []TableRow) []Doc {
items := make([]Doc, 0, len(rows))
for _, r := range rows {
if r.Doc == nil || (r.Label == "" && r.Doc == Nil) {
continue
}
if r.Label != "" {
d := simplifyNil(docFn(r.Label), r.Doc,
func(a, b Doc) Doc {
return Concat(a, NestT(Concat(Line, Group(b))))
})
items = append(items, d)
} else {
items = append(items, Group(r.Doc))
}
}
return items
}