From dabbf9f88c560286e150e9b136a27c3ac23c5ec1 Mon Sep 17 00:00:00 2001
From: Frederic Branczyk <fbranczyk@gmail.com>
Date: Fri, 3 Feb 2023 16:05:56 +0100
Subject: [PATCH] cmd/compile/internal/pgo: fix hard-coded PGO sample data
 position

This patch detects at which index position profiling samples that have
the value-type samples count are, instead of the previously hard-coded
position of index 1. Runtime generated profiles always generate CPU
profiling data with the 0 index being CPU nanoseconds, and samples count
at index 1, which is why this previously hasn't come up.

Fixes #58292
---
 src/cmd/compile/internal/pgo/irgraph.go       | 18 ++++-
 src/cmd/compile/internal/test/pgo_inl_test.go | 68 +++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/src/cmd/compile/internal/pgo/irgraph.go b/src/cmd/compile/internal/pgo/irgraph.go
index ca9e2f3b5adde3..a0319f3962b2f3 100644
--- a/src/cmd/compile/internal/pgo/irgraph.go
+++ b/src/cmd/compile/internal/pgo/irgraph.go
@@ -140,9 +140,25 @@ func New(profileFile string) *Profile {
 		return nil
 	}
 
+	samplesCountIndex := -1
+	for i, s := range profile.SampleType {
+		// Samples count is the raw data collected, and CPU nanoseconds is just
+		// a scaled version of it, so either one we can find is fine.
+		if (s.Type == "samples" && s.Unit == "count") ||
+			(s.Type == "cpu" && s.Unit == "nanoseconds") {
+			samplesCountIndex = i
+			break
+		}
+	}
+
+	if samplesCountIndex == -1 {
+		log.Fatal("failed to find CPU samples count or CPU nanoseconds value-types in profile.")
+		return nil
+	}
+
 	g := newGraph(profile, &Options{
 		CallTree:    false,
-		SampleValue: func(v []int64) int64 { return v[1] },
+		SampleValue: func(v []int64) int64 { return v[samplesCountIndex] },
 	})
 
 	p := &Profile{
diff --git a/src/cmd/compile/internal/test/pgo_inl_test.go b/src/cmd/compile/internal/test/pgo_inl_test.go
index 2f6391fded265d..4d6b5a134a0e28 100644
--- a/src/cmd/compile/internal/test/pgo_inl_test.go
+++ b/src/cmd/compile/internal/test/pgo_inl_test.go
@@ -7,6 +7,7 @@ package test
 import (
 	"bufio"
 	"fmt"
+	"internal/profile"
 	"internal/testenv"
 	"io"
 	"os"
@@ -213,6 +214,73 @@ func TestPGOIntendedInliningShiftedLines(t *testing.T) {
 	testPGOIntendedInlining(t, dir)
 }
 
+// TestPGOSingleIndex tests that the sample index can not be 1 and compilation
+// will not fail. All it should care about is that the sample type is either
+// CPU nanoseconds or samples count, whichever it finds first.
+func TestPGOSingleIndex(t *testing.T) {
+	for _, tc := range []struct {
+		originalIndex int
+	}{{
+		// The `testdata/pgo/inline/inline_hot.pprof` file is a standard CPU
+		// profile as the runtime would generate. The 0 index contains the
+		// value-type samples and value-unit count. The 1 index contains the
+		// value-type cpu and value-unit nanoseconds. These tests ensure that
+		// the compiler can work with profiles that only have a single index,
+		// but are either samples count or CPU nanoseconds.
+		originalIndex: 0,
+	}, {
+		originalIndex: 1,
+	}} {
+		t.Run(fmt.Sprintf("originalIndex=%d", tc.originalIndex), func(t *testing.T) {
+			wd, err := os.Getwd()
+			if err != nil {
+				t.Fatalf("error getting wd: %v", err)
+			}
+			srcDir := filepath.Join(wd, "testdata/pgo/inline")
+
+			// Copy the module to a scratch location so we can add a go.mod.
+			dir := t.TempDir()
+
+			originalPprofFile, err := os.Open(filepath.Join(srcDir, "inline_hot.pprof"))
+			if err != nil {
+				t.Fatalf("error opening inline_hot.pprof: %v", err)
+			}
+			defer originalPprofFile.Close()
+
+			p, err := profile.Parse(originalPprofFile)
+			if err != nil {
+				t.Fatalf("error parsing inline_hot.pprof: %v", err)
+			}
+
+			// Move the samples count value-type to the 0 index.
+			p.SampleType = []*profile.ValueType{p.SampleType[tc.originalIndex]}
+
+			// Ensure we only have a single set of sample values.
+			for _, s := range p.Sample {
+				s.Value = []int64{s.Value[tc.originalIndex]}
+			}
+
+			modifiedPprofFile, err := os.Create(filepath.Join(dir, "inline_hot.pprof"))
+			if err != nil {
+				t.Fatalf("error creating inline_hot.pprof: %v", err)
+			}
+			defer modifiedPprofFile.Close()
+
+			if err := p.Write(modifiedPprofFile); err != nil {
+				t.Fatalf("error writing inline_hot.pprof: %v", err)
+			}
+
+			for _, file := range []string{"inline_hot.go", "inline_hot_test.go"} {
+				if err := copyFile(filepath.Join(dir, file), filepath.Join(srcDir, file)); err != nil {
+					t.Fatalf("error copying %s: %v", file, err)
+				}
+			}
+
+			testPGOIntendedInlining(t, dir)
+		})
+	}
+}
+
 func copyFile(dst, src string) error {
 	s, err := os.Open(src)
 	if err != nil {