apply review suggestions regarding imap

src/swarm-topography/Swarm/Game/Scenario/Topography/Grid.hs

@@ -4,7 +4,6 @@ module Swarm.Game.Scenario.Topography.Grid (
   Grid (..),
   NonEmptyGrid (..),
   gridToVec,
-  zipNumberedNE,
   mapWithCoordsNE,
   mapWithCoords,
   allMembers,
@@ -14,12 +13,14 @@ module Swarm.Game.Scenario.Topography.Grid (
 )
 where
 
+import Control.Lens.Indexed (FunctorWithIndex, imap)
 import Data.Aeson (ToJSON (..))
 import Data.Foldable qualified as F
 import Data.List.NonEmpty (NonEmpty)
 import Data.List.NonEmpty qualified as NE
 import Data.Maybe (fromMaybe)
 import Data.Semigroup
+import Data.Tuple.Extra (both)
 import Data.Vector qualified as V
 import GHC.Generics (Generic)
 import Swarm.Game.World.Coords
@@ -28,6 +29,11 @@ import Prelude hiding (zipWith)
 newtype NonEmptyGrid c = NonEmptyGrid (NonEmpty (NonEmpty c))
   deriving (Generic, Show, Eq, Functor, Foldable, Traversable, ToJSON)
 
+instance FunctorWithIndex Coords NonEmptyGrid where
+  imap f (NonEmptyGrid g) =
+    NonEmptyGrid $
+      imap (\i -> imap (\j -> f (Coords $ both fromIntegral (i, j)))) g
+
 data Grid c
   = EmptyGrid
   | Grid (NonEmptyGrid c)
@@ -51,25 +57,15 @@ mapRowsNE ::
   NonEmptyGrid b
 mapRowsNE f (NonEmptyGrid rows) = NonEmptyGrid $ f rows
 
+allMembersNE :: NonEmptyGrid c -> NonEmpty c
+allMembersNE (NonEmptyGrid g) = sconcat g
+
 allMembers :: Grid a -> [a]
 allMembers EmptyGrid = []
 allMembers g = F.toList g
 
-nonemptyCount :: (Integral i) => NonEmpty i
-nonemptyCount = NE.iterate succ 0
-
-zipNumberedNE ::
-  Integral i =>
-  (i -> a -> b) ->
-  NonEmpty a ->
-  NonEmpty b
-zipNumberedNE f = NE.zipWith f nonemptyCount
-
 mapWithCoordsNE :: (Coords -> a -> b) -> NonEmptyGrid a -> NonEmpty b
-mapWithCoordsNE f (NonEmptyGrid g) =
-  sconcat $ NE.zipWith outer nonemptyCount g
- where
-  outer i = zipNumberedNE $ \j -> f (Coords (i, j))
+mapWithCoordsNE f = allMembersNE . imap f
 
 mapWithCoords :: (Coords -> a -> b) -> Grid a -> [b]
 mapWithCoords _ EmptyGrid = []
@@ -78,8 +74,7 @@ mapWithCoords f (Grid g) = NE.toList $ mapWithCoordsNE f g
 -- | Converts linked lists to vectors to facilitate
 -- random access when assembling the image
 gridToVec :: Grid a -> V.Vector (V.Vector a)
-gridToVec EmptyGrid = V.empty
-gridToVec (Grid (NonEmptyGrid g)) = V.fromList . map (V.fromList . NE.toList) $ NE.toList g
+gridToVec g = V.fromList . map V.fromList $ getRows g
 
 instance (ToJSON a) => ToJSON (Grid a) where
   toJSON EmptyGrid = toJSON ([] :: [a])

src/swarm-topography/Swarm/Game/Scenario/Topography/Structure/Recognition/Prep.hs

@@ -6,6 +6,7 @@ module Swarm.Game.Scenario.Topography.Structure.Recognition.Prep (
 ) where
 
 import Control.Arrow ((&&&))
+import Control.Lens.Indexed (imap)
 import Data.HashMap.Strict qualified as HM
 import Data.HashSet qualified as HS
 import Data.Hashable (Hashable)
@@ -31,7 +32,7 @@ allStructureRows :: [StructureWithGrid b a] -> [StructureRow b a]
 allStructureRows =
   concatMap $ NE.toList . transformRows
  where
-  transformRows g = zipNumberedNE (StructureRow g) rows
+  transformRows g = imap (StructureRow g . fromIntegral) rows
    where
     NonEmptyGrid rows = entityGrid g
 
@@ -150,7 +151,7 @@ explodeRowEntities annotatedRow@(ConsolidatedRowReferences rowMembers _ width) =
     map swap
       . catMaybes
       . NE.toList
-      $ zipNumberedNE (\idx -> fmap (PositionWithinRow idx annotatedRow,)) rowMembers
+      $ imap (\idx -> fmap (PositionWithinRow (fromIntegral idx) annotatedRow,)) rowMembers
 
   deriveEntityOffsets :: PositionWithinRow b a -> InspectionOffsets
   deriveEntityOffsets (PositionWithinRow pos _) = mkOffsets pos width