From fe6fc3dc7babb14c28c995c979dff927d019e65c Mon Sep 17 00:00:00 2001
From: Tnze <cjd001113@outlook.com>
Date: Tue, 14 Dec 2021 14:12:26 +0800
Subject: [PATCH] Enhance BitStorage

---
 save/bitstorage.go      | 94 +++++++++++++++++++++++++++++------------
 save/bitstorage_test.go | 18 ++++++++
 2 files changed, 85 insertions(+), 27 deletions(-)

diff --git a/save/bitstorage.go b/save/bitstorage.go
index 53ec03d..c9e4f9c 100644
--- a/save/bitstorage.go
+++ b/save/bitstorage.go
@@ -5,47 +5,64 @@ import (
 	"math"
 )
 
-// BitStorage implement the compacted data array used in chunk storage.
-// https://wiki.vg/Chunk_Format
-// This implement the format since Minecraft 1.16
+const indexOutOfBounds = "index out of bounds"
+const valueOutOfBounds = "value out of bounds"
+
+// BitStorage implement the compacted data array used in chunk storage and heightmaps.
+// You can think of this as a []intN whose N is called "bits" in NewBitStorage.
+// For more info, see: https://wiki.vg/Chunk_Format
+// This is implementation of the format since Minecraft 1.16
 type BitStorage struct {
 	data []uint64
 	mask uint64
 
-	bits, size    int
+	bits, length  int
 	valuesPerLong int
 }
 
-// NewBitStorage create a new BitStorage.
-// bits is the number of bits per value.
-// size is the number of values.
-// arrl is optional data for initializing.
-// It's length must match the bits and size if it's not nil.
-func NewBitStorage(bits, size int, arrl []uint64) (b *BitStorage) {
+// NewBitStorage create a new BitStorage. Return nil if bits == 0.
+//
+// The "bits" is the number of bits per value, which can be calculated by math/bits.Len()
+// The "length" is the number of values.
+// The "data" is optional for initializing. Panic if data != nil && len(data) != BitStorageSize(bits, length).
+func NewBitStorage(bits, length int, data []uint64) (b *BitStorage) {
+	if bits == 0 {
+		return nil
+	}
+
 	b = &BitStorage{
 		mask:          1<<bits - 1,
 		bits:          bits,
-		size:          size,
+		length:        length,
 		valuesPerLong: 64 / bits,
 	}
-	dataLen := (size + b.valuesPerLong - 1) / b.valuesPerLong
-	if arrl != nil {
-		if len(arrl) != dataLen {
-			panic(initBitStorageErr{ArrlLen: len(arrl), WantLen: dataLen})
+	dataLen := BitStorageSize(bits, length)
+	if data != nil {
+		if len(data) != dataLen {
+			panic(newBitStorageErr{ArrlLen: len(data), WantLen: dataLen})
 		}
-		b.data = arrl
+		b.data = data
 	} else {
 		b.data = make([]uint64, dataLen)
 	}
 	return
 }
 
-type initBitStorageErr struct {
+// BitStorageSize calculate how many uint64 is needed for given bits and length.
+func BitStorageSize(bits, length int) (size int) {
+	if bits == 0 {
+		return 0
+	}
+	valuesPerLong := 64 / bits
+	return (length + valuesPerLong - 1) / valuesPerLong
+}
+
+type newBitStorageErr struct {
 	ArrlLen int
 	WantLen int
 }
 
-func (i initBitStorageErr) Error() string {
+func (i newBitStorageErr) Error() string {
 	return fmt.Sprintf("invalid length given for storage, got: %d but expected: %d", i.ArrlLen, i.WantLen)
 }
 
@@ -57,10 +74,13 @@ func (b *BitStorage) calcIndex(n int) (c, o int) {
 
 // Swap sets v into [i], and return the previous [i] value.
 func (b *BitStorage) Swap(i, v int) (old int) {
-	if i < 0 || i > b.size-1 ||
-		v < 0 || uint64(v) > b.mask {
-		panic("out of bounds")
+	if b == nil || i < 0 || i > b.length-1 {
+		panic(indexOutOfBounds)
 	}
+	if v < 0 || uint64(v) > b.mask {
+		panic(valueOutOfBounds)
+	}
+
 	c, offset := b.calcIndex(i)
 	l := b.data[c]
 	old = int(l >> offset & b.mask)
@@ -68,12 +88,15 @@ func (b *BitStorage) Swap(i, v int) (old int) {
 	return
 }
 
-// Set sets v into [i]
+// Set sets v into [i].
 func (b *BitStorage) Set(i, v int) {
-	if i < 0 || i > b.size-1 ||
-		v < 0 || uint64(v) > b.mask {
-		panic("out of bounds")
+	if b == nil || i < 0 || i > b.length-1 {
+		panic(indexOutOfBounds)
 	}
+	if v < 0 || uint64(v) > b.mask {
+		panic(valueOutOfBounds)
+	}
+
 	c, offset := b.calcIndex(i)
 	l := b.data[c]
 	b.data[c] = l&(b.mask<<offset^math.MaxUint64) | (uint64(v)&b.mask)<<offset
@@ -81,10 +104,27 @@ func (b *BitStorage) Set(i, v int) {
 
 // Get gets [i] value.
 func (b *BitStorage) Get(i int) int {
-	if i < 0 || i > b.size-1 {
-		panic("out of bounds")
+	if b == nil || i < 0 || i > b.length-1 {
+		panic(indexOutOfBounds)
 	}
+
 	c, offset := b.calcIndex(i)
 	l := b.data[c]
 	return int(l >> offset & b.mask)
 }
+
+// Len is the number of stored values.
+func (b *BitStorage) Len() int {
+	if b == nil {
+		return 0
+	}
+	return b.length
+}
+
+// Longs return the underling array of uint64 for encoding/decoding.
+func (b *BitStorage) Longs() []uint64 {
+	if b == nil {
+		return []uint64{}
+	}
+	return b.data
+}
diff --git a/save/bitstorage_test.go b/save/bitstorage_test.go
index 382fb93..ab1bd93 100644
--- a/save/bitstorage_test.go
+++ b/save/bitstorage_test.go
@@ -1,6 +1,8 @@
 package save
 
 import (
+	pk "github.com/Tnze/go-mc/net/packet"
+	"math/bits"
 	"reflect"
 	"testing"
 )
@@ -26,3 +28,19 @@ func TestBitStorage_Set(t *testing.T) {
 		t.Errorf("Encode error, got %v but expected: %v", bs.data, data)
 	}
 }
+
+func ExampleNewBitStorage_heightmaps() {
+	// Create a BitStorage
+	bs := NewBitStorage(bits.Len(256), 16*16, nil)
+	// Fill your data
+	for i := 0; i < 16; i++ {
+		for j := 0; j < 16; j++ {
+			bs.Set(i*16+j, 0)
+		}
+	}
+	// Encode as NBT, and this is ready for packet.Marshal
+	type HeightMaps struct {
+		MotionBlocking []uint64 `nbt:"MOTION_BLOCKING"`
+	}
+	_ = pk.NBT(HeightMaps{bs.Longs()})
+}