1// Copyright 2022 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
5package stringtab
6
7import (
8	"fmt"
9	"internal/coverage/slicereader"
10	"internal/coverage/uleb128"
11	"io"
12)
13
14// This package implements string table writer and reader utilities,
15// for use in emitting and reading/decoding coverage meta-data and
16// counter-data files.
17
18// Writer implements a string table writing utility.
19type Writer struct {
20	stab   map[string]uint32
21	strs   []string
22	tmp    []byte
23	frozen bool
24}
25
26// InitWriter initializes a stringtab.Writer.
27func (stw *Writer) InitWriter() {
28	stw.stab = make(map[string]uint32)
29	stw.tmp = make([]byte, 64)
30}
31
32// Nentries returns the number of strings interned so far.
33func (stw *Writer) Nentries() uint32 {
34	return uint32(len(stw.strs))
35}
36
37// Lookup looks up string 's' in the writer's table, adding
38// a new entry if need be, and returning an index into the table.
39func (stw *Writer) Lookup(s string) uint32 {
40	if idx, ok := stw.stab[s]; ok {
41		return idx
42	}
43	if stw.frozen {
44		panic("internal error: string table previously frozen")
45	}
46	idx := uint32(len(stw.strs))
47	stw.stab[s] = idx
48	stw.strs = append(stw.strs, s)
49	return idx
50}
51
52// Size computes the memory in bytes needed for the serialized
53// version of a stringtab.Writer.
54func (stw *Writer) Size() uint32 {
55	rval := uint32(0)
56	stw.tmp = stw.tmp[:0]
57	stw.tmp = uleb128.AppendUleb128(stw.tmp, uint(len(stw.strs)))
58	rval += uint32(len(stw.tmp))
59	for _, s := range stw.strs {
60		stw.tmp = stw.tmp[:0]
61		slen := uint(len(s))
62		stw.tmp = uleb128.AppendUleb128(stw.tmp, slen)
63		rval += uint32(len(stw.tmp)) + uint32(slen)
64	}
65	return rval
66}
67
68// Write writes the string table in serialized form to the specified
69// io.Writer.
70func (stw *Writer) Write(w io.Writer) error {
71	wr128 := func(v uint) error {
72		stw.tmp = stw.tmp[:0]
73		stw.tmp = uleb128.AppendUleb128(stw.tmp, v)
74		if nw, err := w.Write(stw.tmp); err != nil {
75			return fmt.Errorf("writing string table: %v", err)
76		} else if nw != len(stw.tmp) {
77			return fmt.Errorf("short write emitting stringtab uleb")
78		}
79		return nil
80	}
81	if err := wr128(uint(len(stw.strs))); err != nil {
82		return err
83	}
84	for _, s := range stw.strs {
85		if err := wr128(uint(len(s))); err != nil {
86			return err
87		}
88		if nw, err := w.Write([]byte(s)); err != nil {
89			return fmt.Errorf("writing string table: %v", err)
90		} else if nw != len([]byte(s)) {
91			return fmt.Errorf("short write emitting stringtab")
92		}
93	}
94	return nil
95}
96
97// Freeze sends a signal to the writer that no more additions are
98// allowed, only lookups of existing strings (if a lookup triggers
99// addition, a panic will result). Useful as a mechanism for
100// "finalizing" a string table prior to writing it out.
101func (stw *Writer) Freeze() {
102	stw.frozen = true
103}
104
105// Reader is a helper for reading a string table previously
106// serialized by a Writer.Write call.
107type Reader struct {
108	r    *slicereader.Reader
109	strs []string
110}
111
112// NewReader creates a stringtab.Reader to read the contents
113// of a string table from 'r'.
114func NewReader(r *slicereader.Reader) *Reader {
115	str := &Reader{
116		r: r,
117	}
118	return str
119}
120
121// Read reads/decodes a string table using the reader provided.
122func (str *Reader) Read() {
123	numEntries := int(str.r.ReadULEB128())
124	str.strs = make([]string, 0, numEntries)
125	for idx := 0; idx < numEntries; idx++ {
126		slen := str.r.ReadULEB128()
127		str.strs = append(str.strs, str.r.ReadString(int64(slen)))
128	}
129}
130
131// Entries returns the number of decoded entries in a string table.
132func (str *Reader) Entries() int {
133	return len(str.strs)
134}
135
136// Get returns string 'idx' within the string table.
137func (str *Reader) Get(idx uint32) string {
138	return str.strs[idx]
139}
140