From 9402b78ff27ec15d70654dcabe6dbe9ecc8beb3e Mon Sep 17 00:00:00 2001
From: trevan <trevan@users.noreply.github.com>
Date: Wed, 1 Jul 2020 16:13:25 -0600
Subject: [PATCH] Save the owner of the new units in case another action
 changes the ownership (#6943)

* Save the owner of the new units in case another action changes the ownership

* Handle case where unit isn't in the gameData yet

* Improve naming and add documentation for the unit-owner map
---
 .../engine/data/changefactory/AddUnits.java   | 50 +++++++++++++++++--
 1 file changed, 46 insertions(+), 4 deletions(-)

diff --git a/game-core/src/main/java/games/strategy/engine/data/changefactory/AddUnits.java b/game-core/src/main/java/games/strategy/engine/data/changefactory/AddUnits.java
index b6ce7fc127d..8fc41067769 100644
--- a/game-core/src/main/java/games/strategy/engine/data/changefactory/AddUnits.java
+++ b/game-core/src/main/java/games/strategy/engine/data/changefactory/AddUnits.java
@@ -2,11 +2,16 @@
 
 import games.strategy.engine.data.Change;
 import games.strategy.engine.data.GameData;
+import games.strategy.engine.data.GamePlayer;
 import games.strategy.engine.data.Unit;
 import games.strategy.engine.data.UnitCollection;
 import games.strategy.engine.data.UnitHolder;
-import java.util.ArrayList;
+import java.io.IOException;
+import java.io.ObjectInputStream;
 import java.util.Collection;
+import java.util.Map;
+import java.util.UUID;
+import java.util.stream.Collectors;
 
 /** Add units. */
 public class AddUnits extends Change {
@@ -15,19 +20,31 @@ public class AddUnits extends Change {
   private final String name;
   private final Collection<Unit> units;
   private final String type;
+  /**
+   * The unit's owner can be modified sometime after this Change is created but before it is
+   * performed. To ensure that the newly created units have the correct ownership, their original
+   * owners are stored in this separate map.
+   */
+  private Map<UUID, String> unitOwnerMap;
 
   AddUnits(final UnitCollection collection, final Collection<Unit> units) {
-    this.units = new ArrayList<>(units);
+    this.units = units;
+    unitOwnerMap = buildUnitOwnerMap(units);
     name = collection.getHolder().getName();
     type = collection.getHolder().getType();
   }
 
   AddUnits(final String name, final String type, final Collection<Unit> units) {
-    this.units = new ArrayList<>(units);
+    this.units = units;
+    unitOwnerMap = buildUnitOwnerMap(units);
     this.type = type;
     this.name = name;
   }
 
+  private Map<UUID, String> buildUnitOwnerMap(final Collection<Unit> units) {
+    return units.stream().collect(Collectors.toMap(Unit::getId, unit -> unit.getOwner().getName()));
+  }
+
   @Override
   public Change invert() {
     return new RemoveUnits(name, type, units);
@@ -36,11 +53,36 @@ public Change invert() {
   @Override
   protected void perform(final GameData data) {
     final UnitHolder holder = data.getUnitHolder(name, type);
-    holder.getUnitCollection().addAll(units);
+    final Collection<Unit> unitsWithCorrectOwner = buildUnitsWithOwner(data);
+    holder.getUnitCollection().addAll(unitsWithCorrectOwner);
+  }
+
+  private Collection<Unit> buildUnitsWithOwner(final GameData data) {
+    final Map<UUID, Unit> uuidToUnits =
+        units.stream().collect(Collectors.toMap(Unit::getId, unit -> unit));
+    return unitOwnerMap.entrySet().stream()
+        .map(
+            entry -> {
+              Unit unit = data.getUnits().get(entry.getKey());
+              if (unit == null) {
+                unit = uuidToUnits.get(entry.getKey());
+              }
+              final GamePlayer player = data.getPlayerList().getPlayerId(entry.getValue());
+              unit.setOwner(player);
+              return unit;
+            })
+        .collect(Collectors.toList());
   }
 
   @Override
   public String toString() {
     return "Add unit change.  Add to:" + name + " units:" + units;
   }
+
+  private void readObject(final ObjectInputStream in) throws IOException, ClassNotFoundException {
+    in.defaultReadObject();
+    if (unitOwnerMap == null) {
+      unitOwnerMap = buildUnitOwnerMap(units);
+    }
+  }
 }