2 files changed, 41 insertions, 24 deletions

diff --git a/src/Blocks/BlockPiston.cpp b/src/Blocks/BlockPiston.cpp
index 94782a7ed..83a1b83c4 100644
--- a/src/Blocks/BlockPiston.cpp
+++ b/src/Blocks/BlockPiston.cpp
@@ -86,28 +86,38 @@ bool cBlockPistonHandler::GetPlacementBlockTypeMeta(
 
 
 
-int cBlockPistonHandler::FirstPassthroughBlock(int a_PistonX, int a_PistonY, int a_PistonZ, NIBBLETYPE pistonmeta, cWorld * a_World)
+bool cBlockPistonHandler::CanPushBlock(
+	int a_BlockX, int a_BlockY, int a_BlockZ, cWorld * a_World, bool a_RequirePushable,
+	std::unordered_set<Vector3i, VectorHasher<int>> & a_BlocksPushed, NIBBLETYPE a_PistonMeta
+)
 {
-	// Examine each of the 12 blocks ahead of the piston:
-	for (int ret = 0; ret < PISTON_MAX_PUSH_DISTANCE; ret++)
+	BLOCKTYPE currBlock;
+	NIBBLETYPE currMeta;
+	a_World->GetBlockTypeMeta(a_BlockX, a_BlockY, a_BlockZ, currBlock, currMeta);
+	
+	if (currBlock == E_BLOCK_AIR)
 	{
-		BLOCKTYPE currBlock;
-		NIBBLETYPE currMeta;
-		AddPistonDir(a_PistonX, a_PistonY, a_PistonZ, pistonmeta, 1);
-		a_World->GetBlockTypeMeta(a_PistonX, a_PistonY, a_PistonZ, currBlock, currMeta);
-		if (cBlockInfo::IsPistonBreakable(currBlock))
-		{
-			// This block breaks when pushed, extend up to here
-			return ret;
-		}
-		if (!CanPush(currBlock, currMeta))
-		{
-			// This block cannot be pushed at all, the piston can't extend
-			return -1;
-		}
+		// Air can be pushed
+		return true;
+	}
+	
+	if (!CanPush(currBlock, currMeta))
+	{
+		return !a_RequirePushable;
+	}
+
+	if (a_BlocksPushed.size() >= PISTON_MAX_PUSH_DISTANCE)
+	{
+		return false;
 	}
-	// There is no space for the blocks to move, piston can't extend
-	return -1;
+	
+	if (!a_BlocksPushed.emplace(a_BlockX, a_BlockY, a_BlockZ).second || cBlockInfo::IsPistonBreakable(currBlock))
+	{
+		return true;  // Element exist already
+	}
+	
+	AddPistonDir(a_BlockX, a_BlockY, a_BlockZ, a_PistonMeta, 1);
+	return CanPushBlock(a_BlockX, a_BlockY, a_BlockZ, a_World, true, a_BlocksPushed, a_PistonMeta);
 }
 
 
@@ -126,10 +136,12 @@ void cBlockPistonHandler::ExtendPiston(int a_BlockX, int a_BlockY, int a_BlockZ,
 		return;
 	}
 
-	int dist = FirstPassthroughBlock(a_BlockX, a_BlockY, a_BlockZ, pistonMeta, a_World);
-	if (dist < 0)
+	int dist = 1;  // TODO
+	AddPistonDir(a_BlockX, a_BlockY, a_BlockZ, pistonMeta, 1);
+	std::unordered_set<Vector3i, VectorHasher<int>> blocksPushed;
+	if (!CanPushBlock(a_BlockX, a_BlockY, a_BlockZ, a_World, true, blocksPushed, pistonMeta))
 	{
-		// FirstPassthroughBlock says piston can't push anything, bail out
+		// Can't push anything, bail out
 		return;
 	}
 
diff --git a/src/Blocks/BlockPiston.h b/src/Blocks/BlockPiston.h
index 56f7f9951..41ef79aa6 100644
--- a/src/Blocks/BlockPiston.h
+++ b/src/Blocks/BlockPiston.h
@@ -3,6 +3,8 @@
 
 #include "BlockHandler.h"
 
+#include <unordered_set>
+
 
 class cWorld;
 
@@ -152,8 +154,11 @@ private:
 		return CanPush(a_BlockType, a_BlockMeta);
 	}
 
-	/** Returns how many blocks the piston has to push (where the first free space is); < 0 when unpushable */
-	static int FirstPassthroughBlock(int a_PistonX, int a_PistonY, int a_PistonZ, NIBBLETYPE a_PistonMeta, cWorld * a_World);
+	/** Tries to push a block and increases the pushed blocks variable. Returns true if the block is pushable */
+	static bool CanPushBlock(
+		int a_BlockX, int a_BlockY, int a_BlockZ, cWorld * a_World, bool a_RequirePushable,
+		std::unordered_set<Vector3i, VectorHasher<int>> & a_BlocksPushed, NIBBLETYPE a_PistonMeta
+	);
 } ;