Add get rooms by layer method

mpewsey · May 13, 2024 · 7b1ac84 · 7b1ac84
1 parent e63d2f9
commit 7b1ac84
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 1 deletion.
diff --git 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);
         }
     }
 }
diff --git 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;
+        }
     }
 }