2 files changed, 48 insertions, 8 deletions
diff --git a/common/math/big.go b/common/math/big.go
index c0508c102..704ca40a9 100644
--- a/common/math/big.go
+++ b/common/math/big.go
@@ -28,6 +28,13 @@ var (
 	MaxBig256 = new(big.Int).Set(tt256m1)
 )
 
+const (
+	// number of bits in a big.Word
+	wordBits = 32 << (uint64(^big.Word(0)) >> 63)
+	// number of bytes in a big.Word
+	wordBytes = wordBits / 8
+)
+
 // ParseBig256 parses s as a 256 bit integer in decimal or hexadecimal syntax.
 // Leading zeros are accepted. The empty string parses as zero.
 func ParseBig256(s string) (*big.Int, bool) {
@@ -91,12 +98,25 @@ func FirstBitSet(v *big.Int) int {
 // PaddedBigBytes encodes a big integer as a big-endian byte slice. The length
 // of the slice is at least n bytes.
 func PaddedBigBytes(bigint *big.Int, n int) []byte {
-	bytes := bigint.Bytes()
-	if len(bytes) >= n {
-		return bytes
+	if bigint.BitLen()/8 >= n {
+		return bigint.Bytes()
 	}
 	ret := make([]byte, n)
-	return append(ret[:len(ret)-len(bytes)], bytes...)
+	ReadBits(bigint, ret)
+	return ret
+}
+
+// ReadBits encodes the absolute value of bigint as big-endian bytes. Callers must ensure
+// that buf has enough space. If buf is too short the result will be incomplete.
+func ReadBits(bigint *big.Int, buf []byte) {
+	i := len(buf)
+	for _, d := range bigint.Bits() {
+		for j := 0; j < wordBytes && i > 0; j++ {
+			i--
+			buf[i] = byte(d)
+			d >>= 8
+		}
+	}
 }
 
 // U256 encodes as a 256 bit two's complement number. This operation is destructive.
@@ -119,9 +139,6 @@ func S256(x *big.Int) *big.Int {
 	}
 }
 
-// wordSize is the size number of bits in a big.Word.
-const wordSize = 32 << (uint64(^big.Word(0)) >> 63)
-
 // Exp implements exponentiation by squaring.
 // Exp returns a newly-allocated big integer and does not change
 // base or exponent. The result is truncated to 256 bits.
@@ -131,7 +148,7 @@ func Exp(base, exponent *big.Int) *big.Int {
 	result := big.NewInt(1)
 
 	for _, word := range exponent.Bits() {
-		for i := 0; i < wordSize; i++ {
+		for i := 0; i < wordBits; i++ {
 			if word&1 == 1 {
 				U256(result.Mul(result, base))
 			}
diff --git a/common/math/big_test.go b/common/math/big_test.go
index a0f48a8eb..6eb13f4f1 100644
--- a/common/math/big_test.go
+++ b/common/math/big_test.go
@@ -18,6 +18,7 @@ package math
 
 import (
 	"bytes"
+	"encoding/hex"
 	"math/big"
 	"testing"
 )
@@ -131,6 +132,28 @@ func TestPaddedBigBytes(t *testing.T) {
 	}
 }
 
+func BenchmarkPaddedBigBytes(b *testing.B) {
+	bigint := MustParseBig256("123456789123456789123456789123456789")
+	for i := 0; i < b.N; i++ {
+		PaddedBigBytes(bigint, 32)
+	}
+}
+
+func TestReadBits(t *testing.T) {
+	check := func(input string) {
+		want, _ := hex.DecodeString(input)
+		int, _ := new(big.Int).SetString(input, 16)
+		buf := make([]byte, len(want))
+		ReadBits(int, buf)
+		if !bytes.Equal(buf, want) {
+			t.Errorf("have: %x\nwant: %x", buf, want)
+		}
+	}
+	check("000000000000000000000000000000000000000000000000000000FEFCF3F8F0")
+	check("0000000000012345000000000000000000000000000000000000FEFCF3F8F0")
+	check("18F8F8F1000111000110011100222004330052300000000000000000FEFCF3F8F0")
+}
+
 func TestU256(t *testing.T) {
 	tests := []struct{ x, y *big.Int }{
 		{x: big.NewInt(0), y: big.NewInt(0)},