From 7b1ac84fc64257e21a7c666e6160156d0211652d Mon Sep 17 00:00:00 2001
From: Matt Pewsey <23442063+mpewsey@users.noreply.github.com>
Date: Sun, 12 May 2024 21:02:20 -0400
Subject: [PATCH] Add get rooms by layer method

---
 src/ManiaMap.Tests/TestLayout.cs | 26 +++++++++++++++++++++++++-
 src/ManiaMap/Layout.cs           | 21 +++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/src/ManiaMap.Tests/TestLayout.cs b/src/ManiaMap.Tests/TestLayout.cs
index 7654efe..0f954ce 100644
--- a/src/ManiaMap.Tests/TestLayout.cs
+++ b/src/ManiaMap.Tests/TestLayout.cs
@@ -257,7 +257,31 @@ public void TestGetLayers()
             Assert.IsNotNull(layout);
 
             var expected = new int[] { -1, 0, 2 };
-            CollectionAssert.AreEqual(layout.GetLayers(), expected);
+            CollectionAssert.AreEqual(expected, layout.GetLayers());
+        }
+
+        [TestMethod]
+        public void TestGetRoomsByLayer()
+        {
+            var graph = Samples.GraphLibrary.StackedLoopGraph();
+
+            var templateGroups = new TemplateGroups();
+            templateGroups.Add("Default", Samples.TemplateLibrary.Miscellaneous.HyperSquareTemplate());
+
+            var generator = new LayoutGenerator();
+            var random = new RandomSeed(12345);
+            var layout = generator.Generate(1, graph, templateGroups, random);
+
+            Assert.IsNotNull(layout);
+
+            var roomsByLayer = layout.GetRoomsByLayer();
+            Assert.AreEqual(3, roomsByLayer.Count);
+            Assert.IsTrue(roomsByLayer.ContainsKey(-1));
+            Assert.IsTrue(roomsByLayer.ContainsKey(0));
+            Assert.IsTrue(roomsByLayer.ContainsKey(2));
+            Assert.AreEqual(1, roomsByLayer[-1].Count);
+            Assert.AreEqual(4, roomsByLayer[0].Count);
+            Assert.AreEqual(3, roomsByLayer[2].Count);
         }
     }
 }
\ No newline at end of file
diff --git a/src/ManiaMap/Layout.cs b/src/ManiaMap/Layout.cs
index 013b557..2aae473 100644
--- a/src/ManiaMap/Layout.cs
+++ b/src/ManiaMap/Layout.cs
@@ -504,5 +504,26 @@ public int[] GetLayers()
             Array.Sort(result);
             return result;
         }
+
+        /// <summary>
+        /// Returns a dictionary of rooms by layer (Z) coordinate.
+        /// </summary>
+        public Dictionary<int, List<Room>> GetRoomsByLayer()
+        {
+            var result = new Dictionary<int, List<Room>>();
+
+            foreach (var room in Rooms.Values)
+            {
+                if (!result.TryGetValue(room.Position.Z, out var rooms))
+                {
+                    rooms = new List<Room>();
+                    result.Add(room.Position.Z, rooms);
+                }
+
+                rooms.Add(room);
+            }
+
+            return result;
+        }
     }
 }