🐛 Fix distribution issue and add test

dojoengine · Sep 6, 2024 · 69c4b84 · 69c4b84
1 parent c900f66
commit 69c4b84
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 4 deletions.
diff --git a/crates/map/src/helpers/spreader.cairo b/crates/map/src/helpers/spreader.cairo
@@ -54,7 +54,7 @@ pub impl Spreader of SpreaderTrait {
     #[inline]
     fn iter(mut grid: felt252, index: u8, total: u8, mut count: u8, seed: felt252) -> felt252 {
         // [Checl] Stop if all objects are placed
-        if count == 0 || index >= total {
+        if count == 0 {
             return grid;
         };
         // [Check] Skip if the position is already occupied
@@ -64,15 +64,15 @@ pub impl Spreader of SpreaderTrait {
         };
         // [Compute] Uniform random number between 0 and MULTIPLIER
         let random = seed.into() % MULTIPLIER;
-        let probability: u256 = count.into() * MULTIPLIER / (total - index).into();
+        let probability: u256 = count.into() * MULTIPLIER / total.into();
         // [Check] Probability of being an object
         if random <= probability {
             // [Compute] Update grid
             count -= 1;
             // [Effect] Set bit to 0
             grid = Bitmap::unset(grid, index);
         };
-        Self::iter(grid, index + 1, total, count, seed)
+        Self::iter(grid, index + 1, total - 1, count, seed)
     }
 }
 
@@ -87,7 +87,7 @@ mod tests {
     const SEED: felt252 = 'SEED';
 
     #[test]
-    fn test_spreader_generate() {
+    fn test_spreader_generate_large() {
         // 000000000000100000
         // 000010100000000000
         // 000010000100000000
@@ -108,5 +108,21 @@ mod tests {
         let room = Spreader::generate(grid, width, height, 35, SEED);
         assert_eq!(room, 0x802800084000A006408008401003008004308C0002410E01080200008120);
     }
+
+    #[test]
+    fn test_spreader_generate_small() {
+        // 000000001110000000
+        // 000000001110000000
+        // 000000001110000000
+        // Output:
+        // 000000000110000000
+        // 000000000000000000
+        // 000000001100000000
+        let width = 18;
+        let height = 14;
+        let grid: felt252 = 0x38000E000380;
+        let room = Spreader::generate(grid, width, height, 4, SEED);
+        assert_eq!(room, 0x180000000300);
+    }
 }
 
diff --git a/crates/map/src/room.cairo b/crates/map/src/room.cairo
@@ -246,4 +246,43 @@ mod tests {
         room.open_with_maze(250);
         assert_eq!(room.grid, 0x4000100C060730D1FE6E3F8FFFE69FF8A77E6FFF93FFE4FFF9BFE037F800000);
     }
+
+    #[test]
+    fn test_room_compute_distribution() {
+        // 000000000000000000
+        // 000000000011000000
+        // 000000000111001100
+        // 000001000111111110
+        // 000011100011111110
+        // 000011111111111110
+        // 000010011111111110
+        // 000010011101111110
+        // 000011111111111110
+        // 000011111111111110
+        // 000011111111111110
+        // 000011111111100000
+        // 000001111111100000
+        // 000000000000000000
+        // Output:
+        // 000000000000000000
+        // 000000000000000000
+        // 000000000000000000
+        // 000000000000001000
+        // 000000100001000000
+        // 000010000000001000
+        // 000000000000000000
+        // 000000000000000000
+        // 000000000001001000
+        // 000000000100000000
+        // 000000000000000000
+        // 000000010000100000
+        // 000000000000000000
+        // 000000000000000000
+        let width = 18;
+        let height = 14;
+        let steps: u16 = 2 * width.into() * height.into();
+        let mut room: Room = RoomTrait::new_random_walk(width, height, steps, SEED);
+        let distribution = room.compute_distribution(10, SEED);
+        assert_eq!(distribution, 0x8021002008000000000001200100000000420000000000);
+    }
 }