3 files changed, 29 insertions, 6 deletions

diff --git a/source/Generating/ChunkDesc.cpp b/source/Generating/ChunkDesc.cpp
index dc6c74a3c..039f30d9c 100644
--- a/source/Generating/ChunkDesc.cpp
+++ b/source/Generating/ChunkDesc.cpp
@@ -8,6 +8,7 @@
 #include "../BlockArea.h"
 #include "../Cuboid.h"
 #include "../Noise.h"
+#include "../BlockEntities/BlockEntity.h"
 
 
 
@@ -526,9 +527,28 @@ void cChunkDesc::RandomFillRelCuboid(
 
 
 
-void cChunkDesc::AddBlockEntity(cBlockEntity * a_BlockEntity)
+cBlockEntity * cChunkDesc::GetBlockEntity(int a_RelX, int a_RelY, int a_RelZ)
 {
-	m_BlockEntities.push_back(a_BlockEntity);
+	int AbsX = a_RelX + m_ChunkX * cChunkDef::Width;
+	int AbsZ = a_RelZ + m_ChunkZ * cChunkDef::Width;
+	for (cBlockEntityList::iterator itr = m_BlockEntities.begin(), end = m_BlockEntities.end(); itr != end; ++itr)
+	{
+		if (((*itr)->GetPosX() == AbsX) && ((*itr)->GetPosY() == a_RelY) && ((*itr)->GetPosZ() == AbsZ))
+		{
+			// Already in the list, return it:
+			return *itr;
+		}
+	}  // for itr - m_BlockEntities[]
+	
+	// The block entity is not created yet, try to create it and add to list:
+	cBlockEntity * be = cBlockEntity::CreateByBlockType(GetBlockType(a_RelX, a_RelY, a_RelZ), GetBlockMeta(a_RelX, a_RelY, a_RelZ), AbsX, a_RelY, AbsZ);
+	if (be == NULL)
+	{
+		// No block entity for this block type
+		return NULL;
+	}
+	m_BlockEntities.push_back(be);
+	return be;
 }
 
 
diff --git a/source/Generating/ChunkDesc.h b/source/Generating/ChunkDesc.h
index 067d8494a..e130c463f 100644
--- a/source/Generating/ChunkDesc.h
+++ b/source/Generating/ChunkDesc.h
@@ -170,9 +170,12 @@ public:
 		);
 	}
 	
-	// tolua_end
+	/// Returns the block entity at the specified coords.
+	/// If there is no block entity at those coords, tries to create one, based on the block type
+	/// If the blocktype doesn't support a block entity, returns NULL.
+	cBlockEntity * GetBlockEntity(int a_RelX, int a_RelY, int a_RelZ);
 	
-	void AddBlockEntity(cBlockEntity * a_BlockEntity);
+	// tolua_end
 	
 	// Accessors used by cChunkGenerator::Generator descendants:
 	inline cChunkDef::BiomeMap &     GetBiomeMap              (void) { return m_BiomeMap; }
diff --git a/source/Generating/MineShafts.cpp b/source/Generating/MineShafts.cpp
index 3131b5429..159e6b4ea 100644
--- a/source/Generating/MineShafts.cpp
+++ b/source/Generating/MineShafts.cpp
@@ -794,12 +794,12 @@ void cMineShaftCorridor::PlaceChest(cChunkDesc & a_ChunkDesc)
 	)
 	{
 		a_ChunkDesc.SetBlockTypeMeta(x, m_BoundingBox.p1.y + 1, z, E_BLOCK_CHEST, Meta);
-		cChestEntity * ChestEntity = new cChestEntity(BlockX + x, m_BoundingBox.p1.y + 1, BlockZ + z);
+		cChestEntity * ChestEntity = (cChestEntity *)a_ChunkDesc.GetBlockEntity(x, m_BoundingBox.p1.y + 1, z);
+		ASSERT((ChestEntity != NULL) && (ChestEntity->GetBlockType() == E_BLOCK_CHEST));
 		cNoise Noise(a_ChunkDesc.GetChunkX() ^ a_ChunkDesc.GetChunkZ());
 		int NumSlots = 3 + ((Noise.IntNoise3DInt(x, m_BoundingBox.p1.y, z) / 11) % 4);
 		int Seed = Noise.IntNoise2DInt(x, z);
 		ChestEntity->GetContents().GenerateRandomLootWithBooks(LootProbab, ARRAYCOUNT(LootProbab), NumSlots, Seed);
-		a_ChunkDesc.AddBlockEntity(ChestEntity);
 	}
 }