Repository: ferranbt/fastssz
Pull request #100 (2022-08-02 08:48:37)

Support time.Time
This allows encoding and decoding of the `time.Time` type in structs, to `uint64` values.


Patch:
diff --git a/encode.go b/encode.go
index a4112cc..5848023 100644
--- a/encode.go
+++ b/encode.go
@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math/bits"
+	"time"
 )
 
 // MarshalSSZ marshals an object
@@ -66,6 +67,11 @@ func UnmarshalBool(src []byte) bool {
 	return false
 }
 
+// UnmarshalTime unmarshals a time.Time from the src input
+func UnmarshalTime(src []byte) time.Time {
+	return time.Unix(int64(UnmarshallUint64(src)), 0).UTC()
+}
+
 // ---- Marshal functions ----
 
 // MarshalUint64 marshals a little endian uint64 to dst
@@ -108,6 +114,11 @@ func MarshalBool(dst []byte, b bool) []byte {
 	return dst
 }
 
+// MarshalTime marshals a time to dst
+func MarshalTime(dst []byte, t time.Time) []byte {
+	return MarshalUint64(dst, uint64(t.Unix()))
+}
+
 // ---- offset functions ----
 
 // WriteOffset writes an offset to dst
diff --git a/sszgen/.gitignore b/sszgen/.gitignore
new file mode 100644
index 0000000..01a263b
--- /dev/null
+++ b/sszgen/.gitignore
@@ -0,0 +1,1 @@
+sszgen
diff --git a/sszgen/generator/generator.go b/sszgen/generator/generator.go
index b0efd97..a60fbf5 100644
--- a/sszgen/generator/generator.go
+++ b/sszgen/generator/generator.go
@@ -209,6 +209,8 @@ const (
 	TypeContainer
 	// TypeReference is a SSZ reference
 	TypeReference
+	// TypeTime is a timestamp
+	TypeTime
 )
 
 func (t Type) String() string {
@@ -231,6 +233,8 @@ func (t Type) String() string {
 		return "container"
 	case TypeReference:
 		return "reference"
+	case TypeTime:
+		return "time.Time"
 	default:
 		panic("not found")
 	}
@@ -1056,7 +1060,9 @@ func (e *env) parseASTFieldType(name, tags string, expr ast.Expr) (*Value, error
 		name := obj.X.(*ast.Ident).Name
 		sel := obj.Sel.Name
 
-		if sel == "Bitlist" {
+		if name == "time" && sel == "Time" {
+			return &Value{t: TypeTime, s: 8}, nil
+		} else if sel == "Bitlist" {
 			// go-bitfield/Bitlist
 			maxSize, ok := getTagsInt(tags, "ssz-max")
 			if !ok {
@@ -1174,6 +1180,8 @@ func (v *Value) isFixed() bool {
 			return true
 		}
 		return false
+	case TypeTime:
+		return true
 	default:
 		// TypeUndefined should be the only type to fallthrough to this case
 		// TypeUndefined always means there is a fatal error in the parsing logic
diff --git a/sszgen/generator/hash.go b/sszgen/generator/hash.go
index 03d2bc8..a28bb3d 100755
--- a/sszgen/generator/hash.go
+++ b/sszgen/generator/hash.go
@@ -209,6 +209,9 @@ func (v *Value) hashTreeRoot(name string, appendBytes bool) string {
 			"htrCall": htrCall,
 		})
 
+	case TypeTime:
+		return fmt.Sprintf("hh.PutUint64(uint64(%s.Unix()))", name)
+
 	default:
 		panic(fmt.Errorf("hash not implemented for type %s", v.t.String()))
 	}
diff --git a/sszgen/generator/marshal.go b/sszgen/generator/marshal.go
index a2d284a..b0b1f8e 100644
--- a/sszgen/generator/marshal.go
+++ b/sszgen/generator/marshal.go
@@ -77,6 +77,9 @@ func (v *Value) marshal() string {
 	case TypeList:
 		return v.marshalList()
 
+	case TypeTime:
+		return fmt.Sprintf("dst = ssz.MarshalTime(dst, ::.%s)", v.name)
+
 	default:
 		panic(fmt.Errorf("marshal not implemented for type %s", v.t.String()))
 	}
diff --git a/sszgen/generator/unmarshal.go b/sszgen/generator/unmarshal.go
index 6592246..b6640ad 100644
--- a/sszgen/generator/unmarshal.go
+++ b/sszgen/generator/unmarshal.go
@@ -104,6 +104,9 @@ func (v *Value) unmarshal(dst string) string {
 	case TypeBool:
 		return fmt.Sprintf("::.%s = ssz.UnmarshalBool(%s)", v.name, dst)
 
+	case TypeTime:
+		return fmt.Sprintf("::.%s = ssz.UnmarshalTime(%s)", v.name, dst)
+
 	default:
 		panic(fmt.Errorf("unmarshal not implemented for type %d", v.t))
 	}
@@ -352,6 +355,6 @@ func (v *Value) createSlice(useNumVariable bool) string {
 		return fmt.Sprintf("::.%s = make([][]byte, %s)", v.name, size)
 
 	default:
-		panic(fmt.Sprintf("create not implemented for type %s", v.e.t.String()))
+		panic(fmt.Sprintf("create not implemented for %s type %s", v.name, v.e.t.String()))
 	}
 }


Test patch:
diff --git a/sszgen/testcases/time.go b/sszgen/testcases/time.go
new file mode 100644
index 0000000..a2b402f
--- /dev/null
+++ b/sszgen/testcases/time.go
@@ -0,0 +1,13 @@
+package testcases
+
+import "time"
+
+type TimeType struct {
+	Timestamp time.Time
+	Int       uint64
+}
+
+type TimeRawType struct {
+	Timestamp uint64
+	Int       uint64
+}
diff --git a/sszgen/testcases/time_encoding.go b/sszgen/testcases/time_encoding.go
new file mode 100644
index 0000000..b85ebc2
--- /dev/null
+++ b/sszgen/testcases/time_encoding.go
@@ -0,0 +1,137 @@
+// Code generated by fastssz. DO NOT EDIT.
+// Hash: 04ac916e8f15c6e500617c8195a73e18b9b9f9a76652d4bec1f96b4bb045bef4
+package testcases
+
+import (
+	ssz "github.com/ferranbt/fastssz"
+)
+
+// MarshalSSZ ssz marshals the TimeType object
+func (t *TimeType) MarshalSSZ() ([]byte, error) {
+	return ssz.MarshalSSZ(t)
+}
+
+// MarshalSSZTo ssz marshals the TimeType object to a target array
+func (t *TimeType) MarshalSSZTo(buf []byte) (dst []byte, err error) {
+	dst = buf
+
+	// Field (0) 'Timestamp'
+	dst = ssz.MarshalTime(dst, t.Timestamp)
+
+	// Field (1) 'Int'
+	dst = ssz.MarshalUint64(dst, t.Int)
+
+	return
+}
+
+// UnmarshalSSZ ssz unmarshals the TimeType object
+func (t *TimeType) UnmarshalSSZ(buf []byte) error {
+	var err error
+	size := uint64(len(buf))
+	if size != 16 {
+		return ssz.ErrSize
+	}
+
+	// Field (0) 'Timestamp'
+	t.Timestamp = ssz.UnmarshalTime(buf[0:8])
+
+	// Field (1) 'Int'
+	t.Int = ssz.UnmarshallUint64(buf[8:16])
+
+	return err
+}
+
+// SizeSSZ returns the ssz encoded size in bytes for the TimeType object
+func (t *TimeType) SizeSSZ() (size int) {
+	size = 16
+	return
+}
+
+// HashTreeRoot ssz hashes the TimeType object
+func (t *TimeType) HashTreeRoot() ([32]byte, error) {
+	return ssz.HashWithDefaultHasher(t)
+}
+
+// HashTreeRootWith ssz hashes the TimeType object with a hasher
+func (t *TimeType) HashTreeRootWith(hh ssz.HashWalker) (err error) {
+	indx := hh.Index()
+
+	// Field (0) 'Timestamp'
+	hh.PutUint64(uint64(t.Timestamp.Unix()))
+
+	// Field (1) 'Int'
+	hh.PutUint64(t.Int)
+
+	hh.Merkleize(indx)
+	return
+}
+
+// GetTree ssz hashes the TimeType object
+func (t *TimeType) GetTree() (*ssz.Node, error) {
+	return ssz.ProofTree(t)
+}
+
+// MarshalSSZ ssz marshals the TimeRawType object
+func (t *TimeRawType) MarshalSSZ() ([]byte, error) {
+	return ssz.MarshalSSZ(t)
+}
+
+// MarshalSSZTo ssz marshals the TimeRawType object to a target array
+func (t *TimeRawType) MarshalSSZTo(buf []byte) (dst []byte, err error) {
+	dst = buf
+
+	// Field (0) 'Timestamp'
+	dst = ssz.MarshalUint64(dst, t.Timestamp)
+
+	// Field (1) 'Int'
+	dst = ssz.MarshalUint64(dst, t.Int)
+
+	return
+}
+
+// UnmarshalSSZ ssz unmarshals the TimeRawType object
+func (t *TimeRawType) UnmarshalSSZ(buf []byte) error {
+	var err error
+	size := uint64(len(buf))
+	if size != 16 {
+		return ssz.ErrSize
+	}
+
+	// Field (0) 'Timestamp'
+	t.Timestamp = ssz.UnmarshallUint64(buf[0:8])
+
+	// Field (1) 'Int'
+	t.Int = ssz.UnmarshallUint64(buf[8:16])
+
+	return err
+}
+
+// SizeSSZ returns the ssz encoded size in bytes for the TimeRawType object
+func (t *TimeRawType) SizeSSZ() (size int) {
+	size = 16
+	return
+}
+
+// HashTreeRoot ssz hashes the TimeRawType object
+func (t *TimeRawType) HashTreeRoot() ([32]byte, error) {
+	return ssz.HashWithDefaultHasher(t)
+}
+
+// HashTreeRootWith ssz hashes the TimeRawType object with a hasher
+func (t *TimeRawType) HashTreeRootWith(hh ssz.HashWalker) (err error) {
+	indx := hh.Index()
+
+	// Field (0) 'Timestamp'
+	hh.PutUint64(t.Timestamp)
+
+	// Field (1) 'Int'
+	hh.PutUint64(t.Int)
+
+	hh.Merkleize(indx)
+	return
+}
+
+// GetTree ssz hashes the TimeRawType object
+func (t *TimeRawType) GetTree() (*ssz.Node, error) {
+	return ssz.ProofTree(t)
+}
diff --git a/sszgen/testcases/time_test.go b/sszgen/testcases/time_test.go
new file mode 100644
index 0000000..e98e7c7
--- /dev/null
+++ b/sszgen/testcases/time_test.go
@@ -0,0 +1,144 @@
+package testcases
+
+import (
+	"bytes"
+	"testing"
+	"time"
+)
+
+func TestTimeRoot(t *testing.T) {
+	tests := []struct {
+		name      string
+		timestamp int64
+		root      []byte
+	}{
+		{
+			name:      "Zero",
+			timestamp: 0,
+			root:      []byte{0xf5, 0xa5, 0xfd, 0x42, 0xd1, 0x6a, 0x20, 0x30, 0x27, 0x98, 0xef, 0x6e, 0xd3, 0x09, 0x97, 0x9b, 0x43, 0x00, 0x3d, 0x23, 0x20, 0xd9, 0xf0, 0xe8, 0xea, 0x98, 0x31, 0xa9, 0x27, 0x59, 0xfb, 0x4b},
+		},
+		{
+			name:      "Max",
+			timestamp: 0x7fffffffffffffff,
+			root:      []byte{0xd6, 0xa4, 0x84, 0x7f, 0x18, 0x2d, 0xec, 0x2b, 0xd3, 0x99, 0xd9, 0x7d, 0xb5, 0x96, 0xb4, 0x83, 0x15, 0x7c, 0x5c, 0x62, 0x4c, 0x12, 0x0b, 0x25, 0x0a, 0xde, 0xa9, 0xf5, 0xe9, 0x93, 0x6c, 0x4f},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			timeStruct := TimeType{Timestamp: time.Unix(test.timestamp, 0)}
+			timeRoot, err := timeStruct.HashTreeRoot()
+			if err != nil {
+				t.Fatalf("unexpected error: %v\n", err)
+			}
+			// Expect the root to match the expectation.
+			if !bytes.Equal(test.root, timeRoot[:]) {
+				t.Fatalf("root mismatch with time type")
+			}
+
+			rawStruct := TimeRawType{Timestamp: uint64(test.timestamp)}
+			rawRoot, err := rawStruct.HashTreeRoot()
+			if err != nil {
+				t.Fatalf("unexpected error: %v\n", err)
+			}
+			// Expect the root to match the expectation.
+			if !bytes.Equal(test.root, rawRoot[:]) {
+				t.Fatalf("root mismatch with raw type")
+			}
+		})
+	}
+}
+
+func TestTimeEncode(t *testing.T) {
+	tests := []struct {
+		name       string
+		timestamp  int64
+		marshalled []byte
+	}{
+		{
+			name:       "Zero",
+			timestamp:  0,
+			marshalled: []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+		},
+		{
+			name:       "Real",
+			timestamp:  0x62e8e1b1,
+			marshalled: []byte{0xb1, 0xe1, 0xe8, 0x62, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+		},
+		{
+			name:       "Max",
+			timestamp:  0x7fffffffffffffff,
+			marshalled: []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			timeStruct := TimeType{Timestamp: time.Unix(test.timestamp, 0)}
+			timeMarshalled, err := timeStruct.MarshalSSZ()
+			if err != nil {
+				t.Fatalf("unexpected error: %v\n", err)
+			}
+			if !bytes.Equal(test.marshalled, timeMarshalled) {
+				t.Fatalf("marshal mismatch with time type (%v != %v)", test.marshalled, timeMarshalled)
+			}
+
+			rawStruct := TimeRawType{Timestamp: uint64(test.timestamp)}
+			rawMarshalled, err := rawStruct.MarshalSSZ()
+			if err != nil {
+				t.Fatalf("unexpected error: %v\n", err)
+			}
+			if !bytes.Equal(test.marshalled, rawMarshalled) {
+				t.Fatalf("marshal mismatch with raw type")
+			}
+		})
+	}
+}
+
+func TestTimeDecode(t *testing.T) {
+	tests := []struct {
+		name       string
+		marshalled []byte
+		timeStruct *TimeType
+		rawStruct  *TimeRawType
+	}{
+		{
+			name:       "Zero",
+			marshalled: []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+			timeStruct: &TimeType{
+				Timestamp: time.Unix(0, 0),
+			},
+			rawStruct: &TimeRawType{},
+		},
+		{
+			name:       "Real",
+			marshalled: []byte{0xb1, 0xe1, 0xe8, 0x62, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+			timeStruct: &TimeType{
+				Timestamp: time.Unix(0x62e8e1b1, 0),
+			},
+			rawStruct: &TimeRawType{
+				Timestamp: 0x62e8e1b1,
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			timeStruct := TimeType{}
+			if err := timeStruct.UnmarshalSSZ(test.marshalled); err != nil {
+				t.Fatalf("unexpected error: %v\n", err)
+			}
+			if !test.timeStruct.Timestamp.Equal(timeStruct.Timestamp) {
+				t.Fatalf("unmarshal mismatch with time type (%v != %v)", test.timeStruct.Timestamp, timeStruct.Timestamp)
+			}
+
+			rawStruct := TimeRawType{}
+			if err := rawStruct.UnmarshalSSZ(test.marshalled); err != nil {
+				t.Fatalf("unexpected error: %v\n", err)
+			}
+			if test.rawStruct.Timestamp != rawStruct.Timestamp {
+				t.Fatalf("unmarshal mismatch with time type")
+			}
+		})
+	}
+}