RegisterInterface() so we can implement custom JSON codec

2026-09-02 22:27:19 +00:00 · 2015-01-04 17:33:18 -08:00
parent 7a8a0fefc7
commit 13ca6fb241
9 changed files with 446 additions and 178 deletions
@@ -96,62 +96,15 @@ foo := Foo{Dog{}}
 buf, n, err := new(bytes.Buffer), new(int64), new(error)
 WriteBinary(foo, buf, n, err)

-// This errors with "Unknown field type interface"
+// This errors because we don't know whether to read a Dog or Cat.
 foo2 := ReadBinary(Foo{}, buf, n, err)
 ```

 In the above example, `ReadBinary()` fails because the `Greeter` field for `Foo{}`
 is ambiguous -- it could be either a `Dog{}` or a `Cat{}`, like a union structure.
-In this case, the user must define a custom encoder/decoder as follows:
+The solution is to declare the concrete implementation types for interfaces:

 ```go
-const (
-	GreeterTypeDog = byte(0x01)
-	GreeterTypeCat = byte(0x02)
-)
-
-func GreeterEncoder(o interface{}, w io.Writer, n *int64, err *error) {
-	switch o.(type) {
-	case Dog:
-		WriteByte(GreeterTypeDog, w, n, err)
-		WriteBinary(o, w, n, err)
-	case Cat:
-		WriteByte(GreeterTypeCat, w, n, err)
-		WriteBinary(o, w, n, err)
-	default:
-		*err = errors.New(fmt.Sprintf("Unknown greeter type %v", o))
-	}
-}
-
-func GreeterDecoder(r Unreader, n *int64, err *error) interface{} {
-	switch t := ReadByte(r, n, err); t {
-	case GreeterTypeDog:
-		return ReadBinary(Dog{}, r, n, err)
-	case GreeterTypeCat:
-		return ReadBinary(Cat{}, r, n, err)
-	default:
-		*err = errors.New(fmt.Sprintf("Unknown greeter type byte %X", t))
-		return nil
-	}
-}
-
-// This tells the reflection system to use the custom encoder/decoder functions
-// for encoding/decoding interface struct-field types.
-var _ = RegisterType(&TypeInfo{
-	Type: reflect.TypeOf((*Greeter)(nil)).Elem(),
-	Encoder: GreeterEncoder,
-	Decoder: GreeterDecoder,
-})
-```
-
-Sometimes you want to always prefix a globally unique type byte while encoding,
-whether or not the declared type is an interface or concrete type.
-In this case, you can declare a "TypeByte() byte" function on the struct (as
-a value receiver, not a pointer receiver!), and you can skip the declaration of
-a custom encoder.  The decoder must "peek" the byte instead of consuming it.
-
-```go
-
 type Dog struct{}
 func (d Dog) TypeByte() byte { return GreeterTypeDog }
 func (d Dog) Greet() string { return "Woof!" }
@@ -160,21 +113,18 @@ type Cat struct{}
 func (c Cat) TypeByte() byte { return GreeterTypeCat }
 func (c Cat) Greet() string { return "Meow!" }

