1 files changed, 276 insertions, 0 deletions
diff --git a/users/aspen/xanthous/src/Xanthous/Data/EntityMap.hs b/users/aspen/xanthous/src/Xanthous/Data/EntityMap.hs
new file mode 100644
index 0000000000..33a98f1ae5
--- /dev/null
+++ b/users/aspen/xanthous/src/Xanthous/Data/EntityMap.hs
@@ -0,0 +1,276 @@
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE DeriveTraversable  #-}
+{-# LANGUAGE TupleSections      #-}
+{-# LANGUAGE TemplateHaskell    #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE DeriveFunctor      #-}
+--------------------------------------------------------------------------------
+module Xanthous.Data.EntityMap
+  ( EntityMap
+  , _EntityMap
+  , EntityID
+  , emptyEntityMap
+  , insertAt
+  , insertAtReturningID
+  , fromEIDsAndPositioned
+  , toEIDsAndPositioned
+  , atPosition
+  , atPositionWithIDs
+  , positions
+  , lookup
+  , lookupWithPosition
+  , positionOf
+  -- , positionedEntities
+  , neighbors
+  , Deduplicate(..)
+
+  -- * debug
+  , byID
+  , byPosition
+  , lastID
+
+  ) where
+--------------------------------------------------------------------------------
+import Xanthous.Prelude hiding (lookup)
+import Xanthous.Data
+  ( Position
+  , Positioned(..)
+  , positioned
+  , Neighbors(..)
+  , neighborPositions, position
+  )
+import Xanthous.Data.VectorBag
+import Xanthous.Orphans ()
+import Xanthous.Util (EqEqProp(..))
+--------------------------------------------------------------------------------
+import Data.Monoid (Endo(..))
+import Test.QuickCheck (Arbitrary(..), CoArbitrary, Function)
+import Test.QuickCheck.Checkers (EqProp)
+import Test.QuickCheck.Instances.UnorderedContainers ()
+import Test.QuickCheck.Instances.Vector ()
+import Text.Show (showString, showParen)
+import Data.Aeson
+--------------------------------------------------------------------------------
+
+type EntityID = Word32
+type NonNullSet a = NonNull (Set a)
+
+data EntityMap a where
+  EntityMap ::
+    { _byPosition :: Map Position (NonNullSet EntityID)
+    , _byID       :: HashMap EntityID (Positioned a)
+    , _lastID     :: EntityID
+    } -> EntityMap a
+  deriving stock (Functor, Foldable, Traversable, Generic)
+  deriving anyclass (NFData, CoArbitrary, Function)
+deriving via (EqEqProp (EntityMap a)) instance (Eq a, Ord a) => EqProp (EntityMap a)
+makeLenses ''EntityMap
+
+instance ToJSON a => ToJSON (EntityMap a) where
+  toJSON = toJSON . toEIDsAndPositioned
+
+
+instance FromJSON a => FromJSON (EntityMap a) where
+  parseJSON = fmap (fromEIDsAndPositioned @[_]) . parseJSON
+
+byIDInvariantError :: forall a. a
+byIDInvariantError = error $ "Invariant violation: All EntityIDs in byPosition "
+  <> "must point to entityIDs in byID"
+
+instance (Ord a, Eq a) => Eq (EntityMap a) where
+  -- em₁ == em₂ = em₁ ^. _EntityMap == em₂ ^. _EntityMap
+  (==) = (==) `on` view (_EntityMap . to sort)
+
+deriving stock instance (Ord a) => Ord (EntityMap a)
+
+instance Show a => Show (EntityMap a) where
+  showsPrec pr em
+    = showParen (pr > 10)
+    $ showString
+    . ("fromEIDsAndPositioned " <>)
+    . show
+    . toEIDsAndPositioned
+    $ em
+
+instance Arbitrary a => Arbitrary (EntityMap a) where
+  arbitrary = review _EntityMap <$> arbitrary
+  shrink em = review _EntityMap <$> shrink (em ^. _EntityMap)
+
+type instance Index (EntityMap a) = EntityID
+type instance IxValue (EntityMap a) = (Positioned a)
+instance Ixed (EntityMap a) where ix eid = at eid . traverse
+
+instance At (EntityMap a) where
+  at eid = lens (view $ byID . at eid) setter
+    where
+      setter :: EntityMap a -> Maybe (Positioned a) -> EntityMap a
+      setter m Nothing = fromMaybe m $ do
+        Positioned pos _ <- m ^. byID . at eid
+        pure $ m
+          & removeEIDAtPos pos
+          & byID . at eid .~ Nothing
+      setter m (Just pe@(Positioned pos _)) = m
+        & (case lookupWithPosition eid m of
+             Nothing -> id
+             Just (Positioned origPos _) -> removeEIDAtPos origPos
+          )
+        & byID . at eid ?~ pe
+        & byPosition . at pos %~ \case
+            Nothing -> Just $ opoint eid
+            Just es -> Just $ ninsertSet eid es
+      removeEIDAtPos pos =
+        byPosition . at pos %~ (>>= fromNullable . ndeleteSet eid)
+
+instance Semigroup (EntityMap a) where
+  em₁ <> em₂ = alaf Endo foldMap (uncurry insertAt) (em₂ ^. _EntityMap) em₁
+
+instance Monoid (EntityMap a) where
+  mempty = emptyEntityMap
+
+instance FunctorWithIndex EntityID EntityMap
+
+instance FoldableWithIndex EntityID EntityMap
+
+instance TraversableWithIndex EntityID EntityMap where
+  itraverse = itraverseOf itraversed
+
+type instance Element (EntityMap a) = a
+instance MonoFoldable (EntityMap a)
+
+emptyEntityMap :: EntityMap a
+emptyEntityMap = EntityMap mempty mempty 0
+
+newtype Deduplicate a = Deduplicate (EntityMap a)
+  deriving stock (Show, Traversable, Generic)
+  deriving newtype (Eq, Functor, Foldable, EqProp, Arbitrary)
+
+instance Semigroup (Deduplicate a) where
+  (Deduplicate em₁) <> (Deduplicate em₂) =
+    let _byID = em₁ ^. byID <> em₂ ^. byID
+        _byPosition = mempty &~ do
+          ifor_ _byID $ \eid (Positioned pos _) ->
+            at pos %= \case
+              Just eids -> Just $ ninsertSet eid eids
+              Nothing -> Just $ opoint eid
+        _lastID = fromMaybe 1 $ maximumOf (ifolded . asIndex) _byID
+    in Deduplicate EntityMap{..}
+
+
+--------------------------------------------------------------------------------
+
+_EntityMap :: Iso' (EntityMap a) [(Position, a)]
+_EntityMap = iso hither yon
+  where
+    hither :: EntityMap a -> [(Position, a)]
+    hither em = do
+       (pos, eids) <- em ^. byPosition . _Wrapped
+       eid <- toList eids
+       ent <- em ^.. byID . at eid . folded . positioned
+       pure (pos, ent)
+    yon :: [(Position, a)] -> EntityMap a
+    yon poses = alaf Endo foldMap (uncurry insertAt) poses emptyEntityMap
+
+
+insertAtReturningID :: forall a. Position -> a -> EntityMap a -> (EntityID, EntityMap a)
+insertAtReturningID pos e em =
+  let (eid, em') = em & lastID <+~ 1
+  in em'
+     & byID . at eid ?~ Positioned pos e
+     & byPosition . at pos %~ \case
+       Nothing -> Just $ opoint eid
+       Just es -> Just $ ninsertSet eid es
+     & (eid, )
+
+insertAt :: forall a. Position -> a -> EntityMap a -> EntityMap a
+insertAt pos e = snd . insertAtReturningID pos e
+
+atPosition :: forall a. (Ord a, Show a) => Position -> Lens' (EntityMap a) (VectorBag a)
+atPosition pos = lens getter setter
+  where
+    getter em =
+      let eids :: VectorBag EntityID
+          eids = maybe mempty (VectorBag . toVector . toNullable)
+                 $ em ^. byPosition . at pos
+      in getEIDAssume em <$> eids
+    setter em Empty = em & byPosition . at pos .~ Nothing
+    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
+  $ em ^? byID . ix eid . positioned
+
+atPositionWithIDs :: Position -> EntityMap a -> Vector (EntityID, Positioned a)
+atPositionWithIDs pos em =
+  let eids = maybe mempty (toVector . toNullable)
+             $ em ^. byPosition . at pos
+  in (id &&& Positioned pos . getEIDAssume em) <$> eids
+
+fromEIDsAndPositioned
+  :: forall mono a. (MonoFoldable mono, Element mono ~ (EntityID, Positioned a))
+  => mono
+  -> EntityMap a
+fromEIDsAndPositioned eps = newLastID $ alaf Endo foldMap insert' eps mempty
+  where
+    insert' (eid, pe@(Positioned pos _))
+      = (byID . at eid ?~ pe)
+      . (byPosition . at pos %~ \case
+            Just eids -> Just $ ninsertSet eid eids
+            Nothing   -> Just $ opoint eid
+        )
+    newLastID em = em & lastID
+      .~ fromMaybe 1
+         (maximumOf (ifolded . asIndex) (em ^. byID))
+
+toEIDsAndPositioned :: EntityMap a -> [(EntityID, Positioned a)]
+toEIDsAndPositioned = itoListOf $ byID . ifolded
+
+positions :: EntityMap a -> [Position]
+positions = toListOf $ byPosition . to keys . folded
+
+lookupWithPosition :: EntityID -> EntityMap a -> Maybe (Positioned a)
+lookupWithPosition eid = view $ byID . at eid
+
+lookup :: EntityID -> EntityMap a -> Maybe a
+lookup eid = fmap (view positioned) . lookupWithPosition eid
+
+-- unlawful :(
+-- positionedEntities :: IndexedTraversal EntityID (EntityMap a) (EntityMap b) (Positioned a) (Positioned b)
+-- positionedEntities = byID . itraversed
+
+neighbors :: (Ord a, Show a) => Position -> EntityMap a -> Neighbors (VectorBag a)
+neighbors pos em = (\p -> view (atPosition p) em) <$> neighborPositions pos
+
+-- | Traversal to the position of the entity with the given ID
+positionOf :: EntityID -> Traversal' (EntityMap a) Position
+positionOf eid = ix eid . position
+
+--------------------------------------------------------------------------------
+makeWrapped ''Deduplicate