Preserve entityIDs in atPosition's setter

Make the setter for the atPosition lens preserve entityIDs for
already-existing entities at the position, so that when we plop
something in the same tile as the character the character's entity ID
doesn't disappear.

2 files changed, 50 insertions, 15 deletions

diff --git a/src/Xanthous/Data/EntityMap.hs b/src/Xanthous/Data/EntityMap.hs
index 619b4b05c4b9..4e7796b1f415 100644
--- a/src/Xanthous/Data/EntityMap.hs
+++ b/src/Xanthous/Data/EntityMap.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE ViewPatterns #-}
 {-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE DeriveTraversable  #-}
@@ -185,7 +186,7 @@ insertAtReturningID pos e em =
 insertAt :: forall a. Position -> a -> EntityMap a -> EntityMap a
 insertAt pos e = snd . insertAtReturningID pos e
 
-atPosition :: forall a. Position -> Lens' (EntityMap a) (VectorBag a)
+atPosition :: forall a. (Ord a, Show a) => Position -> Lens' (EntityMap a) (VectorBag a)
 atPosition pos = lens getter setter
   where
     getter em =
@@ -194,16 +195,34 @@ atPosition pos = lens getter setter
                  $ em ^. byPosition . at pos
       in getEIDAssume em <$> eids
     setter em Empty = em & byPosition . at pos .~ Nothing
-    setter em entities =
-      alaf Endo foldMap (insertAt pos) entities
-      . removeAllAt pos
-      $ em
-      where
-        removeAllAt p e =
-          let eids = e ^.. byPosition . at p >>= toList >>= toList
-          in alaf Endo foldMap (\eid -> byID . at eid .~ Nothing) eids
-           . (byPosition . at pos .~ Nothing)
-           $ e
+    setter em (sort -> entities) =
+      let origEIDs = maybe Empty toNullable $ em ^. byPosition . at pos
+          origEntitiesWithIDs =
+            sortOn snd $ toList origEIDs <&> \eid -> (eid, getEIDAssume em eid)
+          go alles₁@((eid, e₁) :< es₁) -- orig
+             (e₂ :< es₂)               -- new
+            | e₁ == e₂
+              -- same, do nothing
+            = let (eids, lastEID, byID') = go es₁ es₂
+              in (insertSet eid eids, lastEID, byID')
+            | otherwise
+              -- e₂ is new, generate a new ID for it
+            = let (eids, lastEID, byID') = go alles₁ es₂
+                  eid' = succ lastEID
+              in (insertSet eid' eids, eid', byID' & at eid' ?~ Positioned pos e₂)
+          go Empty Empty = (mempty, em ^. lastID, em ^. byID)
+          go orig Empty =
+            let byID' = foldr deleteMap (em ^. byID) $ map fst orig
+            in (mempty, em ^. lastID, byID')
+          go Empty (new :< news) =
+            let (eids, lastEID, byID') = go Empty news
+                eid' = succ lastEID
+            in (insertSet eid' eids, eid', byID' & at eid' ?~ Positioned pos new)
+          go _ _ = error "unreachable"
+          (eidsAtPosition, newLastID, newByID) = go origEntitiesWithIDs entities
+      in em & byPosition . at pos .~ fromNullable eidsAtPosition
+            & byID .~ newByID
+            & lastID .~ newLastID
 
 getEIDAssume :: EntityMap a -> EntityID -> a
 getEIDAssume em eid = fromMaybe byIDInvariantError
@@ -247,7 +266,7 @@ lookup eid = fmap (view positioned) . lookupWithPosition eid
 -- positionedEntities :: IndexedTraversal EntityID (EntityMap a) (EntityMap b) (Positioned a) (Positioned b)
 -- positionedEntities = byID . itraversed
 
-neighbors :: Position -> EntityMap a -> Neighbors (VectorBag a)
+neighbors :: (Ord a, Show a) => Position -> EntityMap a -> Neighbors (VectorBag a)
 neighbors pos em = (\p -> view (atPosition p) em) <$> neighborPositions pos
 
 --------------------------------------------------------------------------------
diff --git a/test/Xanthous/Data/EntityMapSpec.hs b/test/Xanthous/Data/EntityMapSpec.hs
index 8317f5f51f8b..7c5cad019616 100644
--- a/test/Xanthous/Data/EntityMapSpec.hs
+++ b/test/Xanthous/Data/EntityMapSpec.hs
@@ -3,11 +3,11 @@
 module Xanthous.Data.EntityMapSpec where
 --------------------------------------------------------------------------------
 import           Test.Prelude
-import           Control.Lens.Properties
 --------------------------------------------------------------------------------
 import qualified Data.Aeson as JSON
 --------------------------------------------------------------------------------
 import           Xanthous.Data.EntityMap
+import           Xanthous.Data (Positioned(..))
 --------------------------------------------------------------------------------
 
 main :: IO ()
@@ -47,7 +47,23 @@ test = localOption (QuickCheckTests 20)
         in toEIDsAndPositioned em' === toEIDsAndPositioned em
     ]
 
-  , testGroup "atPosition"
-    [ testProperty "lens laws" $ \pos -> isLens $ atPosition @Int pos
+  , localOption (QuickCheckTests 50)
+  $ testGroup "atPosition"
+    [ testProperty "setget" $ \pos (em :: EntityMap Int) es ->
+        view (atPosition pos) (set (atPosition pos) es em) === es
+    , testProperty "getset" $ \pos (em :: EntityMap Int) ->
+        set (atPosition pos) (view (atPosition pos) em) em === em
+    , testProperty "setset" $ \pos (em :: EntityMap Int) es ->
+        (set (atPosition pos) es . set (atPosition pos) es) em
+        ===
+        set (atPosition pos) es em
+      -- testProperty "lens laws" $ \pos -> isLens $ atPosition @Int pos
+    , testProperty "preserves IDs" $ \(em :: EntityMap Int) e1 e2 p ->
+        let (eid, em') = insertAtReturningID p e1 em
+            em'' = em' & atPosition p %~ (e2 <|)
+        in
+          counterexample ("em': " <> show em')
+          . counterexample ("em'': " <> show em'')
+          $ em'' ^. at eid === Just (Positioned p e1)
     ]
   ]