Support marshalling value via function (#124)

Author: kamatama41

marshal.go

@@ -266,6 +266,31 @@ func marshalImpl(vo reflect.Value, w io.Writer, pos uint64, parent *Element, opt
 					}
 					pos, err = writeOne(lst[0])
 				}
+			case reflect.Func:
+				ret := vn.Call(nil)
+				lenRet := len(ret)
+				if lenRet != 1 && lenRet != 2 {
+					return pos, wrapErrorf(ErrIncompatibleType, "number of return value must be 1 or 2 but %d", lenRet)
+				}
+				val := ret[0]
+				if lenRet == 2 {
+					errVal := ret[1]
+					if errVal.Type().String() != "error" {
+						return pos, wrapErrorf(ErrIncompatibleType, "2nd return value must be error but %s", errVal.Type())
+					}
+					if iFace := errVal.Interface(); iFace != nil {
+						return pos, iFace.(error)
+					}
+				}
+				lst, ok := pealElem(val, e.t == DataTypeBinary, tag.omitEmpty)
+				if !ok {
+					return pos, wrapErrorf(
+						ErrIncompatibleType, "marshalling %s from func", val.Type(),
+					)
+				}
+				for _, l := range lst {
+					pos, err = writeOne(l)
+				}
 			default:
 				pos, err = writeOne(vn)
 			}

marshal_test.go

@@ -543,6 +543,151 @@ func TestMarshal_Chan(t *testing.T) {
 	})
 }
 
+func TestMarshal_Func(t *testing.T) {
+	expected := []byte{
+		0x18, 0x53, 0x80, 0x67, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+		0x1F, 0x43, 0xB6, 0x75, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+		0xE7, 0x81, 0x01,
+	}
+	type Cluster struct {
+		Timecode uint64 `ebml:"Timecode"`
+	}
+
+	t.Run("FuncStruct", func(t *testing.T) {
+		input := &struct {
+			Segment struct {
+				Cluster func() Cluster `ebml:"Cluster,size=unknown"`
+			} `ebml:"Segment,size=unknown"`
+		}{}
+		input.Segment.Cluster = func() Cluster {
+			return Cluster{Timecode: 0x01}
+		}
+
+		var b bytes.Buffer
+		if err := Marshal(input, &b); err != nil {
+			t.Fatalf("Unexpected error: '%v'", err)
+		}
+		if !bytes.Equal(expected, b.Bytes()) {
+			t.Errorf("Marshaled binary doesn't match:\n expected: %v,\n      got: %v", expected, b.Bytes())
+		}
+	})
+	t.Run("FuncStructPtr", func(t *testing.T) {
+		input := &struct {
+			Segment struct {
+				Cluster func() *Cluster `ebml:"Cluster,size=unknown"`
+			} `ebml:"Segment,size=unknown"`
+		}{}
+
+		t.Run("Valid", func(t *testing.T) {
+			input.Segment.Cluster = func() *Cluster {
+				return &Cluster{Timecode: 0x01}
+			}
+
+			var b bytes.Buffer
+			if err := Marshal(input, &b); err != nil {
+				t.Fatalf("Unexpected error: '%v'", err)
+			}
+			if !bytes.Equal(expected, b.Bytes()) {
+				t.Errorf("Marshaled binary doesn't match:\n expected: %v,\n      got: %v", expected, b.Bytes())
+			}
+		})
+		t.Run("Nil", func(t *testing.T) {
+			input.Segment.Cluster = func() *Cluster {
+				return nil
+			}
+
+			if err := Marshal(input, &bytes.Buffer{}); !errs.Is(err, ErrIncompatibleType) {
+				t.Fatalf("Expected error: '%v', got: '%v'", ErrIncompatibleType, err)
+			}
+		})
+	})
+	t.Run("FuncStructSlice", func(t *testing.T) {
+		input := &struct {
+			Segment struct {
+				Cluster func() []Cluster `ebml:"Cluster,size=unknown"`
+			} `ebml:"Segment,size=unknown"`
+		}{}
+		input.Segment.Cluster = func() []Cluster {
+			return []Cluster{
+				{Timecode: 0x01},
+				{Timecode: 0x02},
+			}
+		}
+
+		var b bytes.Buffer
+		if err := Marshal(input, &b); err != nil {
+			t.Fatalf("Unexpected error: '%v'", err)
+		}
+		expected := []byte{
+			0x18, 0x53, 0x80, 0x67, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+			0x1F, 0x43, 0xB6, 0x75, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+			0xE7, 0x81, 0x01,
+			0x1F, 0x43, 0xB6, 0x75, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
+			0xE7, 0x81, 0x02,
+		}
+		if !bytes.Equal(expected, b.Bytes()) {
+			t.Errorf("Marshaled binary doesn't match:\n expected: %v,\n      got: %v", expected, b.Bytes())
+		}
+	})
+	t.Run("FuncWithError", func(t *testing.T) {
+		input := &struct {
+			Segment struct {
+				Cluster func() (Cluster, error) `ebml:"Cluster,size=unknown"`
+			} `ebml:"Segment,size=unknown"`
+		}{}
+
+		t.Run("Valid", func(t *testing.T) {
+			input.Segment.Cluster = func() (Cluster, error) {
+				return Cluster{Timecode: 0x01}, nil
+			}
+
+			var b bytes.Buffer
+			if err := Marshal(input, &b); err != nil {
+				t.Fatalf("Unexpected error: '%v'", err)
+			}
+			if !bytes.Equal(expected, b.Bytes()) {
+				t.Errorf("Marshaled binary doesn't match:\n expected: %v,\n      got: %v", expected, b.Bytes())
+			}
+		})
+		t.Run("Error", func(t *testing.T) {
+			expectedErr := errors.New("an error")
+			input.Segment.Cluster = func() (Cluster, error) {
+				return Cluster{Timecode: 0x01}, expectedErr
+			}
+
+			if err := Marshal(input, &bytes.Buffer{}); !errs.Is(err, expectedErr) {
+				t.Fatalf("Expected error: '%v', got: '%v'", expectedErr, err)
+			}
+		})
+		t.Run("NonErrorType", func(t *testing.T) {
+			input := &struct {
+				Segment struct {
+					Cluster func() (*Cluster, int) `ebml:"Cluster,size=unknown"`
+				} `ebml:"Segment,size=unknown"`
+			}{}
+			input.Segment.Cluster = func() (*Cluster, int) {
+				return nil, 1
+			}
+
+			if err := Marshal(input, &bytes.Buffer{}); !errs.Is(err, ErrIncompatibleType) {
+				t.Fatalf("Expected error: '%v', got: '%v'", ErrIncompatibleType, err)
+			}
+		})
+	})
+	t.Run("InvalidFunc", func(t *testing.T) {
+		input := &struct {
+			Segment struct {
+				Cluster func() `ebml:"Cluster,size=unknown"`
+			} `ebml:"Segment,size=unknown"`
+		}{}
+		input.Segment.Cluster = func() {}
+
+		if err := Marshal(input, &bytes.Buffer{}); !errs.Is(err, ErrIncompatibleType) {
+			t.Fatalf("Expected error: '%v', got: '%v'", ErrIncompatibleType, err)
+		}
+	})
+}
+
 func BenchmarkMarshal(b *testing.B) {
 	type EBMLHeader struct {
 		DocType            string `ebml:"EBMLDocType"`