Add marshaling via reflect (#11)

encode.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"reflect"
 	"strconv"
 )
 
@@ -40,66 +41,84 @@ 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{
 			K: strconv.Itoa(i),
 			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)
+}

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()
+}