support encoding map #98

nbt/decode_test.go

@@ -35,6 +35,7 @@ func TestUnmarshal_string(t *testing.T) {
 		t.Errorf("Unmarshal NBT fail: get %q, want %q", Name, "Bananrama")
 	}
 }
+
 func TestUnmarshal_simple(t *testing.T) {
 	var data = []byte{
 		0x0a, 0x00, 0x0b, 0x68, 0x65, 0x6c, 0x6c, 0x6f,
@@ -218,21 +219,21 @@ func TestUnmarshal_bigTest(t *testing.T) {
 	// t.Log(inf)
 }
 
-func TestMarshal_bigTest(t *testing.T) {
-	var b bytes.Buffer
-	err := MarshalCompound(&b, MakeBigTestStruct(), "Level")
-	if err != nil {
-		t.Error(err)
+func BenchmarkUnmarshal_bigTest(b *testing.B) {
+	var value BigTestStruct
+	for i := 0; i < b.N; i++ {
+		r, err := gzip.NewReader(bytes.NewReader(bigTestData[:]))
+		if err != nil {
+			b.Fatal(err)
+		}
+		if err := NewDecoder(r).Decode(&value); err != nil {
+			b.Fatal(err)
+		}
 	}
 
-	rd, _ := gzip.NewReader(bytes.NewReader(bigTestData[:]))
-	want, err := io.ReadAll(rd)
-	if err != nil {
-		t.Error(err)
-	}
-
-	if !bytes.Equal(b.Bytes(), want) {
-		t.Errorf("got:\n[% 2x]\nwant:\n[% 2x]", b.Bytes(), want)
+	want := MakeBigTestStruct()
+	if !reflect.DeepEqual(value, want) {
+		b.Errorf("parse fail, expect %v, get %v", want, value)
 	}
 }
 
@@ -330,6 +331,7 @@ func TestUnmarshal_LongArray(t *testing.T) {
 	}
 	// t.Log(infValue)
 }
+
 func TestUnmarshal_ByteArray(t *testing.T) {
 	data := []byte{
 		TagByteArray, 0, 0,

nbt/encode.go

@@ -2,9 +2,11 @@ package nbt
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"math"
 	"reflect"
+	"strconv"
 	"strings"
 )
 
@@ -58,7 +60,7 @@ func (e *Encoder) writeHeader(val reflect.Value, tagType byte, tagName string) (
 func (e *Encoder) writeValue(val reflect.Value, tagType byte) error {
 	switch tagType {
 	default:
-		return errors.New("unsupported type " + val.Type().Kind().String())
+		return errors.New("unsupported type 0x" + strconv.FormatUint(uint64(tagType), 16))
 	case TagByte:
 		_, err := e.w.Write([]byte{byte(val.Uint())})
 		return err
@@ -114,24 +116,46 @@ func (e *Encoder) writeValue(val reflect.Value, tagType byte) error {
 		return err
 
 	case TagCompound:
-		if val.Kind() == reflect.Interface {
-			val = reflect.ValueOf(val.Interface())
+		for val.Kind() == reflect.Interface {
+			val = val.Elem()
 		}
 
-		n := val.NumField()
-		for i := 0; i < n; i++ {
-			f := val.Type().Field(i)
-			tag := f.Tag.Get("nbt")
-			if (f.PkgPath != "" && !f.Anonymous) || tag == "-" {
-				continue // Private field
-			}
+		switch val.Kind() {
+		case reflect.Struct:
+			n := val.NumField()
+			for i := 0; i < n; i++ {
+				f := val.Type().Field(i)
+				tag := f.Tag.Get("nbt")
+				if (f.PkgPath != "" && !f.Anonymous) || tag == "-" {
+					continue // Private field
+				}
 
-			tagProps := parseTag(f, tag)
-			err := e.marshal(val.Field(i), tagProps.Type, tagProps.Name)
-			if err != nil {
-				return err
+				tagProps := parseTag(f, tag)
+				if err := e.marshal(val.Field(i), tagProps.Type, tagProps.Name); err != nil {
+					return err
+				}
+			}
+		case reflect.Map:
+			r := val.MapRange()
+			for r.Next() {
+				var tagName string
+				if tn, ok := r.Key().Interface().(fmt.Stringer); ok {
+					tagName = tn.String()
+				} else {
+					tagName = r.Key().String()
+				}
+				tagValue := r.Value()
+				tagType := getTagType(tagValue.Type())
+				if tagType == TagNone {
+					return errors.New("unsupported value " + tagValue.String())
+				}
+
+				if err := e.marshal(tagValue, tagType, tagName); err != nil {
+					return err
+				}
 			}
 		}
+
 		_, err := e.w.Write([]byte{TagEnd})
 		return err
 	}
@@ -154,7 +178,7 @@ func getTagType(vk reflect.Type) byte {
 		return TagDouble
 	case reflect.String:
 		return TagString
-	case reflect.Struct, reflect.Interface:
+	case reflect.Struct, reflect.Interface, reflect.Map:
 		return TagCompound
 	case reflect.Array, reflect.Slice:
 		switch vk.Elem().Kind() {

nbt/encode_test.go

@@ -2,7 +2,10 @@ package nbt
 
 import (
 	"bytes"
+	"compress/gzip"
+	"io"
 	"math"
+	"reflect"
 	"testing"
 )
 
@@ -166,3 +169,48 @@ func TestMarshal_StructArray(t *testing.T) {
 		})
 	}
 }
+
+func TestMarshal_bigTest(t *testing.T) {
+	var b bytes.Buffer
+	err := MarshalCompound(&b, MakeBigTestStruct(), "Level")
+	if err != nil {
+		t.Error(err)
+	}
+
+	rd, _ := gzip.NewReader(bytes.NewReader(bigTestData[:]))
+	want, err := io.ReadAll(rd)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if !bytes.Equal(b.Bytes(), want) {
+		t.Errorf("got:\n[% 2x]\nwant:\n[% 2x]", b.Bytes(), want)
+	}
+}
+
+func TestMarshal_map(t *testing.T) {
+	v := map[string][]int32{
+		"Tnze":     {1, 2, 3, 4, 5},
+		"Xi_Xi_Mi": {0, 0, 4, 7, 2},
+	}
+
+	var buf bytes.Buffer
+	if err := Marshal(&buf, v); err != nil {
+		t.Fatal(err)
+	}
+
+	var data struct {
+		Tnze []int32
+		XXM  []int32 `nbt:"Xi_Xi_Mi"`
+	}
+
+	if err := NewDecoder(&buf).Decode(&data); err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(data.Tnze, v["Tnze"]) {
+		t.Fatalf("Marshal map error: got: %q, want %q", data.Tnze, v["Tnze"])
+	}
+	if !reflect.DeepEqual(data.XXM, v["Xi_Xi_Mi"]) {
+		t.Fatalf("Marshal map error: got: %#v, want %#v", data.XXM, v["Xi_Xi_Mi"])
+	}
+}