1 files changed, 141 insertions, 0 deletions
diff --git a/src/Xanthous/Data/EntityMap.hs b/src/Xanthous/Data/EntityMap.hs
new file mode 100644
index 000000000000..e3ceb6f65182
--- /dev/null
+++ b/src/Xanthous/Data/EntityMap.hs
@@ -0,0 +1,141 @@
+{-# LANGUAGE DeriveTraversable #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE DeriveFunctor #-}
+
+module Xanthous.Data.EntityMap
+  ( EntityMap
+  , EntityID
+  , emptyEntityMap
+  , insertAt
+  , insertAtReturningID
+  , atPosition
+  , positions
+  , lookup
+  , lookupWithPosition
+  -- , positionedEntities
+  ) where
+
+import Data.Monoid (Endo(..))
+import Test.QuickCheck (Arbitrary(..))
+import Test.QuickCheck.Checkers (EqProp)
+
+import Xanthous.Prelude hiding (lookup)
+import Xanthous.Data (Position, Positioned(..), positioned, position)
+import Xanthous.Orphans ()
+import Xanthous.Util (EqEqProp(..))
+
+type EntityID = Word32
+type NonNullVector a = NonNull (Vector a)
+
+data EntityMap a where
+  EntityMap ::
+    { _byPosition :: Map Position (NonNullVector EntityID)
+    , _byID       :: HashMap EntityID (Positioned a)
+    , _lastID     :: EntityID
+    } -> EntityMap a
+  deriving stock (Functor, Foldable, Traversable)
+deriving via (EqEqProp (EntityMap a)) instance Eq a => EqProp (EntityMap a)
+makeLenses ''EntityMap
+
+byIDInvariantError :: forall a. a
+byIDInvariantError = error $ "Invariant violation: All EntityIDs in byPosition "
+  <> "must point to entityIDs in byID"
+
+instance Eq a => Eq (EntityMap a) where
+  em₁ == em₂ = em₁ ^. _EntityMap == em₂ ^. _EntityMap
+
+instance Show a => Show (EntityMap a) where
+  show em = "_EntityMap # " <> show (em ^. _EntityMap)
+
+instance Arbitrary a => Arbitrary (EntityMap a) where
+  arbitrary = review _EntityMap <$> arbitrary
+
+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 (Positioned pos e)) =
+        case lookupWithPosition eid m of
+          Nothing -> insertAt pos e m
+          Just (Positioned origPos _) -> m
+            & removeEIDAtPos origPos
+            & byID . ix eid . position .~ pos
+            & byPosition . at pos %~ \case
+              Nothing -> Just $ ncons eid mempty
+              Just es -> Just $ eid <| es
+      removeEIDAtPos pos =
+        byPosition . at pos %~ (>>= fromNullable . nfilter (/= eid))
+
+emptyEntityMap :: EntityMap a
+emptyEntityMap = EntityMap mempty mempty 0
+
+_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
+
+instance Semigroup (EntityMap a) where
+  em₁ <> em₂ = alaf Endo foldMap (uncurry insertAt) (em₁ ^. _EntityMap) em₂
+
+instance Monoid (EntityMap a) where
+  mempty = 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 $ ncons eid mempty
+       Just es -> Just $ eid <| es
+     & (eid, )
+
+insertAt :: forall a. Position -> a -> EntityMap a -> EntityMap a
+insertAt pos e = snd . insertAtReturningID pos e
+
+atPosition :: forall a. Position -> Lens' (EntityMap a) (Vector a)
+atPosition pos = lens getter setter
+  where
+    getter em =
+      let
+        eids :: Vector EntityID
+        eids = maybe mempty toNullable $ em ^. byPosition . at pos
+
+        getEIDAssume :: EntityID -> a
+        getEIDAssume eid = fromMaybe byIDInvariantError
+          $ em ^? byID . ix eid . positioned
+      in getEIDAssume <$> eids
+    setter em Empty = em & byPosition . at pos .~ Nothing
+    setter em entities = alaf Endo foldMap (insertAt pos) entities em
+
+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