From e8c6b936617f6111ca87bb5728a9f58e7c7e8ea9 Mon Sep 17 00:00:00 2001
From: DevelopingTom <tom@rouillard.org>
Date: Mon, 28 Jul 2025 06:51:02 +0200
Subject: [PATCH] Improve UnitGrid performances by grouping units by type
 (#1592)

## Description:

Even with the internal grid, `UnitGrid::nearbyUnits` can consume a
significant amount of CPU, where it can account for ~5% of the rendering
thread. Most of this overhead comes from filtering units by type.

This PR improves performance by storing units in type-specific sets,
significantly reducing lookup time.

It reduces the usage to ~1% with this change:
![Sans
titre](https://github.com/user-attachments/assets/f206c2bb-4cdb-4b6e-a628-91b1ec7f56e3)

## Please complete the following:

- [x] I have added screenshots for all UI updates
- [x] I process any text displayed to the user through translateText()
and I've added it to the en.json file
- [x] I have added relevant tests to the test directory
- [x] I confirm I have thoroughly tested these changes and take full
responsibility for any bugs introduced
- [x] I have read and accepted the CLA agreement (only required once).

## Please put your Discord username so you can be contacted if a bug or
regression is found:

IngloriousTom

---------

Co-authored-by: Tom Rouillard <trouilla@mathworks.com>
Co-authored-by: Scott Anderson <662325+scottanderson@users.noreply.github.com>
---
 src/core/game/UnitGrid.ts | 48 ++++++++++++++++++++++++---------------
 1 file changed, 30 insertions(+), 18 deletions(-)

diff --git a/src/core/game/UnitGrid.ts b/src/core/game/UnitGrid.ts
index 62e703a68..027f1b33f 100644
--- a/src/core/game/UnitGrid.ts
+++ b/src/core/game/UnitGrid.ts
@@ -3,7 +3,7 @@ import { GameMap, TileRef } from "./GameMap";
 import { UnitView } from "./GameView";
 
 export class UnitGrid {
-  private grid: Set<Unit | UnitView>[][];
+  private grid: Map<UnitType, Set<Unit | UnitView>>[][];
   private readonly cellSize = 100;
 
   constructor(private gm: GameMap) {
@@ -12,7 +12,7 @@ export class UnitGrid {
       .map(() =>
         Array(Math.ceil(gm.width() / this.cellSize))
           .fill(null)
-          .map(() => new Set<Unit | UnitView>()),
+          .map(() => new Map<UnitType, Set<Unit | UnitView>>()),
       );
   }
 
@@ -27,7 +27,15 @@ export class UnitGrid {
     const [gridX, gridY] = this.getGridCoords(this.gm.x(tile), this.gm.y(tile));
 
     if (this.isValidCell(gridX, gridY)) {
-      this.grid[gridY][gridX].add(unit);
+      const unitSet = this.grid[gridY][gridX].get(unit.type());
+      if (unitSet !== undefined) {
+        unitSet.add(unit);
+      } else {
+        this.grid[gridY][gridX].set(
+          unit.type(),
+          new Set<Unit | UnitView>([unit]),
+        );
+      }
     }
   }
 
@@ -41,7 +49,10 @@ export class UnitGrid {
     const [gridX, gridY] = this.getGridCoords(this.gm.x(tile), this.gm.y(tile));
 
     if (this.isValidCell(gridX, gridY)) {
-      this.grid[gridY][gridX].delete(unit);
+      const unitSet = this.grid[gridY][gridX].get(unit.type());
+      if (unitSet !== undefined) {
+        unitSet.delete(unit);
+      }
     }
   }
 
@@ -134,14 +145,17 @@ export class UnitGrid {
     const typeSet = Array.isArray(types) ? new Set(types) : new Set([types]);
     for (let cy = startGridY; cy <= endGridY; cy++) {
       for (let cx = startGridX; cx <= endGridX; cx++) {
-        for (const unit of this.grid[cy][cx]) {
-          if (!typeSet.has(unit.type())) continue;
-          if (!unit.isActive()) continue;
-          const distSquared = this.squaredDistanceFromTile(unit, tile);
-          if (distSquared > rangeSquared) continue;
-          const value = { unit, distSquared };
-          if (predicate !== undefined && !predicate(value)) continue;
-          nearby.push(value);
+        for (const type of typeSet) {
+          const unitSet = this.grid[cy][cx].get(type);
+          if (unitSet === undefined) continue;
+          for (const unit of unitSet) {
+            if (!unit.isActive()) continue;
+            const distSquared = this.squaredDistanceFromTile(unit, tile);
+            if (distSquared > rangeSquared) continue;
+            const value = { unit, distSquared };
+            if (predicate !== undefined && !predicate(value)) continue;
+            nearby.push(value);
+          }
         }
       }
     }
@@ -162,12 +176,10 @@ export class UnitGrid {
     const rangeSquared = searchRange * searchRange;
     for (let cy = startGridY; cy <= endGridY; cy++) {
       for (let cx = startGridX; cx <= endGridX; cx++) {
-        for (const unit of this.grid[cy][cx]) {
-          if (
-            unit.type() === type &&
-            unit.owner().id() === playerId &&
-            unit.isActive()
-          ) {
+        const unitSet = this.grid[cy][cx].get(type);
+        if (unitSet === undefined) continue;
+        for (const unit of unitSet) {
+          if (unit.owner().id() === playerId && unit.isActive()) {
             const distSquared = this.squaredDistanceFromTile(unit, tile);
             if (distSquared <= rangeSquared) {
               return true;