1// Copyright 2024 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 pgo
6
7import (
8	"bytes"
9	"encoding/binary"
10	"fmt"
11	"reflect"
12	"strings"
13	"testing"
14)
15
16// equal returns an error if got and want are not equal.
17func equal(got, want *Profile) error {
18	if got.TotalWeight != want.TotalWeight {
19		return fmt.Errorf("got.TotalWeight %d != want.TotalWeight %d", got.TotalWeight, want.TotalWeight)
20	}
21	if !reflect.DeepEqual(got.NamedEdgeMap.ByWeight, want.NamedEdgeMap.ByWeight) {
22		return fmt.Errorf("got.NamedEdgeMap.ByWeight != want.NamedEdgeMap.ByWeight\ngot = %+v\nwant = %+v", got.NamedEdgeMap.ByWeight, want.NamedEdgeMap.ByWeight)
23	}
24	if !reflect.DeepEqual(got.NamedEdgeMap.Weight, want.NamedEdgeMap.Weight) {
25		return fmt.Errorf("got.NamedEdgeMap.Weight != want.NamedEdgeMap.Weight\ngot = %+v\nwant = %+v", got.NamedEdgeMap.Weight, want.NamedEdgeMap.Weight)
26	}
27
28	return nil
29}
30
31func testRoundTrip(t *testing.T, d *Profile) []byte {
32	var buf bytes.Buffer
33	n, err := d.WriteTo(&buf)
34	if err != nil {
35		t.Fatalf("WriteTo got err %v want nil", err)
36	}
37	if n != int64(buf.Len()) {
38		t.Errorf("WriteTo got n %d want %d", n, int64(buf.Len()))
39	}
40
41	b := buf.Bytes()
42
43	got, err := FromSerialized(&buf)
44	if err != nil {
45		t.Fatalf("processSerialized got err %v want nil", err)
46	}
47	if err := equal(got, d); err != nil {
48		t.Errorf("processSerialized output does not match input: %v", err)
49	}
50
51	return b
52}
53
54func TestEmpty(t *testing.T) {
55	d := emptyProfile()
56	b := testRoundTrip(t, d)
57
58	// Contents should consist of only a header.
59	if string(b) != serializationHeader {
60		t.Errorf("WriteTo got %q want %q", string(b), serializationHeader)
61	}
62}
63
64func TestRoundTrip(t *testing.T) {
65	d := &Profile{
66		TotalWeight: 3,
67		NamedEdgeMap: NamedEdgeMap{
68			ByWeight: []NamedCallEdge{
69				{
70					CallerName: "a",
71					CalleeName: "b",
72					CallSiteOffset: 14,
73				},
74				{
75					CallerName: "c",
76					CalleeName: "d",
77					CallSiteOffset: 15,
78				},
79			},
80			Weight: map[NamedCallEdge]int64{
81				{
82					CallerName: "a",
83					CalleeName: "b",
84					CallSiteOffset: 14,
85				}: 2,
86				{
87					CallerName: "c",
88					CalleeName: "d",
89					CallSiteOffset: 15,
90				}: 1,
91			},
92		},
93	}
94
95	testRoundTrip(t, d)
96}
97
98func constructFuzzProfile(t *testing.T, b []byte) *Profile {
99	// The fuzzer can't construct an arbitrary structure, so instead we
100	// consume bytes from b to act as our edge data.
101	r := bytes.NewReader(b)
102	consumeString := func() (string, bool) {
103		// First byte: how many bytes to read for this string? We only
104		// use a byte to avoid making humongous strings.
105		length, err := r.ReadByte()
106		if err != nil {
107			return "", false
108		}
109		if length == 0 {
110			return "", false
111		}
112
113		b := make([]byte, length)
114		_, err = r.Read(b)
115		if err != nil {
116			return "", false
117		}
118
119		return string(b), true
120	}
121	consumeInt64 := func() (int64, bool) {
122		b := make([]byte, 8)
123		_, err := r.Read(b)
124		if err != nil {
125			return 0, false
126		}
127
128		return int64(binary.LittleEndian.Uint64(b)), true
129	}
130
131	d := emptyProfile()
132
133	for {
134		caller, ok := consumeString()
135		if !ok {
136			break
137		}
138		if strings.ContainsAny(caller, " \r\n") {
139			t.Skip("caller contains space or newline")
140		}
141
142		callee, ok := consumeString()
143		if !ok {
144			break
145		}
146		if strings.ContainsAny(callee, " \r\n") {
147			t.Skip("callee contains space or newline")
148		}
149
150		line, ok := consumeInt64()
151		if !ok {
152			break
153		}
154		weight, ok := consumeInt64()
155		if !ok {
156			break
157		}
158
159		edge := NamedCallEdge{
160			CallerName: caller,
161			CalleeName: callee,
162			CallSiteOffset: int(line),
163		}
164
165		if _, ok := d.NamedEdgeMap.Weight[edge]; ok {
166			t.Skip("duplicate edge")
167		}
168
169		d.NamedEdgeMap.Weight[edge] = weight
170		d.TotalWeight += weight
171	}
172
173	byWeight := make([]NamedCallEdge, 0, len(d.NamedEdgeMap.Weight))
174	for namedEdge := range d.NamedEdgeMap.Weight {
175		byWeight = append(byWeight, namedEdge)
176	}
177	sortByWeight(byWeight, d.NamedEdgeMap.Weight)
178	d.NamedEdgeMap.ByWeight = byWeight
179
180	return d
181}
182
183func FuzzRoundTrip(f *testing.F) {
184	f.Add([]byte("")) // empty profile
185
186	f.Fuzz(func(t *testing.T, b []byte) {
187		d := constructFuzzProfile(t, b)
188		testRoundTrip(t, d)
189	})
190}
191