uint256: optimize WriteToArray (#190)

This change optimizes WriteToArray, and also implements PutUint256 for use in writing to "sufficiently large slices". --------- Co-authored-by: Martin Holst Swende <[email protected]>
holiman · Nov 29, 2024 · 439fbd4 · 439fbd4
1 parent 555918b
commit 439fbd4
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 13 deletions.
diff --git a/benchmarks_test.go b/benchmarks_test.go
@@ -1020,3 +1020,38 @@ func BenchmarkExtendSign(b *testing.B) {
 		result.ExtendSign(a, n)
 	}
 }
+
+func BenchmarkWriteTo(b *testing.B) {
+	fa, err := FromHex("0x1100030405060708090a0b0c0d0ed1e870eec79504c60144cc7f5fc2bad1e611")
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.Run("fixed-20", func(b *testing.B) {
+		dest := [20]byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
+		for i := 0; i < b.N; i++ {
+			fa.WriteToArray20(&dest)
+		}
+		_ = (string(dest[:])) // Prevent the compiler from optimizing away the op
+	})
+	b.Run("fixed-32", func(b *testing.B) {
+		dest := [32]byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
+		for i := 0; i < b.N; i++ {
+			fa.WriteToArray32(&dest)
+		}
+		_ = (string(dest[:])) // Prevent the compiler from optimizing away the op
+	})
+	b.Run("slice", func(b *testing.B) {
+		dest := make([]byte, 64)
+		for i := 0; i < b.N; i++ {
+			fa.WriteToSlice(dest)
+		}
+		_ = (string(dest[:])) // Prevent the compiler from optimizing away the op
+	})
+	b.Run("put256", func(b *testing.B) {
+		dest := make([]byte, 64)
+		for i := 0; i < b.N; i++ {
+			fa.PutUint256(dest)
+		}
+		_ = (string(dest[:])) // Prevent the compiler from optimizing away the op
+	})
+}
diff --git a/uint256.go b/uint256.go
@@ -144,18 +144,35 @@ func (z *Int) WriteToSlice(dest []byte) {
 	}
 }
 
+// PutUint256 writes all 32 bytes of z to the destination slice, including zero-bytes.
+// If dest is larger than 32 bytes, z will fill the first parts, and leave
+// the end untouched.
+// Note: The dest slice must be at least 32 bytes large, otherwise this
+// method will panic. The method WriteToSlice, which is slower,  should be used
+// if the destination slice is smaller or of unknown size.
+func (z *Int) PutUint256(dest []byte) {
+	_ = dest[31]
+	binary.BigEndian.PutUint64(dest[0:8], z[3])
+	binary.BigEndian.PutUint64(dest[8:16], z[2])
+	binary.BigEndian.PutUint64(dest[16:24], z[1])
+	binary.BigEndian.PutUint64(dest[24:32], z[0])
+}
+
 // WriteToArray32 writes all 32 bytes of z to the destination array, including zero-bytes
 func (z *Int) WriteToArray32(dest *[32]byte) {
-	for i := 0; i < 32; i++ {
-		dest[31-i] = byte(z[i/8] >> uint64(8*(i%8)))
-	}
+	// The PutUint64()s are inlined and we get 4x (load, bswap, store) instructions.
+	binary.BigEndian.PutUint64(dest[0:8], z[3])
+	binary.BigEndian.PutUint64(dest[8:16], z[2])
+	binary.BigEndian.PutUint64(dest[16:24], z[1])
+	binary.BigEndian.PutUint64(dest[24:32], z[0])
 }
 
 // WriteToArray20 writes the last 20 bytes of z to the destination array, including zero-bytes
 func (z *Int) WriteToArray20(dest *[20]byte) {
-	for i := 0; i < 20; i++ {
-		dest[19-i] = byte(z[i/8] >> uint64(8*(i%8)))
-	}
+	// The PutUint*()s are inlined and we get 3x (load, bswap, store) instructions.
+	binary.BigEndian.PutUint32(dest[0:4], uint32(z[2]))
+	binary.BigEndian.PutUint64(dest[4:12], z[1])
+	binary.BigEndian.PutUint64(dest[12:20], z[0])
 }
 
 // Uint64 returns the lower 64-bits of z

diff --git a/uint256_test.go b/uint256_test.go
@@ -479,7 +479,7 @@ func TestUdivremQuick(t *testing.T) {
 	var (
 		u        = []uint64{1, 0, 0, 0, 0}
 		expected = new(Int)
-		rem Int
+		rem      Int
 	)
 	udivrem([]uint64{}, u, &Int{0, 1, 0, 0}, &rem)
 	copy(expected[:], u)
@@ -760,30 +760,37 @@ func TestWriteToSlice(t *testing.T) {
 	}
 
 }
+
 func TestInt_WriteToArray(t *testing.T) {
-	x1 := hex2Bytes("0000000000000000000000000000d1e870eec79504c60144cc7f5fc2bad1e611")
+	x1 := hex2Bytes("0102030405060708090a0b0c0d0ed1e870eec79504c60144cc7f5fc2bad1e611")
 	a := big.NewInt(0).SetBytes(x1)
 	fa, _ := FromBig(a)
 
 	{
 		dest := [20]byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
 		fa.WriteToArray20(&dest)
-		exp := hex2Bytes("0000d1e870eec79504c60144cc7f5fc2bad1e611")
+		exp := hex2Bytes("0d0ed1e870eec79504c60144cc7f5fc2bad1e611")
 		if !bytes.Equal(dest[:], exp) {
 			t.Errorf("got %x, expected %x", dest, exp)
 		}
-
 	}
-
 	{
 		dest := [32]byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
 			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
 		fa.WriteToArray32(&dest)
-		exp := hex2Bytes("0000000000000000000000000000d1e870eec79504c60144cc7f5fc2bad1e611")
+		exp := hex2Bytes("0102030405060708090a0b0c0d0ed1e870eec79504c60144cc7f5fc2bad1e611")
+		if !bytes.Equal(dest[:], exp) {
+			t.Errorf("got %x, expected %x", dest, exp)
+		}
+	}
+	{ // A 36-byte slice: the first 32 bytes are overwritten
+		dest := []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}
+		fa.PutUint256(dest)
+		exp := hex2Bytes("0102030405060708090a0b0c0d0ed1e870eec79504c60144cc7f5fc2bad1e611ffffffff")
 		if !bytes.Equal(dest[:], exp) {
 			t.Errorf("got %x, expected %x", dest, exp)
 		}
-
 	}
 }