Add Marshal option to reserve data size length (#54)

marshal.go

@@ -52,13 +52,17 @@ var (
 //   // Field appears as element "EBMLVersion" and
 //   // the size of the element data is reserved by 4 bytes.
 //   Field uint64 `ebml:EBMLVersion,size=4`
-func Marshal(val interface{}, w io.Writer) error {
+func Marshal(val interface{}, w io.Writer, opts ...MarshalOption) error {
+	options := &MarshalOptions{}
+	for _, o := range opts {
+		o(options)
+	}
 	vo := reflect.ValueOf(val).Elem()
 
-	return marshalImpl(vo, w)
+	return marshalImpl(vo, w, options)
 }
 
-func marshalImpl(vo reflect.Value, w io.Writer) error {
+func marshalImpl(vo reflect.Value, w io.Writer, options *MarshalOptions) error {
 	for i := 0; i < vo.NumField(); i++ {
 		vn := vo.Field(i)
 		tn := vo.Type().Field(i)
@@ -118,7 +122,7 @@ func marshalImpl(vo reflect.Value, w io.Writer) error {
 				}
 
 				if e.t == TypeMaster {
-					if err := marshalImpl(vn, bw); err != nil {
+					if err := marshalImpl(vn, bw, options); err != nil {
 						return err
 					}
 				} else {
@@ -133,7 +137,7 @@ func marshalImpl(vo reflect.Value, w io.Writer) error {
 
 				// Write element with length
 				if !unknown {
-					bsz := encodeDataSize(uint64(bw.(*bytes.Buffer).Len()), 0)
+					bsz := encodeDataSize(uint64(bw.(*bytes.Buffer).Len()), options.dataSizeLen)
 					if _, err := w.Write(bsz); err != nil {
 						return err
 					}
@@ -146,3 +150,18 @@ func marshalImpl(vo reflect.Value, w io.Writer) error {
 	}
 	return nil
 }
+
+// MarshalOption configures a MarshalOptions struct
+type MarshalOption func(*MarshalOptions)
+
+// MarshalOptions stores options for marshalling
+type MarshalOptions struct {
+	dataSizeLen uint64
+}
+
+// WithDataSizeLen returns an MarshalOption which sets number of reserved bytes of element data size
+func WithDataSizeLen(l int) MarshalOption {
+	return func(opts *MarshalOptions) {
+		opts.dataSizeLen = uint64(l)
+	}
+}

marshal_test.go

@@ -114,6 +114,34 @@ func ExampleMarshal() {
 	// 0x1a, 0x45, 0xdf, 0xa3, 0x90, 0x42, 0x82, 0x85, 0x77, 0x65, 0x62, 0x6d, 0x00, 0x42, 0x87, 0x81, 0x02, 0x42, 0x85, 0x81, 0x02,
 }
 
+func ExampleWithDataSizeLen() {
+	type EBMLHeader struct {
+		DocType            string `ebml:"EBMLDocType"`
+		DocTypeVersion     uint64 `ebml:"EBMLDocTypeVersion"`
+		DocTypeReadVersion uint64 `ebml:"EBMLDocTypeReadVersion"`
+	}
+	type TestEBML struct {
+		Header EBMLHeader `ebml:"EBML"`
+	}
+	s := TestEBML{
+		Header: EBMLHeader{
+			DocType:            "webm",
+			DocTypeVersion:     2,
+			DocTypeReadVersion: 2,
+		},
+	}
+
+	var b bytes.Buffer
+	if err := Marshal(&s, &b, WithDataSizeLen(2)); err != nil {
+		panic(err)
+	}
+	for _, b := range b.Bytes() {
+		fmt.Printf("0x%02x, ", int(b))
+	}
+	// Output:
+	// 0x1a, 0x45, 0xdf, 0xa3, 0x40, 0x13, 0x42, 0x82, 0x40, 0x05, 0x77, 0x65, 0x62, 0x6d, 0x00, 0x42, 0x87, 0x40, 0x01, 0x02, 0x42, 0x85, 0x40, 0x01, 0x02,
+}
+
 func TestMarshal_Tag(t *testing.T) {
 	tagged := struct {
 		DocCustomNamedType string `ebml:"EBMLDocType"`