gomsgpack_polymorphic_type.go

package main

import (
	"bytes"
	"fmt"

	"github.com/hashicorp/go-msgpack/codec"
)

var h codec.MsgpackHandle

type EntryOp struct {
	Op    int
	Entry EntryType
}

func (o *EntryOp) MarshalBinary() (data []byte, err error) {
	// bs will grow if needed but allocate enough to avoid reallocation in common
	// case.
	bs := make([]byte, 128)
	enc := codec.NewEncoderBytes(&bs, &h)
	// Encode kind first
	err = enc.Encode(o.Entry.Kind())
	if err != nil {
		return nil, err
	}
	// Then actual value using alisa trick to avoid infinite recursion
	type Alias EntryOp
	err = enc.Encode(struct {
		*Alias
	}{
		Alias: (*Alias)(&o),
	})
	if err != nil {
		return nil, err
	}
	return bs, nil
}
func (o *EntryOp) UnmarshalBinary(data []byte) error {
	// First decode the kind prefix
	var kind int
	dec := codec.NewDecoderBytes(data, &h)
	if err := dec.Decode(&kind); err != nil {
		return err
	}

	// Then decode the real thing with appropriate kind of Entry
	o.Entry = makeEntry(kind)

	// Alias juggling to prevent infinite recursive calls back to this decode
	// method.
	type Alias EntryOp
	as := struct {
		*Alias
	}{
		Alias: (*Alias)(o),
	}
	if err := dec.Decode(&as); err != nil {
		return err
	}
	return nil
}

type Foo struct {
	Foo int
}

func (f *Foo) Kind() int { return 1 }

type Bar struct {
	Bar int
}

func (f *Bar) Kind() int { return 2 }

type EntryType interface {
	Kind() int
}

func makeEntry(kind int) EntryType {
	switch kind {
	case 1:
		return &Foo{}
	case 2:
		return &Bar{}
	default:
		panic("invalid kind")
	}
}

func main() {

	var buf bytes.Buffer

	// note this needs to be a pointer since that is the type the MarshalBinary method is defined on
	e := &EntryOp{
		Op:    1,
		Entry: &Foo{Foo: 3},
	}

	enc := codec.NewEncoder(&buf, &h)
	if err := enc.Encode(e); err != nil {
		panic(err)
	}

	fmt.Println(buf.String())

	var e2 EntryOp

	dec := codec.NewDecoder(&buf, &h)
	if err := dec.Decode(&e2); err != nil {
		panic(err)
	}

	fmt.Printf("%q", e2)
	// Output:
	// {'\x01' %!q(*main.Foo=&{3})}
	// 
	// Which shows it preserved the Foo type and it's values.
}