diff --git a/internal/staging/primitives/uint256/uint256.go b/internal/staging/primitives/uint256/uint256.go index afa6be3f..c2e70ecb 100644 --- a/internal/staging/primitives/uint256/uint256.go +++ b/internal/staging/primitives/uint256/uint256.go @@ -21,6 +21,11 @@ var ( // zero32 is an array of 32 bytes used for the purposes of zeroing and is // defined here to avoid extra allocations. zero32 = [32]byte{} + + // bigUint256Mask is the value 2^256 - 1 (aka max uint256) as a stdlib big + // integer. It is defined here to save allocations in the conversion code. + bigTwo256 = new(big.Int).Lsh(big.NewInt(1), 256) + bigUint256Mask = new(big.Int).Sub(bigTwo256, big.NewInt(1)) ) // Uint256 implements high-performance, zero-allocation, unsigned 256-bit @@ -1776,3 +1781,40 @@ func (n *Uint256) ToBig() *big.Int { n.PutBig(&out) return &out } + +// SetBig sets the uint256 to the passed standard library big integer modulo +// 2^256. +// +// The resulting uint256 will be set to the 2's complement of the provided value +// when it is negative. +// +// The uint256 is returned to support chaining. This enables syntax like: +// n := new(Uint256).SetBig(n2).AddUint64(1) so that n = n2 + 1 where n2 is not +// modified. +// +// PERFORMANCE NOTE: When the caller expects values to potentially be larger +// than a max uint256, it is _highly_ recommended to reduce the value mod 2^256 +// prior to calling this method for better performance. +// +// The reason is that this method requires an allocation and copy when the +// provided big integer is larger than a max uint256 in order to reduce it +// without modifying the provided arg. The caller can avoid this allocation by +// performing the mod 2^256 prior to calling this method with the value. +// +// More concretely, it is around 3 to 4 times faster to perform the reduction +// caller side as well as avoiding the allocation. +func (n *Uint256) SetBig(n2 *big.Int) *Uint256 { + // Take the value mod 2^256 if needed. + tmp := n2 + if n2.BitLen() > 256 { + tmp = new(big.Int).And(n2, bigUint256Mask) + } + + var buf [32]byte + tmp.FillBytes(buf[:]) // Requires Go 1.15. + n.SetBytes(&buf) + if tmp.Sign() < 0 { + n.Negate() + } + return n +} diff --git a/internal/staging/primitives/uint256/uint256_test.go b/internal/staging/primitives/uint256/uint256_test.go index c6a7f106..28c0d79b 100644 --- a/internal/staging/primitives/uint256/uint256_test.go +++ b/internal/staging/primitives/uint256/uint256_test.go @@ -4077,3 +4077,102 @@ func TestUint256ToBigRandom(t *testing.T) { } } } + +// TestUint256SetBig ensures that setting a uint256 to a standard library big +// integer works as expected for edge cases. +func TestUint256SetBig(t *testing.T) { + t.Parallel() + + tests := []struct { + name string // test description + in string // hex encoded big int test value + want string // expected hex encoded uint256 + }{{ + name: "0", + in: "00", + want: "0", + }, { + name: "1", + in: "1", + want: "1", + }, { + name: "-1", + in: "-1", + want: "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + }, { + name: "2", + in: "2", + want: "2", + }, { + name: "-2", + in: "-2", + want: "fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffe", + }, { + name: "max int256 (2^255 - 1)", + in: "7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + want: "7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + }, { + name: "min int256 (-2^255)", + in: "-8000000000000000000000000000000000000000000000000000000000000000", + want: "8000000000000000000000000000000000000000000000000000000000000000", + }, { + name: "max uint256 (2^256 - 1)", + in: "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + want: "ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff", + }, { + name: "max uint256 + 1 (2^256)", + in: "10000000000000000000000000000000000000000000000000000000000000000", + want: "0", + }, { + name: "max uint256 + 2 (2^256 + 1)", + in: "10000000000000000000000000000000000000000000000000000000000000001", + want: "1", + }} + + for _, test := range tests { + // Parse test vals. + in, ok := new(big.Int).SetString(test.in, 16) + if !ok { + t.Errorf("%q: big int parse err for string %s", test.name, test.in) + continue + } + want := hexToUint256(test.want) + + // Ensure setting the uint256 to the test big int produces the expected + // value. + var n Uint256 + n.SetBig(in) + if !n.Eq(want) { + t.Errorf("%q: unexpected result -- got: %x, want: %x", test.name, n, + want) + continue + } + } +} + +// TestUint256SetBigRandom ensures that converting bit ints created from random +// values to uint256s works as expected. +func TestUint256SetBigRandom(t *testing.T) { + t.Parallel() + + // Use a unique random seed each test instance and log it if the tests fail. + seed := time.Now().Unix() + rng := rand.New(rand.NewSource(seed)) + defer func(t *testing.T, seed int64) { + if t.Failed() { + t.Logf("random seed: %d", seed) + } + }(t, seed) + + for i := 0; i < 100; i++ { + // Generate big integer and uint256 pair. + bigN, wantUint256 := randBigIntAndUint256(t, rng) + + // Convert the big int to a uint256 and ensure they match. + uint256Result := new(Uint256).SetBig(bigN) + if !uint256Result.Eq(wantUint256) { + t.Fatalf("mismatched uint256 conversion: %x -- got %x, want %x", + bigN, uint256Result, wantUint256) + } + } +}