From cf150c16fa2ee29c293a8ba1be47fcfe9fd4bbaa Mon Sep 17 00:00:00 2001
From: Tnze <cjd001113@outlook.com>
Date: Wed, 8 Apr 2020 16:02:01 +0800
Subject: [PATCH] support marshal []int16, []float32, []float64 and []string.
 More work needed for #50

---
 nbt/marshal.go      | 75 +++++++++++++++++++++++++++++++++++++++++++++
 nbt/marshal_test.go | 23 ++++++++++++--
 2 files changed, 95 insertions(+), 3 deletions(-)

diff --git a/nbt/marshal.go b/nbt/marshal.go
index d19d336..a3da95e 100644
--- a/nbt/marshal.go
+++ b/nbt/marshal.go
@@ -106,6 +106,81 @@ func (e *Encoder) marshal(val reflect.Value, tagName string) error {
 				}
 			}
 
+		case reflect.Int16:
+			if err := e.writeTag(TagList, tagName); err != nil {
+				return err
+			}
+			if _, err := e.w.Write([]byte{TagShort}); err != nil {
+				return err
+			}
+			n := val.Len()
+			if err := e.writeInt32(int32(n)); err != nil {
+				return err
+			}
+			for i := 0; i < n; i++ {
+				if err := e.writeInt16(int16(val.Index(i).Int())); err != nil {
+					return err
+				}
+			}
+
+		case reflect.Float32:
+			if err := e.writeTag(TagList, tagName); err != nil {
+				return err
+			}
+			if _, err := e.w.Write([]byte{TagFloat}); err != nil {
+				return err
+			}
+			n := val.Len()
+			if err := e.writeInt32(int32(n)); err != nil {
+				return err
+			}
+			for i := 0; i < n; i++ {
+				if err := e.writeInt32(int32(math.Float32bits(float32(val.Index(i).Float())))); err != nil {
+					return err
+				}
+			}
+
+		case reflect.Float64:
+			if err := e.writeTag(TagList, tagName); err != nil {
+				return err
+			}
+			if _, err := e.w.Write([]byte{TagFloat}); err != nil {
+				return err
+			}
+			n := val.Len()
+			if err := e.writeInt32(int32(n)); err != nil {
+				return err
+			}
+			for i := 0; i < n; i++ {
+				if err := e.writeInt64(int64(math.Float64bits(val.Index(i).Float()))); err != nil {
+					return err
+				}
+			}
+
+		case reflect.String:
+			if err := e.writeTag(TagList, tagName); err != nil {
+				return err
+			}
+			if _, err := e.w.Write([]byte{TagString}); err != nil {
+				return err
+			}
+			n := val.Len()
+			// Write length of strings
+			if err := e.writeInt32(int32(n)); err != nil {
+				return err
+			}
+			for i := 0; i < n; i++ {
+				// Write length of this string
+				s := val.Index(i).String()
+				if err := e.writeInt16(int16(len(s))); err != nil {
+					return err
+				}
+				// Write string
+				if _, err := e.w.Write([]byte(s)); err != nil {
+					return err
+				}
+			}
+
 		default:
 			return errors.New("unknown type " + val.Type().String() + " slice")
 		}
diff --git a/nbt/marshal_test.go b/nbt/marshal_test.go
index 59b10f4..b1db812 100644
--- a/nbt/marshal_test.go
+++ b/nbt/marshal_test.go
@@ -2,6 +2,7 @@ package nbt
 
 import (
 	"bytes"
+	"math"
 	"testing"
 )
 
@@ -26,9 +27,9 @@ func TestMarshal_IntArray(t *testing.T) {
 	}{[]int32{0, -10, 3}}
 	out = []byte{TagCompound, 0x00, 0x00,
 		TagIntArray, 0x00, 0x03, 'a', 'r', 'y', 0, 0, 0, 3,
-		0x00, 0x00, 0x00, 0x00,
-		0xff, 0xff, 0xff, 0xf6,
-		0x00, 0x00, 0x00, 0x03,
+		0x00, 0x00, 0x00, 0x00, // 0
+		0xff, 0xff, 0xff, 0xf6, // -10
+		0x00, 0x00, 0x00, 0x03, // 3
 		TagEnd,
 	}
 	buf.Reset()
@@ -38,3 +39,19 @@ func TestMarshal_IntArray(t *testing.T) {
 		t.Errorf("output binary not right: get % 02x, want % 02x ", buf.Bytes(), out)
 	}
 }
+
+func TestMarshal_FloatArray(t *testing.T) {
+	// Test marshal pure Int array
+	v := []float32{0.3, -100, float32(math.NaN())}
+	out := []byte{TagList, 0x00, 0x00, TagFloat, 0, 0, 0, 3,
+		0x3e, 0x99, 0x99, 0x9a, // 0.3
+		0xc2, 0xc8, 0x00, 0x00, // -100
+		0x7f, 0xc0, 0x00, 0x00, // NaN
+	}
+	var buf bytes.Buffer
+	if err := Marshal(&buf, v); err != nil {
+		t.Error(err)
+	} else if !bytes.Equal(buf.Bytes(), out) {
+		t.Errorf("output binary not right: get % 02x, want % 02x ", buf.Bytes(), out)
+	}
+}