rlp/rlpgen: add support for uint256

rlp/rlpgen/gen.go

@@ -283,7 +283,7 @@ func (op byteArrayOp) genDecode(ctx *genContext) (string, string) {
 	return resultV, b.String()
 }
 
-// bigIntNoPtrOp handles non-pointer big.Int.
+// bigIntOp handles big.Int.
 // This exists because big.Int has it's own decoder operation on rlp.Stream,
 // but the decode method returns *big.Int, so it needs to be dereferenced.
 type bigIntOp struct {
@@ -330,6 +330,49 @@ func (op bigIntOp) genDecode(ctx *genContext) (string, string) {
 	return result, b.String()
 }
 
+// uint256Op handles "github.com/holiman/uint256".Int
+type uint256Op struct {
+	pointer bool
+}
+
+func (op uint256Op) genWrite(ctx *genContext, v string) string {
+	var b bytes.Buffer
+
+	dst := v
+	if !op.pointer {
+		dst = "&" + v
+	}
+	fmt.Fprintf(&b, "w.WriteUint256(%s)\n", dst)
+
+	// Wrap with nil check.
+	if op.pointer {
+		code := b.String()
+		b.Reset()
+		fmt.Fprintf(&b, "if %s == nil {\n", v)
+		fmt.Fprintf(&b, "  w.Write(rlp.EmptyString)")
+		fmt.Fprintf(&b, "} else {\n")
+		fmt.Fprint(&b, code)
+		fmt.Fprintf(&b, "}\n")
+	}
+
+	return b.String()
+}
+
+func (op uint256Op) genDecode(ctx *genContext) (string, string) {
+	ctx.addImport("github.com/holiman/uint256")
+
+	var b bytes.Buffer
+	resultV := ctx.temp()
+	fmt.Fprintf(&b, "var %s uint256.Int\n", resultV)
+	fmt.Fprintf(&b, "if err := dec.ReadUint256(&%s); err != nil { return err }\n", resultV)
+
+	result := resultV
+	if op.pointer {
+		result = "&" + resultV
+	}
+	return result, b.String()
+}
+
 // encoderDecoderOp handles rlp.Encoder and rlp.Decoder.
 // In order to be used with this, the type must implement both interfaces.
 // This restriction may be lifted in the future by creating separate ops for
@@ -635,6 +678,9 @@ func (bctx *buildContext) makeOp(name *types.Named, typ types.Type, tags rlpstru
 		if isBigInt(typ) {
 			return bigIntOp{}, nil
 		}
+		if isUint256(typ) {
+			return uint256Op{}, nil
+		}
 		if typ == bctx.rawValueType {
 			return bctx.makeRawValueOp(), nil
 		}
@@ -647,6 +693,9 @@ func (bctx *buildContext) makeOp(name *types.Named, typ types.Type, tags rlpstru
 		if isBigInt(typ.Elem()) {
 			return bigIntOp{pointer: true}, nil
 		}
+		if isUint256(typ.Elem()) {
+			return uint256Op{pointer: true}, nil
+		}
 		// Encoder/Decoder interfaces.
 		if bctx.isEncoder(typ) {
 			if bctx.isDecoder(typ) {

rlp/rlpgen/gen_test.go

@@ -47,7 +47,7 @@ func init() {
 	}
 }
 
-var tests = []string{"uints", "nil", "rawvalue", "optional", "bigint"}
+var tests = []string{"uints", "nil", "rawvalue", "optional", "bigint", "uint256"}
 
 func TestOutput(t *testing.T) {
 	for _, test := range tests {

rlp/rlpgen/testdata/uint256.in.txt

@@ -0,0 +1,10 @@
+// -*- mode: go -*-
+
+package test
+
+import "github.com/holiman/uint256"
+
+type Test struct {
+	Int      *uint256.Int
+	IntNoPtr uint256.Int
+}

rlp/rlpgen/testdata/uint256.out.txt

@@ -0,0 +1,44 @@
+package test
+
+import "github.com/ethereum/go-ethereum/rlp"
+import "github.com/holiman/uint256"
+import "io"
+
+func (obj *Test) EncodeRLP(_w io.Writer) error {
+	w := rlp.NewEncoderBuffer(_w)
+	_tmp0 := w.List()
+	if obj.Int == nil {
+		w.Write(rlp.EmptyString)
+	} else {
+		w.WriteUint256(obj.Int)
+	}
+	w.WriteUint256(&obj.IntNoPtr)
+	w.ListEnd(_tmp0)
+	return w.Flush()
+}
+
+func (obj *Test) DecodeRLP(dec *rlp.Stream) error {
+	var _tmp0 Test
+	{
+		if _, err := dec.List(); err != nil {
+			return err
+		}
+		// Int:
+		var _tmp1 uint256.Int
+		if err := dec.DecodeUint256(&_tmp1); err != nil {
+			return err
+		}
+		_tmp0.Int = &_tmp1
+		// IntNoPtr:
+		var _tmp2 uint256.Int
+		if err := dec.DecodeUint256(&_tmp2); err != nil {
+			return err
+		}
+		_tmp0.IntNoPtr = _tmp2
+		if err := dec.ListEnd(); err != nil {
+			return err
+		}
+	}
+	*obj = _tmp0
+	return nil
+}

rlp/rlpgen/types.go

@@ -97,6 +97,16 @@ func isBigInt(typ types.Type) bool {
 	return name.Pkg().Path() == "math/big" && name.Name() == "Int"
 }
 
+// isUint256 checks whether 'typ' is "github.com/holiman/uint256".Int.
+func isUint256(typ types.Type) bool {
+	named, ok := typ.(*types.Named)
+	if !ok {
+		return false
+	}
+	name := named.Obj()
+	return name.Pkg().Path() == "github.com/holiman/uint256" && name.Name() == "Int"
+}
+
 // isByte checks whether the underlying type of 'typ' is uint8.
 func isByte(typ types.Type) bool {
 	basic, ok := resolveUnderlying(typ).(*types.Basic)