From 068b534e7a51b4e0f91314eb3cea7c21b188f432 Mon Sep 17 00:00:00 2001 From: Oleg Kovalov Date: Fri, 18 Aug 2023 18:34:43 +0200 Subject: [PATCH] Add marshaling via reflect (#11) --- encode.go | 87 ++++++++++++++++++++++++++++++++++++++++---------- encode_test.go | 74 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 17 deletions(-) diff --git a/encode.go b/encode.go index 504c582..71b96e6 100644 --- a/encode.go +++ b/encode.go @@ -4,6 +4,7 @@ import ( "bytes" "fmt" "io" + "reflect" "strconv" ) @@ -40,18 +41,18 @@ func (enc *Encoder) marshal(v any) error { } enc.buf.Write(raw) case A: - enc.marshalArray(v) + enc.marshalArray(enc.buf, v) case D: - enc.marshalDoc(v) + enc.marshalDoc(enc.buf, v) case M: - enc.marshalDoc(v.AsD()) + enc.marshalDoc(enc.buf, v.AsD()) default: - return fmt.Errorf("type %T is not supported yet", v) + return enc.marshalReflect(enc.buf, v) } return nil } -func (enc *Encoder) marshalArray(arr A) error { +func (enc *Encoder) marshalArray(w io.Writer, arr A) error { doc := make(D, len(arr)) for i := range arr { doc[i] = e{ @@ -59,47 +60,65 @@ func (enc *Encoder) marshalArray(arr A) error { V: arr[i], } } - return enc.marshalDoc(doc) + return enc.marshalDoc(w, doc) } -func (enc *Encoder) marshalDoc(doc D) error { +func (enc *Encoder) marshalDoc(w io.Writer, doc D) error { // TODO(cristaloleg): prealloc or smarter way. - var elist bytes.Buffer + elist := bytes.NewBuffer(make([]byte, 0, 128)) for i := range doc { - pair := doc[i] key := doc[i].K + val := doc[i].V - switch v := pair.V.(type) { + switch v := val.(type) { case string: - enc.writeKey(&elist, TypeString, key) + enc.writeKey(elist, TypeString, key) b := putUint32(uint32(len(v) + 1)) elist.Write(b[:]) elist.WriteString(v) elist.WriteByte(0) case int32: - enc.writeKey(&elist, TypeInt32, key) + enc.writeKey(elist, TypeInt32, key) b := putUint32(uint32(v)) elist.Write(b[:]) case int64: - enc.writeKey(&elist, TypeInt64, key) + enc.writeKey(elist, TypeInt64, key) b := putUint64(uint64(v)) elist.Write(b[:]) case bool: - enc.writeKey(&elist, TypeBool, key) + enc.writeKey(elist, TypeBool, key) elist.WriteByte(putBool(v)) + + default: + var err error + switch rv := reflect.ValueOf(val); rv.Kind() { + case reflect.Map: + enc.writeKey(elist, TypeDocument, key) + err = enc.marshalMap(elist, rv) + + case reflect.Array, reflect.Slice: + enc.writeKey(elist, TypeArray, key) + err = enc.marshalSlice(elist, rv) + + default: + return fmt.Errorf("type %T is not supported yet", v) + } + if err != nil { + return err + } } } size := 4 + elist.Len() + 1 // header + len + null. b := putUint32(uint32(size)) - enc.buf.Write(b[:]) + w.Write(b[:]) - io.Copy(enc.buf, &elist) - enc.buf.WriteByte(0) + io.Copy(w, elist) + w.Write([]byte{0}) return nil } @@ -108,3 +127,37 @@ func (enc *Encoder) writeKey(buf *bytes.Buffer, t Type, s string) { buf.WriteString(s) buf.WriteByte(0) } + +func (enc *Encoder) marshalReflect(w io.Writer, v any) error { + switch rv := reflect.ValueOf(v); rv.Kind() { + // TODO(cristaloleg): add reflect.Struct + case reflect.Map: + return enc.marshalMap(w, rv) + case reflect.Array, reflect.Slice: + return enc.marshalSlice(w, rv) + default: + return fmt.Errorf("type %T is not supported yet", v) + } +} + +func (enc *Encoder) marshalMap(w io.Writer, v reflect.Value) error { + doc := make(D, v.Len()) + for i, iter := 0, v.MapRange(); iter.Next(); i++ { + doc[i] = e{ + K: iter.Key().String(), + V: iter.Value().Interface(), + } + } + return enc.marshalDoc(w, doc) +} + +func (enc *Encoder) marshalSlice(w io.Writer, v reflect.Value) error { + doc := make(D, v.Len()) + for i := 0; i < v.Len(); i++ { + doc[i] = e{ + K: strconv.Itoa(i), + V: v.Index(i).Interface(), + } + } + return enc.marshalDoc(w, doc) +} diff --git a/encode_test.go b/encode_test.go index a8ed24a..e2a663e 100644 --- a/encode_test.go +++ b/encode_test.go @@ -94,3 +94,77 @@ func TestEncodeM(t *testing.T) { wantBytes(t, buf.Bytes(), "1b0000001061000a000000126200a0af9b00000000000863000100") buf.Reset() } + +func TestEncodeReflectMap(t *testing.T) { + var buf bytes.Buffer + enc := NewEncoder(&buf) + + var err error + var m map[string]any + + m = map[string]any{} + err = enc.Encode(m) + mustOk(t, err) + wantBytes(t, buf.Bytes(), "0500000000") + buf.Reset() + + m = map[string]any{"abc": int32(123)} + err = enc.Encode(m) + mustOk(t, err) + wantBytes(t, buf.Bytes(), "0e00000010616263007b00000000") + buf.Reset() + + m = map[string]any{"hello": "world", "foo": int32(123)} + err = enc.Encode(m) + mustOk(t, err) + wantBytes(t, buf.Bytes(), "1f0000000268656c6c6f0006000000776f726c640010666f6f007b00000000") + buf.Reset() +} + +func TestEncodeReflectSlice(t *testing.T) { + var buf bytes.Buffer + enc := NewEncoder(&buf) + + var err error + var s []string + + s = []string{} + err = enc.Encode(s) + mustOk(t, err) + wantBytes(t, buf.Bytes(), "0500000000") + buf.Reset() + + s = []string{"abc"} + err = enc.Encode(s) + mustOk(t, err) + wantBytes(t, buf.Bytes(), "10000000023000040000006162630000") + buf.Reset() + + s = []string{"hello", "world"} + err = enc.Encode(s) + mustOk(t, err) + wantBytes(t, buf.Bytes(), "1f0000000230000600000068656c6c6f0002310006000000776f726c640000") + buf.Reset() +} + +func TestEncodeReflectArray(t *testing.T) { + var buf bytes.Buffer + enc := NewEncoder(&buf) + + var err error + + err = enc.Encode([0]string{}) + mustOk(t, err) + wantBytes(t, buf.Bytes(), "0500000000") + buf.Reset() + + err = enc.Encode([1]string{"abc"}) + mustOk(t, err) + wantBytes(t, buf.Bytes(), "10000000023000040000006162630000") + buf.Reset() + + err = enc.Encode([...]string{"hello", "world"}) + mustOk(t, err) + wantBytes(t, buf.Bytes(), "1f0000000230000600000068656c6c6f0002310006000000776f726c640000") + buf.Reset() +}