(svn r23512) -Change [FS#4872]: Allow to place locks also on river rapids and restore rivers, if locks are deleted

2 files changed, 13 insertions, 9 deletions

diff --git a/src/water_cmd.cpp b/src/water_cmd.cpp
index 5e18908a1..a62ae7040 100644
--- a/src/water_cmd.cpp
+++ b/src/water_cmd.cpp
@@ -246,6 +246,7 @@ static CommandCost DoBuildLock(TileIndex tile, DiagDirection dir, DoCommandFlag
 	if (ret.Failed()) return ret;
 
 	/* middle tile */
+	WaterClass wc_middle = IsWaterTile(tile) ? GetWaterClass(tile) : WATER_CLASS_CANAL;
 	ret = DoCommand(tile, 0, 0, flags, CMD_LANDSCAPE_CLEAR);
 	if (ret.Failed()) return ret;
 	cost.AddCost(ret);
@@ -295,7 +296,7 @@ static CommandCost DoBuildLock(TileIndex tile, DiagDirection dir, DoCommandFlag
 			DirtyCompanyInfrastructureWindows(_current_company);
 		}
 
-		MakeLock(tile, _current_company, dir, wc_lower, wc_upper);
+		MakeLock(tile, _current_company, dir, wc_lower, wc_upper, wc_middle);
 		MarkTileDirtyByTile(tile);
 		MarkTileDirtyByTile(tile - delta);
 		MarkTileDirtyByTile(tile + delta);
@@ -336,9 +337,14 @@ static CommandCost RemoveLock(TileIndex tile, DoCommandFlag flags)
 			DirtyCompanyInfrastructureWindows(c->index);
 		}
 
-		DoClearSquare(tile);
+		if (GetWaterClass(tile) == WATER_CLASS_RIVER) {
+			MakeRiver(tile, Random());
+		} else {
+			DoClearSquare(tile);
+		}
 		MakeWaterKeepingClass(tile + delta, GetTileOwner(tile + delta));
 		MakeWaterKeepingClass(tile - delta, GetTileOwner(tile - delta));
+		MarkCanalsAndRiversAroundDirty(tile);
 		MarkCanalsAndRiversAroundDirty(tile - delta);
 		MarkCanalsAndRiversAroundDirty(tile + delta);
 	}
@@ -360,9 +366,6 @@ CommandCost CmdBuildLock(TileIndex tile, DoCommandFlag flags, uint32 p1, uint32
 	DiagDirection dir = GetInclinedSlopeDirection(GetTileSlope(tile));
 	if (dir == INVALID_DIAGDIR) return_cmd_error(STR_ERROR_LAND_SLOPED_IN_WRONG_DIRECTION);
 
-	/* Disallow building of locks on river rapids */
-	if (IsWaterTile(tile)) return_cmd_error(STR_ERROR_SITE_UNSUITABLE);
-
 	return DoBuildLock(tile, dir, flags);
 }
 
diff --git a/src/water_map.h b/src/water_map.h
index 14b7c1329..cfec9cc38 100644
--- a/src/water_map.h
+++ b/src/water_map.h
@@ -466,14 +466,15 @@ static inline void MakeLockTile(TileIndex t, Owner o, LockPart part, DiagDirecti
  * @param d Direction of the water lock.
  * @param wc_lower Original water class of the lower part.
  * @param wc_upper Original water class of the upper part.
+ * @param wc_middle Original water class of the middle part.
  */
-static inline void MakeLock(TileIndex t, Owner o, DiagDirection d, WaterClass wc_lower, WaterClass wc_upper)
+static inline void MakeLock(TileIndex t, Owner o, DiagDirection d, WaterClass wc_lower, WaterClass wc_upper, WaterClass wc_middle)
 {
 	TileIndexDiff delta = TileOffsByDiagDir(d);
 
-	MakeLockTile(t, o, LOCK_PART_MIDDLE, d, WATER_CLASS_CANAL);
-	/* Keep the current owner for the upper and lower part if it is a
-	 * water tile so we can restore the owner after deleting the lock. */
+	/* Keep the current waterclass and owner for the tiles.
+	 * It allows to restore them after the lock is deleted */
+	MakeLockTile(t, o, LOCK_PART_MIDDLE, d, wc_middle);
 	MakeLockTile(t - delta, IsWaterTile(t - delta) ? GetTileOwner(t - delta) : o, LOCK_PART_LOWER, d, wc_lower);
 	MakeLockTile(t + delta, IsWaterTile(t + delta) ? GetTileOwner(t + delta) : o, LOCK_PART_UPPER, d, wc_upper);
 }