From caee7238d105114e595a99474711c0c8c58ef983 Mon Sep 17 00:00:00 2001
From: Joe Harvey <joe.harvey@jasper.io>
Date: Tue, 25 Jan 2022 15:54:36 +1300
Subject: [PATCH] Allow custom unmarshal implementation

---
 ion/unmarshal.go      | 12 ++++++++++++
 ion/unmarshal_test.go | 23 +++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/ion/unmarshal.go b/ion/unmarshal.go
index 2d4908c..e5121d3 100644
--- a/ion/unmarshal.go
+++ b/ion/unmarshal.go
@@ -38,6 +38,11 @@ var typesAcceptableKinds = map[Type][]reflect.Kind{
 	ListType:  {reflect.Slice, reflect.Array},
 }
 
+// Unmarshaler is the interface implemented by types that can unmarshal themselves to Ion.
+type Unmarshaler interface {
+	UnmarshalIon(r Reader) error
+}
+
 // Unmarshal unmarshals Ion data to the given object.
 //
 // User must pass the proper object type to the unmarshalled Ion data.
@@ -289,6 +294,8 @@ func (d *Decoder) DecodeTo(v interface{}) error {
 	return d.decodeTo(rv)
 }
 
+var unmarshalerType = reflect.TypeOf((*Unmarshaler)(nil)).Elem()
+
 func (d *Decoder) decodeTo(v reflect.Value) error {
 	if !v.IsValid() {
 		// Don't actually have anywhere to put this value; skip it.
@@ -305,6 +312,11 @@ func (d *Decoder) decodeTo(v reflect.Value) error {
 		return nil
 	}
 
+	t := v.Type()
+	if t.Kind() != reflect.Ptr && v.CanAddr() && reflect.PtrTo(t).Implements(unmarshalerType) {
+		return v.Addr().Interface().(Unmarshaler).UnmarshalIon(d.r)
+	}
+
 	switch d.r.Type() {
 	case BoolType:
 		return d.decodeBoolTo(v)
diff --git a/ion/unmarshal_test.go b/ion/unmarshal_test.go
index 7d81d6a..20a9240 100644
--- a/ion/unmarshal_test.go
+++ b/ion/unmarshal_test.go
@@ -316,6 +316,29 @@ func TestUnmarshalStructBinary(t *testing.T) {
 	test(ionByteValue, "structToStruct", &foo{}, &foo{2}) // unmarshal IonStruct to a Go struct
 }
 
+type unmarshalMe struct {
+	Name   string
+	custom bool
+}
+
+func (u *unmarshalMe) UnmarshalIon(r Reader) error {
+	u.custom = true
+	return nil
+}
+
+func TestUnmarshalCustomMarshaler(t *testing.T) {
+	ionBinary, err := MarshalBinary(unmarshalMe{
+		Name: "John Doe",
+	})
+	require.NoError(t, err)
+
+	dec := NewDecoder(NewReader(bytes.NewReader(ionBinary)))
+	var decodedResult unmarshalMe
+
+	assert.NoError(t, dec.DecodeTo(&decodedResult))
+	assert.True(t, decodedResult.custom)
+}
+
 func TestUnmarshalListSexpBinary(t *testing.T) {
 	test := func(data []byte, testName string, val, eval interface{}) {
 		t.Run("reflect.TypeOf(val).String()", func(t *testing.T) {