From 83c9b8340540a77385f2f7b40cec50bb583053dd Mon Sep 17 00:00:00 2001
From: Keith Randall <khr@golang.org>
Date: Wed, 8 Jun 2022 13:30:02 -0700
Subject: [PATCH] math/big: make NewInt inlineable and zero allocation

Mark the assembly routines as not escaping their arguments.

Add a special case to NewInt that, when inlined, can do all
of its allocations (a big.Int and a [1]Word) on the stack.

Update #29951

Change-Id: I9bd38c262eb97df98c0ed9874da7daac381243ea
Reviewed-on: https://go-review.googlesource.com/c/go/+/411254
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Keith Randall <khr@google.com>
Run-TryBot: Keith Randall <khr@golang.org>
Reviewed-by: Robert Griesemer <gri@google.com>
---
 src/math/big/arith_decl.go | 16 ++++++++++++++++
 src/math/big/int.go        | 15 ++++++++++++++-
 src/math/big/int_test.go   | 26 +++++++++++++++++++++++++-
 3 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/src/math/big/arith_decl.go b/src/math/big/arith_decl.go
index 301aa55f1ab6f..9b254f22130fc 100644
--- a/src/math/big/arith_decl.go
+++ b/src/math/big/arith_decl.go
@@ -8,11 +8,27 @@
 package big
 
 // implemented in arith_$GOARCH.s
+
+//go:noescape
 func addVV(z, x, y []Word) (c Word)
+
+//go:noescape
 func subVV(z, x, y []Word) (c Word)
+
+//go:noescape
 func addVW(z, x []Word, y Word) (c Word)
+
+//go:noescape
 func subVW(z, x []Word, y Word) (c Word)
+
+//go:noescape
 func shlVU(z, x []Word, s uint) (c Word)
+
+//go:noescape
 func shrVU(z, x []Word, s uint) (c Word)
+
+//go:noescape
 func mulAddVWW(z, x []Word, y, r Word) (c Word)
+
+//go:noescape
 func addMulVVW(z, x []Word, y Word) (c Word)
diff --git a/src/math/big/int.go b/src/math/big/int.go
index ec168f8ffe571..ca4c3561e629f 100644
--- a/src/math/big/int.go
+++ b/src/math/big/int.go
@@ -65,7 +65,20 @@ func (z *Int) SetUint64(x uint64) *Int {
 
 // NewInt allocates and returns a new Int set to x.
 func NewInt(x int64) *Int {
-	return new(Int).SetInt64(x)
+	// This code is arranged to be inlineable and produce
+	// zero allocations when inlined. See issue 29951.
+	u := uint64(x)
+	if x < 0 {
+		u = -u
+	}
+	var abs []Word
+	if x == 0 {
+	} else if _W == 32 && u>>32 != 0 {
+		abs = []Word{Word(u), Word(u >> 32)}
+	} else {
+		abs = []Word{Word(u)}
+	}
+	return &Int{neg: x < 0, abs: abs}
 }
 
 // Set sets z to x and returns z.
diff --git a/src/math/big/int_test.go b/src/math/big/int_test.go
index 3c8557323a032..4ebb09d2c2aff 100644
--- a/src/math/big/int_test.go
+++ b/src/math/big/int_test.go
@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"encoding/hex"
 	"fmt"
+	"math"
 	"math/rand"
 	"strconv"
 	"strings"
@@ -92,7 +93,7 @@ func testFunZZ(t *testing.T, msg string, f funZZ, a argZZ) {
 		t.Errorf("%s%v is not normalized", msg, z)
 	}
 	if (&z).Cmp(a.z) != 0 {
-		t.Errorf("%s%+v\n\tgot z = %v; want %v", msg, a, &z, a.z)
+		t.Errorf("%v %s %v\n\tgot z = %v; want %v", a.x, msg, a.y, &z, a.z)
 	}
 }
 
@@ -1894,3 +1895,26 @@ func TestFillBytes(t *testing.T) {
 		})
 	}
 }
+
+func TestNewIntMinInt64(t *testing.T) {
+	// Test for uint64 cast in NewInt.
+	want := int64(math.MinInt64)
+	if got := NewInt(want).Int64(); got != want {
+		t.Fatalf("wanted %d, got %d", want, got)
+	}
+}
+
+func TestNewIntAllocs(t *testing.T) {
+	for _, n := range []int64{0, 7, -7, 1 << 30, -1 << 30, 1 << 50, -1 << 50} {
+		x := NewInt(3)
+		got := testing.AllocsPerRun(100, func() {
+			// NewInt should inline, and all its allocations
+			// can happen on the stack. Passing the result of NewInt
+			// to Add should not cause any of those allocations to escape.
+			x.Add(x, NewInt(n))
+		})
+		if got != 0 {
+			t.Errorf("x.Add(x, NewInt(%d)), wanted 0 allocations, got %f", n, got)
+		}
+	}
+}