From c4ec74a5c9f258918d2266ea281f3e344c6ea35d Mon Sep 17 00:00:00 2001
From: David Chase <drchase@google.com>
Date: Thu, 11 Aug 2022 13:48:54 -0400
Subject: [PATCH] go/internal/pkgbits: fix performance of rawReloc

Tracking changes in go repo for unified IR.

See CL 422297.

Change-Id: If05aba7693c70cac4969721e5dd8a5bf197670b4
Reviewed-on: https://go-review.googlesource.com/c/tools/+/422899
Reviewed-by: Matthew Dempsky <mdempsky@google.com>
Run-TryBot: David Chase <drchase@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
---
 go/internal/pkgbits/encoder.go | 18 +++++++++++-------
 go/internal/pkgbits/reloc.go   |  4 ++--
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/go/internal/pkgbits/encoder.go b/go/internal/pkgbits/encoder.go
index c50c838caa..e98e41171e 100644
--- a/go/internal/pkgbits/encoder.go
+++ b/go/internal/pkgbits/encoder.go
@@ -147,8 +147,9 @@ func (pw *PkgEncoder) NewEncoderRaw(k RelocKind) Encoder {
 type Encoder struct {
 	p *PkgEncoder
 
-	Relocs []RelocEnt
-	Data   bytes.Buffer // accumulated element bitstream data
+	Relocs   []RelocEnt
+	RelocMap map[RelocEnt]uint32
+	Data     bytes.Buffer // accumulated element bitstream data
 
 	encodingRelocHeader bool
 
@@ -210,15 +211,18 @@ func (w *Encoder) rawVarint(x int64) {
 }
 
 func (w *Encoder) rawReloc(r RelocKind, idx Index) int {
-	// TODO(mdempsky): Use map for lookup; this takes quadratic time.
-	for i, rEnt := range w.Relocs {
-		if rEnt.Kind == r && rEnt.Idx == idx {
-			return i
+	e := RelocEnt{r, idx}
+	if w.RelocMap != nil {
+		if i, ok := w.RelocMap[e]; ok {
+			return int(i)
 		}
+	} else {
+		w.RelocMap = make(map[RelocEnt]uint32)
 	}
 
 	i := len(w.Relocs)
-	w.Relocs = append(w.Relocs, RelocEnt{r, idx})
+	w.RelocMap[e] = uint32(i)
+	w.Relocs = append(w.Relocs, e)
 	return i
 }
 
diff --git a/go/internal/pkgbits/reloc.go b/go/internal/pkgbits/reloc.go
index 7a8f04ab3f..fcdfb97ca9 100644
--- a/go/internal/pkgbits/reloc.go
+++ b/go/internal/pkgbits/reloc.go
@@ -5,11 +5,11 @@
 package pkgbits
 
 // A RelocKind indicates a particular section within a unified IR export.
-type RelocKind int
+type RelocKind int32
 
 // An Index represents a bitstream element index within a particular
 // section.
-type Index int
+type Index int32
 
 // A relocEnt (relocation entry) is an entry in an element's local
 // reference table.