From aaabd5e9141e2e51b81c0dacd32e7e849826204d Mon Sep 17 00:00:00 2001
From: Kara <lunarautomaton6@gmail.com>
Date: Mon, 6 May 2024 04:25:11 -0700
Subject: [PATCH] Fix `TryFindRandomTile` grid weighting (#27724)

---
 .../Rules/GameRuleSystem.Utility.cs           | 21 +++++++++++++------
 .../Random/Helpers/SharedRandomExtensions.cs  |  5 +++--
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/Content.Server/GameTicking/Rules/GameRuleSystem.Utility.cs b/Content.Server/GameTicking/Rules/GameRuleSystem.Utility.cs
index 27a9edbad7..cbd981e99e 100644
--- a/Content.Server/GameTicking/Rules/GameRuleSystem.Utility.cs
+++ b/Content.Server/GameTicking/Rules/GameRuleSystem.Utility.cs
@@ -1,7 +1,10 @@
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using Content.Server.GameTicking.Components;
 using Content.Server.GameTicking.Rules.Components;
 using Content.Server.Station.Components;
+using Content.Shared.Random.Helpers;
+using Robust.Server.GameObjects;
 using Robust.Shared.Collections;
 using Robust.Shared.Map;
 using Robust.Shared.Map.Components;
@@ -82,17 +85,23 @@ public abstract partial class GameRuleSystem<T> where T: IComponent
         targetCoords = EntityCoordinates.Invalid;
         targetGrid = EntityUid.Invalid;
 
-        var possibleTargets = station.Comp.Grids;
-        if (possibleTargets.Count == 0)
+        // Weight grid choice by tilecount
+        var weights = new Dictionary<Entity<MapGridComponent>, float>();
+        foreach (var possibleTarget in station.Comp.Grids)
+        {
+            if (!TryComp<MapGridComponent>(possibleTarget, out var comp))
+                continue;
+
+            weights.Add((possibleTarget, comp), _map.GetAllTiles(possibleTarget, comp).Count());
+        }
+
+        if (weights.Count == 0)
         {
             targetGrid = EntityUid.Invalid;
             return false;
         }
 
-        targetGrid = RobustRandom.Pick(possibleTargets);
-
-        if (!TryComp<MapGridComponent>(targetGrid, out var gridComp))
-            return false;
+        (targetGrid, var gridComp) = RobustRandom.Pick(weights);
 
         var found = false;
         var aabb = gridComp.LocalAABB;
diff --git a/Content.Shared/Random/Helpers/SharedRandomExtensions.cs b/Content.Shared/Random/Helpers/SharedRandomExtensions.cs
index dcdc9d61d9..20e57e9421 100644
--- a/Content.Shared/Random/Helpers/SharedRandomExtensions.cs
+++ b/Content.Shared/Random/Helpers/SharedRandomExtensions.cs
@@ -57,7 +57,8 @@ namespace Content.Shared.Random.Helpers
             throw new InvalidOperationException($"Invalid weighted pick for {prototype.ID}!");
         }
 
-        public static string Pick(this IRobustRandom random, Dictionary<string, float> weights)
+        public static T Pick<T>(this IRobustRandom random, Dictionary<T, float> weights)
+            where T: notnull
         {
             var sum = weights.Values.Sum();
             var accumulated = 0f;
@@ -74,7 +75,7 @@ namespace Content.Shared.Random.Helpers
                 }
             }
 
-            throw new InvalidOperationException($"Invalid weighted pick");
+            throw new InvalidOperationException("Invalid weighted pick");
         }
 
         public static (string reagent, FixedPoint2 quantity) Pick(this WeightedRandomFillSolutionPrototype prototype, IRobustRandom? random = null)