From 65d3a2687a42ed7bfe158b58784ab35c29f3aab4 Mon Sep 17 00:00:00 2001
From: Tnze <cjd001113@outlook.com>
Date: Sun, 20 Nov 2022 23:21:53 +0800
Subject: [PATCH] fix PaletteContainer bug

---
 level/palette.go      | 26 +++++++++++-------------
 level/palette_test.go | 46 ++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 54 insertions(+), 18 deletions(-)

diff --git a/level/palette.go b/level/palette.go
index 516c467..214f709 100644
--- a/level/palette.go
+++ b/level/palette.go
@@ -1,11 +1,12 @@
 package level
 
 import (
-	"github.com/Tnze/go-mc/level/block"
-	pk "github.com/Tnze/go-mc/net/packet"
 	"io"
 	"math/bits"
 	"strconv"
+
+	"github.com/Tnze/go-mc/level/block"
+	pk "github.com/Tnze/go-mc/net/packet"
 )
 
 type State interface {
@@ -102,30 +103,25 @@ func (p *PaletteContainer[T]) Set(i int, v T) {
 	if vv, ok := p.palette.id(v); ok {
 		p.data.Set(i, vv)
 	} else {
+		length := p.data.Len()
 		// resize
-		oldLen := p.data.Len()
-		newPalette := PaletteContainer[T]{
+		newContainer := PaletteContainer[T]{
 			bits:    vv,
 			config:  p.config,
 			palette: p.config.create(vv),
-			data:    NewBitStorage(p.config.bits(vv), oldLen, nil),
+			data:    NewBitStorage(p.config.bits(vv), length, nil),
 		}
 		// copy
-		for i := 0; i < oldLen; i++ {
-			raw := p.data.Get(i)
-			if vv, ok := newPalette.palette.id(T(raw)); !ok {
-				panic("not reachable")
-			} else {
-				newPalette.data.Set(i, vv)
-			}
+		for i := 0; i < length; i++ {
+			newContainer.Set(i, p.Get(i))
 		}
 
-		if vv, ok := newPalette.palette.id(v); !ok {
+		if vv, ok := newContainer.palette.id(v); !ok {
 			panic("not reachable")
 		} else {
-			newPalette.data.Set(i, vv)
+			newContainer.data.Set(i, vv)
 		}
-		*p = newPalette
+		*p = newContainer
 	}
 }
 
diff --git a/level/palette_test.go b/level/palette_test.go
index 501a6fe..689a164 100644
--- a/level/palette_test.go
+++ b/level/palette_test.go
@@ -1,12 +1,14 @@
 package level
 
 import (
+	"math/rand"
 	"testing"
+
+	"github.com/Tnze/go-mc/level/block"
 )
 
-func TestPaletteResize(t *testing.T) {
-	container := NewStatesPaletteContainer(16*16*16, 0)
-
+func TestPaletteContainer_seq(t *testing.T) {
+	container := NewStatesPaletteContainer(4096, 0)
 	for i := 0; i < 4096; i++ {
 		container.Set(i, BlocksState(i))
 	}
@@ -16,3 +18,41 @@ func TestPaletteResize(t *testing.T) {
 		}
 	}
 }
+
+func TestPaletteContainer_rand(t *testing.T) {
+	data := make([]BlocksState, 4096)
+	for i := range data {
+		data[i] = BlocksState(rand.Intn(1 << block.BitsPerBlock))
+	}
+	container := NewStatesPaletteContainer(4096, 0)
+	for i, v := range data {
+		container.Set(i, v)
+	}
+	for i, v := range data {
+		if v2 := container.Get(i); v != v2 {
+			t.Errorf("value not match: got %v, except: %v", v2, v)
+		}
+	}
+}
+
+func BenchmarkPaletteContainer(b *testing.B) {
+	data := make([]BlocksState, 4096)
+	for i := range data {
+		data[i] = BlocksState(rand.Intn(1 << block.BitsPerBlock))
+	}
+	rand.Shuffle(len(data), func(i, j int) { data[i], data[j] = data[j], data[i] })
+	container := NewStatesPaletteContainer(4096, 0)
+	b.ResetTimer()
+	b.Run("Set", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			index := i % 4096
+			container.Set(index, data[index])
+		}
+	})
+	b.Run("Get", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			index := i % 4096
+			container.Get(index)
+		}
+	})
+}