encode.mx raw
1 // Copyright 2021 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 encodecounter
6
7 import (
8 "bufio"
9 "encoding/binary"
10 "fmt"
11 "internal/coverage"
12 "internal/coverage/slicewriter"
13 "internal/coverage/stringtab"
14 "internal/coverage/uleb128"
15 "io"
16 "maps"
17 "os"
18 "slices"
19 )
20
21 // This package contains APIs and helpers for encoding initial portions
22 // of the counter data files emitted at runtime when coverage instrumentation
23 // is enabled. Counter data files may contain multiple segments; the file
24 // header and first segment are written via the "Write" method below, and
25 // additional segments can then be added using "AddSegment".
26
27 type CoverageDataWriter struct {
28 stab *stringtab.Writer
29 w *bufio.Writer
30 csh coverage.CounterSegmentHeader
31 tmp []byte
32 cflavor coverage.CounterFlavor
33 segs uint32
34 debug bool
35 }
36
37 func NewCoverageDataWriter(w io.Writer, flav coverage.CounterFlavor) *CoverageDataWriter {
38 r := &CoverageDataWriter{
39 stab: &stringtab.Writer{},
40 w: bufio.NewWriter(w),
41
42 tmp: []byte{:64},
43 cflavor: flav,
44 }
45 r.stab.InitWriter()
46 r.stab.Lookup("")
47 return r
48 }
49
50 // CounterVisitor describes a helper object used during counter file
51 // writing; when writing counter data files, clients pass a
52 // CounterVisitor to the write/emit routines, then the expectation is
53 // that the VisitFuncs method will then invoke the callback "f" with
54 // data for each function to emit to the file.
55 type CounterVisitor interface {
56 VisitFuncs(f CounterVisitorFn) error
57 }
58
59 // CounterVisitorFn describes a callback function invoked when writing
60 // coverage counter data.
61 type CounterVisitorFn func(pkid uint32, funcid uint32, counters []uint32) error
62
63 // Write writes the contents of the count-data file to the writer
64 // previously supplied to NewCoverageDataWriter. Returns an error
65 // if something went wrong somewhere with the write.
66 func (cfw *CoverageDataWriter) Write(metaFileHash [16]byte, args map[string][]byte, visitor CounterVisitor) error {
67 if err := cfw.writeHeader(metaFileHash); err != nil {
68 return err
69 }
70 return cfw.AppendSegment(args, visitor)
71 }
72
73 func padToFourByteBoundary(ws *slicewriter.WriteSeeker) error {
74 sz := len(ws.BytesWritten())
75 zeros := []byte{0, 0, 0, 0}
76 rem := uint32(sz) % 4
77 if rem != 0 {
78 pad := zeros[:(4 - rem)]
79 if nw, err := ws.Write(pad); err != nil {
80 return err
81 } else if nw != len(pad) {
82 return fmt.Errorf("error: short write")
83 }
84 }
85 return nil
86 }
87
88 func (cfw *CoverageDataWriter) patchSegmentHeader(ws *slicewriter.WriteSeeker) error {
89 // record position
90 off, err := ws.Seek(0, io.SeekCurrent)
91 if err != nil {
92 return fmt.Errorf("error seeking in patchSegmentHeader: %v", err)
93 }
94 // seek back to start so that we can update the segment header
95 if _, err := ws.Seek(0, io.SeekStart); err != nil {
96 return fmt.Errorf("error seeking in patchSegmentHeader: %v", err)
97 }
98 if cfw.debug {
99 fmt.Fprintf(os.Stderr, "=-= writing counter segment header: %+v", cfw.csh)
100 }
101 if err := binary.Write(ws, binary.LittleEndian, cfw.csh); err != nil {
102 return err
103 }
104 // ... and finally return to the original offset.
105 if _, err := ws.Seek(off, io.SeekStart); err != nil {
106 return fmt.Errorf("error seeking in patchSegmentHeader: %v", err)
107 }
108 return nil
109 }
110
111 func (cfw *CoverageDataWriter) writeSegmentPreamble(args map[string][]byte, ws *slicewriter.WriteSeeker) error {
112 if err := binary.Write(ws, binary.LittleEndian, cfw.csh); err != nil {
113 return err
114 }
115 hdrsz := uint32(len(ws.BytesWritten()))
116
117 // Write string table and args to a byte slice (since we need
118 // to capture offsets at various points), then emit the slice
119 // once we are done.
120 cfw.stab.Freeze()
121 if err := cfw.stab.Write(ws); err != nil {
122 return err
123 }
124 cfw.csh.StrTabLen = uint32(len(ws.BytesWritten())) - hdrsz
125
126 akeys := slices.Sorted(maps.Keys(args))
127
128 wrULEB128 := func(v uint) error {
129 cfw.tmp = cfw.tmp[:0]
130 cfw.tmp = uleb128.AppendUleb128(cfw.tmp, v)
131 if _, err := ws.Write(cfw.tmp); err != nil {
132 return err
133 }
134 return nil
135 }
136
137 // Count of arg pairs.
138 if err := wrULEB128(uint(len(args))); err != nil {
139 return err
140 }
141 // Arg pairs themselves.
142 for _, k := range akeys {
143 ki := uint(cfw.stab.Lookup(k))
144 if err := wrULEB128(ki); err != nil {
145 return err
146 }
147 v := args[k]
148 vi := uint(cfw.stab.Lookup(v))
149 if err := wrULEB128(vi); err != nil {
150 return err
151 }
152 }
153 if err := padToFourByteBoundary(ws); err != nil {
154 return err
155 }
156 cfw.csh.ArgsLen = uint32(len(ws.BytesWritten())) - (cfw.csh.StrTabLen + hdrsz)
157
158 return nil
159 }
160
161 // AppendSegment appends a new segment to a counter data, with a new
162 // args section followed by a payload of counter data clauses.
163 func (cfw *CoverageDataWriter) AppendSegment(args map[string][]byte, visitor CounterVisitor) error {
164 cfw.stab = &stringtab.Writer{}
165 cfw.stab.InitWriter()
166 cfw.stab.Lookup("")
167
168 var err error
169 for k, v := range args {
170 cfw.stab.Lookup(k)
171 cfw.stab.Lookup(v)
172 }
173
174 ws := &slicewriter.WriteSeeker{}
175 if err = cfw.writeSegmentPreamble(args, ws); err != nil {
176 return err
177 }
178 if err = cfw.writeCounters(visitor, ws); err != nil {
179 return err
180 }
181 if err = cfw.patchSegmentHeader(ws); err != nil {
182 return err
183 }
184 if err := cfw.writeBytes(ws.BytesWritten()); err != nil {
185 return err
186 }
187 if err = cfw.writeFooter(); err != nil {
188 return err
189 }
190 if err := cfw.w.Flush(); err != nil {
191 return fmt.Errorf("write error: %v", err)
192 }
193 cfw.stab = nil
194 return nil
195 }
196
197 func (cfw *CoverageDataWriter) writeHeader(metaFileHash [16]byte) error {
198 // Emit file header.
199 ch := coverage.CounterFileHeader{
200 Magic: coverage.CovCounterMagic,
201 Version: coverage.CounterFileVersion,
202 MetaHash: metaFileHash,
203 CFlavor: cfw.cflavor,
204 BigEndian: false,
205 }
206 if err := binary.Write(cfw.w, binary.LittleEndian, ch); err != nil {
207 return err
208 }
209 return nil
210 }
211
212 func (cfw *CoverageDataWriter) writeBytes(b []byte) error {
213 if len(b) == 0 {
214 return nil
215 }
216 nw, err := cfw.w.Write(b)
217 if err != nil {
218 return fmt.Errorf("error writing counter data: %v", err)
219 }
220 if len(b) != nw {
221 return fmt.Errorf("error writing counter data: short write")
222 }
223 return nil
224 }
225
226 func (cfw *CoverageDataWriter) writeCounters(visitor CounterVisitor, ws *slicewriter.WriteSeeker) error {
227 // Notes:
228 // - this version writes everything little-endian, which means
229 // a call is needed to encode every value (expensive)
230 // - we may want to move to a model in which we just blast out
231 // all counters, or possibly mmap the file and do the write
232 // implicitly.
233 ctrb := []byte{:4}
234 wrval := func(val uint32) error {
235 var buf []byte
236 var towr int
237 if cfw.cflavor == coverage.CtrRaw {
238 binary.LittleEndian.PutUint32(ctrb, val)
239 buf = ctrb
240 towr = 4
241 } else if cfw.cflavor == coverage.CtrULeb128 {
242 cfw.tmp = cfw.tmp[:0]
243 cfw.tmp = uleb128.AppendUleb128(cfw.tmp, uint(val))
244 buf = cfw.tmp
245 towr = len(buf)
246 } else {
247 panic("internal error: bad counter flavor")
248 }
249 if sz, err := ws.Write(buf); err != nil {
250 return err
251 } else if sz != towr {
252 return fmt.Errorf("writing counters: short write")
253 }
254 return nil
255 }
256
257 // Write out entries for each live function.
258 emitter := func(pkid uint32, funcid uint32, counters []uint32) error {
259 cfw.csh.FcnEntries++
260 if err := wrval(uint32(len(counters))); err != nil {
261 return err
262 }
263
264 if err := wrval(pkid); err != nil {
265 return err
266 }
267
268 if err := wrval(funcid); err != nil {
269 return err
270 }
271 for _, val := range counters {
272 if err := wrval(val); err != nil {
273 return err
274 }
275 }
276 return nil
277 }
278 if err := visitor.VisitFuncs(emitter); err != nil {
279 return err
280 }
281 return nil
282 }
283
284 func (cfw *CoverageDataWriter) writeFooter() error {
285 cfw.segs++
286 cf := coverage.CounterFileFooter{
287 Magic: coverage.CovCounterMagic,
288 NumSegments: cfw.segs,
289 }
290 if err := binary.Write(cfw.w, binary.LittleEndian, cf); err != nil {
291 return err
292 }
293 return nil
294 }
295