-func GreeterDecoder(r Unreader, n *int64, err *error) interface{} {
-	// We must peek the type byte because ReadBinary() expects it.
-	switch t := PeekByte(r, n, err); t {
-	case GreeterTypeDog:
-		return ReadBinary(Dog{}, r, n, err)
-	case GreeterTypeCat:
-		return ReadBinary(Cat{}, r, n, err)
-	default:
-		*err = errors.New(fmt.Sprintf("Unknown greeter type byte %X", t))
-		return nil
-	}
-}
-
-var _ = RegisterType(&TypeInfo{
-	Type: reflect.TypeOf((*Greeter)(nil)).Elem(),
-	Decoder: GreeterDecoder,
+var _ = RegisterInterface(
+	struct{Greeter}{},
+	ConcreteType{Dog{}},
+	ConcreteType{Cat{}},
 })
 ```
+
+NOTE: The TypeByte() is written and expected to be read even when the struct
+is encoded or decoded directly:
+
+```go
+WriteBinary(Dog{}, buf, n, err)        // Writes GreeterTypeDog byte
+dog_ := ReadBinary(Dog{}, buf, n, err) // Expects to read GreeterTypeDog byte
+dog := dog_.(Dog)                      // ok if *err != nil, otherwise dog_ == nil.
+```
@@ -11,21 +11,42 @@ import (
 )

 type TypeInfo struct {
-	Type    reflect.Type // The type
-	Encoder Encoder      // Optional custom encoder function
-	Decoder Decoder      // Optional custom decoder function
+	Type reflect.Type // The type

+	// Custom encoder/decoder
+	Encoder Encoder
+	Decoder Decoder
+
+	// If Type is kind reflect.Interface
+	ConcreteTypes map[byte]reflect.Type
+
+	// If Type is concrete
 	HasTypeByte bool
 	TypeByte    byte
 }

+// e.g. If o is struct{Foo}{}, return is the Foo interface type.
+func GetTypeFromStructDeclaration(o interface{}) reflect.Type {
+	rt := reflect.TypeOf(o)
+	if rt.NumField() != 1 {
+		panic("Unexpected number of fields in struct-wrapped declaration of type")
+	}
+	return rt.Field(0).Type
+}
+
+// If o implements HasTypeByte, returns (true, typeByte)
+func GetTypeByteFromStruct(o interface{}) (hasTypeByte bool, typeByte byte) {
+	if _, ok := o.(HasTypeByte); ok {
+		return true, o.(HasTypeByte).TypeByte()
+	} else {
+		return false, byte(0x00)
+	}
+}
+
 // If a type implements TypeByte, the byte is included
-// as the first byte for encoding.  This is used to encode
-// interfaces/union types.  In this case the decoding should
-// be done manually with a switch statement, and so the
-// reflection-based decoder provided here does not expect this
-// prefix byte.
-// See the reactor implementations for use-cases.
+// as the first byte for encoding and decoding.
+// This is primarily used to encode interfaces types.
+// The interface should be declared with RegisterInterface()
 type HasTypeByte interface {
 	TypeByte() byte
 }
@@ -45,11 +66,50 @@ func GetTypeInfo(rt reflect.Type) *TypeInfo {
 	return info
 }

+// For use with the RegisterInterface declaration
+type ConcreteType struct {
+	O interface{}
+}
+
+// Must use this to register an interface to properly decode the
+// underlying concrete type.
+func RegisterInterface(o interface{}, args ...interface{}) *TypeInfo {
+	it := GetTypeFromStructDeclaration(o)
+	if it.Kind() != reflect.Interface {
+		panic("RegisterInterface expects an interface")
+	}
+	concreteTypes := make(map[byte]reflect.Type, 0)
+	for _, arg := range args {
+		switch arg.(type) {
+		case ConcreteType:
+			concreteTypeInfo := arg.(ConcreteType)
+			concreteType := reflect.TypeOf(concreteTypeInfo.O)
+			hasTypeByte, typeByte := GetTypeByteFromStruct(concreteTypeInfo.O)
+			//fmt.Println(Fmt("HasTypeByte: %v typeByte: %X type: %X", hasTypeByte, typeByte, concreteType))
+			if !hasTypeByte {
+				panic(Fmt("Expected concrete type %v to implement HasTypeByte", concreteType))
+			}
+			if concreteTypes[typeByte] != nil {
+				panic(Fmt("Duplicate TypeByte for type %v and %v", concreteType, concreteTypes[typeByte]))
+			}
+			concreteTypes[typeByte] = concreteType
+		default:
+			panic(Fmt("Unexpected argument type %v", reflect.TypeOf(arg)))
+		}
+	}
+	typeInfo := &TypeInfo{
+		Type:          it,
+		ConcreteTypes: concreteTypes,
+	}
+	typeInfos[it] = typeInfo
+	return typeInfo
+}
+
 // Registers and possibly modifies the TypeInfo.
 // NOTE: not goroutine safe, so only call upon program init.
 func RegisterType(info *TypeInfo) *TypeInfo {

-	// Also register the underlying struct's info, if info.Type is a pointer.
+	// Also register the dereferenced struct if info.Type is a pointer.
 	// Or, if info.Type is not a pointer, register the pointer.
 	var rt, ptrRt reflect.Type
 	if info.Type.Kind() == reflect.Ptr {
@@ -96,8 +156,6 @@ func readReflect(rv reflect.Value, rt reflect.Type, r Unreader, n *int64, err *e
 	// Custom decoder
 	if typeInfo.Decoder != nil {
 		decoded := typeInfo.Decoder(r, n, err)
-		//decodedRv := reflect.Indirect(reflect.ValueOf(decoded))
-		//rv.Set(decodedRv)
 		rv.Set(reflect.ValueOf(decoded))
 		return
 	}
@@ -126,6 +184,19 @@ func readReflect(rv reflect.Value, rt reflect.Type, r Unreader, n *int64, err *e
 	}

 	switch rt.Kind() {
+	case reflect.Interface:
+		typeByte := PeekByte(r, n, err)
+		if *err != nil {
+			return
+		}
+		concreteType, ok := typeInfo.ConcreteTypes[typeByte]
+		if !ok {
+			panic(Fmt("TypeByte %X not registered for interface %v", typeByte, rt))
+		}
+		newRv := reflect.New(concreteType)
+		readReflect(newRv.Elem(), concreteType, r, n, err)
+		rv.Set(newRv.Elem())
+
 	case reflect.Slice:
 		elemRt := rt.Elem()
 		if elemRt.Kind() == reflect.Uint8 {
@@ -0,0 +1,317 @@
+package binary
+
+import (
+	"bytes"
+	"reflect"
+	"testing"
+)
+
+type SimpleStruct struct {
+	String string
+	Bytes  []byte
+}
+
+//-------------------------------------
+
+type Animal interface{}
+
+const (
+	AnimalTypeCat   = byte(0x01)
+	AnimalTypeDog   = byte(0x02)
+	AnimalTypeSnake = byte(0x03)
+	AnimalTypeViper = byte(0x04)
+)
+
+// Implements Animal
+type Cat struct {
+	SimpleStruct
+}
+
+func (cat Cat) TypeByte() byte { return AnimalTypeCat }
+
+// Implements Animal
+type Dog struct {
+	SimpleStruct
+}
+
+func (dog Dog) TypeByte() byte { return AnimalTypeDog }
+
+// Implements Animal
+type Snake []byte
+
+func (snake Snake) TypeByte() byte { return AnimalTypeSnake }
+
+// Implements Animal
+type Viper struct {
+	Bytes []byte
+}
+
+func (viper *Viper) TypeByte() byte { return AnimalTypeViper }
+
+var _ = RegisterInterface(
+	struct{ Animal }{},
+	ConcreteType{Cat{}},
+	ConcreteType{Dog{}},
+	ConcreteType{Snake{}},
+	ConcreteType{&Viper{}},
+)
+
+//-------------------------------------
+
+func TestBasic(t *testing.T) {
+	cat := Cat{
+		SimpleStruct{
+			String: "String",
+			Bytes:  []byte("Bytes"),
+		},
+	}
+
+	buf, n, err := new(bytes.Buffer), new(int64), new(error)
+	WriteBinary(cat, buf, n, err)
+	if *err != nil {
+		t.Fatalf("Failed to write cat: %v", *err)
+	}
+	t.Logf("Wrote bytes: %X", buf.Bytes())
+	bufBytes := buf.Bytes()
+
+	// Read onto a struct
+	cat2_ := ReadBinary(Cat{}, buf, n, err)
+	if *err != nil {
+		t.Fatalf("Failed to read cat: %v", *err)
+	}
+	cat2 := cat2_.(Cat)
+
+	if cat2.String != "String" {
+		t.Errorf("Expected cat2.String == 'String', got %v", cat2.String)
+	}
+	if string(cat2.Bytes) != "Bytes" {
+		t.Errorf("Expected cat2.Bytes == 'Bytes', got %X", cat2.Bytes)
+	}
+
+	// Read onto a ptr
+	r := bytes.NewReader(bufBytes)
+	cat3_ := ReadBinary(&Cat{}, r, n, err)
+	if *err != nil {
+		t.Fatalf("Failed to read cat: %v", *err)
+	}
+	cat3 := cat3_.(*Cat)
+
+	if cat3.String != "String" {
+		t.Errorf("Expected cat3.String == 'String', got %v", cat3.String)
+	}
+	if string(cat3.Bytes) != "Bytes" {
+		t.Errorf("Expected cat3.Bytes == 'Bytes', got %X", cat3.Bytes)
+	}
+
+}
+
+//-------------------------------------
+
+type ComplexStruct struct {
+	Name   string
+	Animal Animal
+}
+
+func TestComplexStruct(t *testing.T) {
+	c := ComplexStruct{
+		Name: "Complex",
+		Animal: Cat{
+			SimpleStruct{
+				String: "String",
+				Bytes:  []byte("Bytes"),
+			},
+		},
+	}
+
+	buf, n, err := new(bytes.Buffer), new(int64), new(error)
+	WriteBinary(c, buf, n, err)
+	if *err != nil {
+		t.Fatalf("Failed to write c: %v", *err)
+	}
+
+	t.Logf("Wrote bytes: %X", buf.Bytes())
+
+	c2_ := ReadBinary(&ComplexStruct{}, buf, n, err)
+	if *err != nil {
+		t.Fatalf("Failed to read c: %v", *err)
+	}
+	c2 := c2_.(*ComplexStruct)
+
+	if cat, ok := c2.Animal.(Cat); ok {
+		if cat.String != "String" {
+			t.Errorf("Expected cat.String == 'String', got %v", cat.String)
+		}
+		if string(cat.Bytes) != "Bytes" {
+			t.Errorf("Expected cat.Bytes == 'Bytes', got %X", cat.Bytes)
+		}
+	} else {
+		t.Errorf("Expected c2.Animal to be of type cat, got %v", reflect.ValueOf(c2.Animal).Elem().Type())
+	}
+}
+
+//-------------------------------------
+
+type ComplexArrayStruct struct {
+	Animals []Animal
+}
+
+func TestComplexArrayStruct(t *testing.T) {
+	c := ComplexArrayStruct{
+		Animals: []Animal{
+			Cat{
+				SimpleStruct{
+					String: "String",
+					Bytes:  []byte("Bytes"),
+				},
+			},
+			&Dog{ // Even though it's a *Dog, we'll get a Dog{} back.
+				SimpleStruct{
+					String: "Woof",
+					Bytes:  []byte("Bark"),
+				},
+			},
+			Snake([]byte("hiss")),
+			&Viper{
+				Bytes: []byte("hizz"),
+			},
+		},
+	}
+
+	buf, n, err := new(bytes.Buffer), new(int64), new(error)
+	WriteBinary(c, buf, n, err)
+	if *err != nil {
+		t.Fatalf("Failed to write c: %v", *err)
+	}
+
+	t.Logf("Wrote bytes: %X", buf.Bytes())
+
+	c2_ := ReadBinary(&ComplexArrayStruct{}, buf, n, err)
+	if *err != nil {
+		t.Fatalf("Failed to read c: %v", *err)
+	}
+	c2 := c2_.(*ComplexArrayStruct)
+
+	if cat, ok := c2.Animals[0].(Cat); ok {
+		if cat.String != "String" {
+			t.Errorf("Expected cat.String == 'String', got %v", cat.String)
+		}
+		if string(cat.Bytes) != "Bytes" {
+			t.Errorf("Expected cat.Bytes == 'Bytes', got %X", cat.Bytes)
+		}
+	} else {
+		t.Errorf("Expected c2.Animals[0] to be of type cat, got %v", reflect.ValueOf(c2.Animals[0]).Elem().Type())
+	}
+
+	if dog, ok := c2.Animals[1].(Dog); ok {
+		if dog.String != "Woof" {
+			t.Errorf("Expected dog.String == 'Woof', got %v", dog.String)
+		}
+		if string(dog.Bytes) != "Bark" {
+			t.Errorf("Expected dog.Bytes == 'Bark', got %X", dog.Bytes)
+		}
+	} else {
+		t.Errorf("Expected c2.Animals[1] to be of type dog, got %v", reflect.ValueOf(c2.Animals[1]).Elem().Type())
+	}
+
+	if snake, ok := c2.Animals[2].(Snake); ok {
+		if string(snake) != "hiss" {
+			t.Errorf("Expected string(snake) == 'hiss', got %v", string(snake))
+		}
+	} else {
+		t.Errorf("Expected c2.Animals[2] to be of type Snake, got %v", reflect.ValueOf(c2.Animals[2]).Elem().Type())
+	}
+
+	if viper, ok := c2.Animals[3].(*Viper); ok {
+		if string(viper.Bytes) != "hizz" {
+			t.Errorf("Expected string(viper.Bytes) == 'hizz', got %v", string(viper.Bytes))
+		}
+	} else {
+		t.Errorf("Expected c2.Animals[3] to be of type *Viper, got %v", reflect.ValueOf(c2.Animals[3]).Elem().Type())
+	}
+}
+
+//-------------------------------------
+
+type ComplexStruct2 struct {
+	Cat    Cat
+	Dog    *Dog
+	Snake  Snake
+	Snake2 *Snake
+	Viper  Viper
+	Viper2 *Viper
+}
+
+func TestComplexStruct2(t *testing.T) {
+
+	snake_ := Snake([]byte("hiss"))
+	snakePtr_ := &snake_
+
+	c := ComplexStruct2{
+		Cat: Cat{
+			SimpleStruct{
+				String: "String",
+				Bytes:  []byte("Bytes"),
+			},
+		},
+		Dog: &Dog{
+			SimpleStruct{
+				String: "Woof",
+				Bytes:  []byte("Bark"),
+			},
+		},
+		Snake:  Snake([]byte("hiss")),
+		Snake2: snakePtr_,
+		Viper:  Viper{Bytes: []byte("hizz")},
+		Viper2: &Viper{Bytes: []byte("hizz")},
+	}
+
+	buf, n, err := new(bytes.Buffer), new(int64), new(error)
+	WriteBinary(c, buf, n, err)
+	if *err != nil {
+		t.Fatalf("Failed to write c: %v", *err)
+	}
+
+	t.Logf("Wrote bytes: %X", buf.Bytes())
+
+	c2_ := ReadBinary(&ComplexStruct2{}, buf, n, err)
+	if *err != nil {
+		t.Fatalf("Failed to read c: %v", *err)
+	}
+	c2 := c2_.(*ComplexStruct2)
+
+	cat := c2.Cat
+	if cat.String != "String" {
+		t.Errorf("Expected cat.String == 'String', got %v", cat.String)
+	}
+	if string(cat.Bytes) != "Bytes" {
+		t.Errorf("Expected cat.Bytes == 'Bytes', got %X", cat.Bytes)
+	}
+
+	dog := c2.Dog
+	if dog.String != "Woof" {
+		t.Errorf("Expected dog.String == 'Woof', got %v", dog.String)
+	}
+	if string(dog.Bytes) != "Bark" {
+		t.Errorf("Expected dog.Bytes == 'Bark', got %X", dog.Bytes)
+	}
+
+	snake := c2.Snake
+	if string(snake) != "hiss" {
+		t.Errorf("Expected string(snake) == 'hiss', got %v", string(snake))
+	}
+
+	snake2 := c2.Snake2
+	if string(*snake2) != "hiss" {
+		t.Errorf("Expected string(snake2) == 'hiss', got %v", string(*snake2))
+	}
+
+	viper := c2.Viper
+	if string(viper.Bytes) != "hizz" {
+		t.Errorf("Expected string(viper.Bytes) == 'hizz', got %v", string(viper.Bytes))
+	}
+
+	viper2 := c2.Viper2
+	if string(viper2.Bytes) != "hizz" {
+		t.Errorf("Expected string(viper2.Bytes) == 'hizz', got %v", string(viper2.Bytes))
+	}
+}