improved world data reading, still broken

treepuncher/world/chunk.py

@@ -24,15 +24,23 @@ class BitStream:
 	def read(self, size:int) -> int:
 		if len(self) < size:
 			raise ValueError(f"Not enough bits ({len(self)} left, {size} requested)")
+		# Calculate splice indexes
 		start_byte = (self.cursor//8)
-		end_byte = math.ceil((self.cursor + size) / 8) + 1
+		end_byte = math.ceil((self.cursor + size) / 8)
+		# Construct int from bytes
 		buf = int.from_bytes(
 			self.data[start_byte:end_byte],
 			byteorder='little', signed=False
 		)
-		cut_right = 8 - ((self.cursor + size) % 8)
+		# Trim extra bytes
+		end_offset = (self.cursor + size) % 8
+		if end_offset > 0:
+			buf = buf >> (8 - end_offset) # There's an extra 1 to the left in air, maybe shift 1 bit less?
+		start_offset = self.cursor % 8
+		buf = buf & (( 1 << size ) - 1)
+		# Increment and return
 		self.cursor += size
-		return ( buf >> cut_right ) & ( 0xFF >> (8 - (size%8)))
+		return buf
 
 class PalettedContainer(Type):
 	pytype : type
@@ -47,29 +55,56 @@ class PalettedContainer(Type):
 		raise NotImplementedError
 
 	def read(self, buffer:io.BytesIO, ctx:Context):
-		bits = max(UnsignedByte.read(buffer, ctx=ctx), 4)
+		bits = UnsignedByte.read(buffer, ctx=ctx)
-		if bits > 13:
+		logging.info("Bits per block : %d", bits)
-			raise ValueError("Bits Per Bit too high : %d", bits)
+		if bits < 4:
-		palette = np.empty((0,), dtype='int32')
+			bits = 4
+		if bits >= self.threshold:
+			bits = 13 # this should not be hardcoded but we have no way to calculate all possible block states
 		palette_len = VarInt.read(buffer, ctx=ctx)
-		if bits < self.threshold:
 		palette = np.zeros((palette_len,), dtype='int32')
 		for i in range(palette_len):
 			palette[i] = VarInt.read(buffer, ctx=ctx)
-		size = VarInt.read(buffer, ctx=ctx)
+		container_size = VarInt.read(buffer, ctx=ctx)
-		stream = BitStream(buffer.read(size * 8), size*8*8) # a Long is 64 bits long
+		stream = BitStream(buffer.read(container_size * 8), container_size*8*8) # a Long is 64 bits long
 		section = np.zeros((self.size, self.size, self.size), dtype='int32')
-		index = 0
 		for y in range(self.size):
 			for z in range(self.size):
 				for x in range(self.size):
 					val = stream.read(bits)
+					if bits > 4:
+						if val >= len(palette):
+							logging.warning("out of bounds : %d (%d)", val, len(palette))
+							section[x, y, z] = val
+							continue
+						logging.info("Reading index when bits > 4")
 					section[x, y, z] = palette[val] if bits < self.threshold else val
 		return section
 
 BiomeContainer = PalettedContainer(4, 4)
 BlockStateContainer = PalettedContainer(9, 16)
 
+class HalfByteArrayType(Type):
+	size : int
+
+	def __init__(self, size:int):
+		self.size = size
+
+	def write(self, data, buffer:io.BytesIO, ctx:Context):
+		raise NotImplementedError
+
+	def read(self, buffer:io.BytesIO, ctx:Context):
+		section = np.empty((self.size, self.size, self.size), dtype='int32')
+		bit_buffer = BitStream(buffer.read((self.size**3)//2), (self.size**3)*4)
+		for y in range(self.size):
+			for z in range(self.size):
+				for x in range(self.size):
+					section[x, y, z] = bit_buffer.read(4)
+		return section
+
+BlockLightSection = HalfByteArrayType(16)
+SkyLightSection = HalfByteArrayType(16)
+
 class NewChunkSectionType(Type):
 	pytype : type
 
@@ -93,22 +128,14 @@ class OldChunkSectionType(Type):
 		raise NotImplementedError
 
 	def read(self, buffer:io.BytesIO, ctx:Context):
-		section = BlockStateContainer.read(buffer, ctx=ctx)
+		block_states = BlockStateContainer.read(buffer, ctx=ctx)
-		block_light = np.empty((16, 16, 16), dtype='int32')
+		block_light = BlockLightSection.read(buffer, ctx=ctx)
-		block_light_buffer = BitStream(buffer.read(2048), 2048*8)
-		for y in range(16):
-			for z in range(16):
-				for x in range(16):
-					block_light[x, y, z] = block_light_buffer.read(4)
-		sky_light = np.empty((16, 16, 16), dtype='int32')
 		if ctx.overworld:
-			sky_light_buffer = BitStream(buffer.read(2048), 2048*8)
+			sky_light = SkyLightSection.read(buffer, ctx=ctx)
-			for y in range(16):
+		else:
-				for z in range(16):
+			sky_light = np.empty((16, 16, 16), dtype='int32')
-					for x in range(16):
-						sky_light[x, y, z] = sky_light_buffer.read(4)
 		return (
-			section,
+			block_states,
 			block_light,
 			sky_light
 		)
@@ -142,12 +169,15 @@ class Chunk(Type):
 		logging.info("Reading chunk")
 		for i in range(16):
 			if (self.bitmask >> i) & 1:
-				section, block_light, sky_light = ChunkSection.read(buffer, ctx=ctx)
+				logging.info("Reading section #%d", i)
-				self.blocks[:, i*16 : (i+1)*16, :] = section
+				block_states, block_light, sky_light = ChunkSection.read(buffer, ctx=ctx)
+				self.blocks[:, i*16 : (i+1)*16, :] = block_states
 				self.block_light[:, i*16 : (i+1)*16, :] = block_light
 				self.sky_light[:, i*16 : (i+1)*16, :] = sky_light
 		if self.ground_up_continuous:
 			self.biomes = buffer.read(256) # 16x16
+		if buffer.read():
+			logging.warning("Leftover data in chunk buffer")
 		return self
 
 class World: