4 files changed, 44 insertions, 1 deletions

diff --git a/src/Generating/PieceGenerator.cpp b/src/Generating/PieceGenerator.cpp
index 1880a20d5..5de231f75 100644
--- a/src/Generating/PieceGenerator.cpp
+++ b/src/Generating/PieceGenerator.cpp
@@ -362,7 +362,31 @@ cPlacedPiece * cPieceGenerator::PlaceStartingPiece(int a_BlockX, int a_BlockY, i
 	
 	// Choose a random one of the starting pieces:
 	cPieces StartingPieces = m_PiecePool.GetStartingPieces();
-	cPiece * StartingPiece = StartingPieces[rnd % StartingPieces.size()];
+	int Total = 0;
+	for (cPieces::const_iterator itr = StartingPieces.begin(), end = StartingPieces.end(); itr != end; ++itr)
+	{
+		Total += m_PiecePool.GetStartingPieceWeight(**itr);
+	}
+	cPiece * StartingPiece;
+	if (Total > 0)
+	{
+		int Chosen = rnd % Total;
+		StartingPiece = StartingPieces.front();
+		for (cPieces::const_iterator itr = StartingPieces.begin(), end = StartingPieces.end(); itr != end; ++itr)
+		{
+			Chosen -= m_PiecePool.GetStartingPieceWeight(**itr);
+			if (Chosen <= 0)
+			{
+				StartingPiece = *itr;
+				break;
+			}
+		}
+	}
+	else
+	{
+		// All pieces returned zero weight, but we need one to start. Choose with equal chance:
+		StartingPiece = StartingPieces[rnd % StartingPieces.size()];
+	}
 	rnd = rnd >> 16;
 	
 	// Choose a random supported rotation:
diff --git a/src/Generating/PieceGenerator.h b/src/Generating/PieceGenerator.h
index 21c155c96..fd8576706 100644
--- a/src/Generating/PieceGenerator.h
+++ b/src/Generating/PieceGenerator.h
@@ -120,6 +120,15 @@ public:
 		const cPiece & a_NewPiece
 	) { return 1; }
 	
+	/** Returns the relative weight with which the a_NewPiece is to be selected for placing as the first piece.
+	This allows the pool to tweak the piece's chances.
+	The higher the number returned, the higher the chance the piece will be chosen. 0 means the piece will not be chosen.
+	If all pieces return 0, a random piece is chosen, with all equal chances.
+	*/
+	virtual int GetStartingPieceWeight(
+		const cPiece & a_NewPiece
+	) { return 1; }
+
 	/** Called after a piece is placed, to notify the pool that it has been used.
 	The pool may adjust the pieces it will return the next time. */
 	virtual void PiecePlaced(const cPiece & a_Piece) = 0;
diff --git a/src/Generating/PrefabPiecePool.cpp b/src/Generating/PrefabPiecePool.cpp
index ed9340815..ad97ab0e7 100644
--- a/src/Generating/PrefabPiecePool.cpp
+++ b/src/Generating/PrefabPiecePool.cpp
@@ -101,6 +101,15 @@ int cPrefabPiecePool::GetPieceWeight(const cPlacedPiece & a_PlacedPiece, const c
 
 
 
+int cPrefabPiecePool::GetStartingPieceWeight(const cPiece & a_NewPiece)
+{
+	return ((const cPrefab &)a_NewPiece).GetDefaultWeight();
+}
+
+
+
+
+
 void cPrefabPiecePool::PiecePlaced(const cPiece & a_Piece)
 {
 	// Do nothing
diff --git a/src/Generating/PrefabPiecePool.h b/src/Generating/PrefabPiecePool.h
index c6a5ad360..695ab4ea5 100644
--- a/src/Generating/PrefabPiecePool.h
+++ b/src/Generating/PrefabPiecePool.h
@@ -70,6 +70,7 @@ protected:
 	virtual cPieces GetPiecesWithConnector(int a_ConnectorType) override;
 	virtual cPieces GetStartingPieces(void) override;
 	virtual int GetPieceWeight(const cPlacedPiece & a_PlacedPiece, const cPiece::cConnector & a_ExistingConnector, const cPiece & a_NewPiece) override;
+	virtual int GetStartingPieceWeight(const cPiece & a_NewPiece) override;
 	virtual void PiecePlaced(const cPiece & a_Piece) override;
 	virtual void Reset(void) override;
 } ;