Added win check and valid play check, Added test for both

1 files changed, 221 insertions, 25 deletions

diff --git a/src/main.rs b/src/main.rs
index 35c3e7a..137115e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,8 +1,16 @@
-const HEIGHT: u64 = 6;
-const WIDTH:  u64 = 7;
+const HEIGHT:      u64 = 6;
+const WIDTH:       u64 = 7;
+const FULL_HEIGHT: u64 = HEIGHT + 1;
 
 type Bits = u64;
 
+#[derive(Debug, Eq, PartialEq)]
+enum Cell {
+    Empty,
+    CurrentPlayer,
+    OtherPlayer,
+}
+
 /**
  * bit order:
  *
@@ -19,6 +27,7 @@ type Bits = u64;
  *
  *
  */
+#[derive(Clone)]
 struct Position {
     /// stones of the current player
     player: u64,
@@ -26,6 +35,7 @@ struct Position {
     mask:   u64,
 }
 
+// https://github.com/PascalPons/connect4/blob/master/Position.hpp
 impl Position {
     fn new() -> Position {
         Position{ player: 0, mask: 0 }
@@ -34,53 +44,239 @@ impl Position {
     fn from_position(col_moves: &[u64]) -> Position {
         let mut position = Position::new();
         for col_pos in col_moves {
-            position.play(*col_pos);
+            position = position.play(*col_pos);
         }
         position
     }
 
-    fn play(&mut self, col_pos: u64) {
-        self.player ^= self.mask;
-        self.mask |= self.mask + (1 << (col_pos * (HEIGHT + 1)));
+    fn play(&self, col_pos: u64) -> Position {
+        Position {
+            player: self.player ^ self.mask,
+            mask:   self.mask | self.mask + Position::bottom_mask(col_pos),
+        }
+    }
+
+    fn is_valid_play(&self, col_pos: u64) -> bool {
+        self.mask & Position::top_mask(col_pos) == 0
+    }
+
+    fn is_winning_play(&self, col_pos: u64) -> bool {
+        // the current player stones with the play
+        // mask + bottom_mask = add a stone in column
+        // & column_mask: only modify the current player stone
+        // (ignoring other player in other columns)
+        let p = self.player | ((self.mask + Position::bottom_mask(col_pos)) & Position::column_mask(col_pos));
+
+        // shift the board in different direction,
+        // if the shifed boards have at least 1 stone in common
+        // it means that the original has 4 stone aligned.
+        // possible optimization with tmp variable to make 2 shifts instead of 4
+
+        // vertical
+        if (p >> 0) &
+           (p >> 1) &
+           (p >> 2) &
+           (p >> 3) != 0 {
+            return true;
+        }
+        // horizontal
+        if (p >> (0 * FULL_HEIGHT)) &
+           (p >> (1 * FULL_HEIGHT)) &
+           (p >> (2 * FULL_HEIGHT)) &
+           (p >> (3 * FULL_HEIGHT)) != 0 {
+            return true;
+        }
+        // diagonal
+        if (p >> (0 * HEIGHT)) &
+           (p >> (1 * HEIGHT)) &
+           (p >> (2 * HEIGHT)) &
+           (p >> (3 * HEIGHT)) != 0 {
+            return true;
+        }
+        // anti diagonal
+        if (p >> (0 * (HEIGHT + 2))) &
+           (p >> (1 * (HEIGHT + 2))) &
+           (p >> (2 * (HEIGHT + 2))) &
+           (p >> (3 * (HEIGHT + 2))) != 0 {
+            return true;
+        }
+        false
     }
 
     fn key(&self) -> u64 {
         self.player + self.mask
     }
+
+    fn bottom_mask(col_pos: u64) -> u64 {
+        1 << (col_pos * FULL_HEIGHT)
+    }
+
+    fn top_mask(col_pos: u64) -> u64 {
+        Position::bottom_mask(col_pos) << (HEIGHT - 1)
+    }
+
+    fn column_mask(col_pos: u64) -> u64 {
+        0b1111111 << (col_pos * FULL_HEIGHT)
+        // ((1 << HEIGHT) - 1) << (col_pos * FULL_HEIGHT)
+    }
+
+    fn at(&self, y: u64, x: u64) -> Cell {
+        let pos_mask = (1 << (x * FULL_HEIGHT)) << y;
+        if self.mask & pos_mask == 0 {
+            return Cell::Empty;
+        }
+        if self.player & pos_mask != 0 {
+            return Cell::CurrentPlayer;
+        }
+        return Cell::OtherPlayer;
+    }
 }
 
 use std::fmt;
 
 impl fmt::Debug for Position {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "mask: {:064b}\n", self.mask)?;
-        for i in 0..WIDTH {
-            write!(f, "\t{:07b}\n", (self.mask & (0b1111111 << (i * WIDTH))) >> (i * WIDTH));
+        use Cell::*;
+        write!(f, "{:10}{:10}{:10}\n", "position", "mask", "player")?;
+        for y in (0..FULL_HEIGHT).rev() {
+            for x in 0..WIDTH {
+                match self.at(y, x) {
+                    Empty         => write!(f, ".")?,
+                    CurrentPlayer => write!(f, "x")?,
+                    OtherPlayer   => write!(f, "o")?,
+                }
+            }
+            write!(f, "  ")?;
+            for x in 0..WIDTH {
+                match self.at(y, x) {
+                    Empty                       => write!(f, ".")?,
+                    CurrentPlayer | OtherPlayer => write!(f, "#")?,
+                }
+            }
+            write!(f, "  ")?;
+            for x in 0..WIDTH {
+                match self.at(y, x) {
+                    Empty | OtherPlayer => write!(f, ".")?,
+                    CurrentPlayer       => write!(f, "#")?,
+                }
+            }
+            write!(f, "\n")?;
         }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
 
-        write!(f, "player: {:064b}\n", self.player)?;
-        for i in 0..WIDTH {
-            write!(f, "\t{:07b}\n", (self.player & (0b1111111 << (i * WIDTH))) >> (i * WIDTH));
+    #[test]
+    fn test_new() {
+        let p = Position::new();
+        for y in 0..HEIGHT {
+            for x in 0..WIDTH {
+                assert_eq!(p.at(y, x), Cell::Empty);
+            }
         }
+    }
+
+    #[test]
+    fn test_play() {
+        let mut p = Position::new();
+        p = p.play(0);
+        assert_eq!(p.at(0, 0), Cell::OtherPlayer,   "\n{:?}", p);
+        p = p.play(0);
+        assert_eq!(p.at(0, 0), Cell::CurrentPlayer, "\n{:?}", p);
+        assert_eq!(p.at(1, 0), Cell::OtherPlayer,   "\n{:?}", p);
+        p = p.play(0);
+        assert_eq!(p.at(0, 0), Cell::OtherPlayer,   "\n{:?}", p);
+        assert_eq!(p.at(1, 0), Cell::CurrentPlayer, "\n{:?}", p);
+        assert_eq!(p.at(2, 0), Cell::OtherPlayer,   "\n{:?}", p);
+        p = p.play(0);
+        assert_eq!(p.at(0, 0), Cell::CurrentPlayer, "\n{:?}", p);
+        assert_eq!(p.at(1, 0), Cell::OtherPlayer,   "\n{:?}", p);
+        assert_eq!(p.at(2, 0), Cell::CurrentPlayer, "\n{:?}", p);
+        assert_eq!(p.at(3, 0), Cell::OtherPlayer,   "\n{:?}", p);
+
+        p = p.play(1);
+        assert_eq!(p.at(0, 0), Cell::OtherPlayer,   "\n{:?}", p);
+        assert_eq!(p.at(1, 0), Cell::CurrentPlayer, "\n{:?}", p);
+        assert_eq!(p.at(2, 0), Cell::OtherPlayer,   "\n{:?}", p);
+        assert_eq!(p.at(3, 0), Cell::CurrentPlayer, "\n{:?}", p);
+        assert_eq!(p.at(0, 1), Cell::OtherPlayer,   "\n{:?}", p);
 
-        write!(f, "key: {:064b}\n", self.key())?;
-        for i in 0..WIDTH {
-            write!(f, "\t{:07b}\n", (self.key() & (0b1111111 << (i * WIDTH))) >> (i * WIDTH));
+        p = p.play(WIDTH - 1);
+        assert_eq!(p.at(0, 0), Cell::CurrentPlayer, "\n{:?}", p);
+        assert_eq!(p.at(1, 0), Cell::OtherPlayer,   "\n{:?}", p);
+        assert_eq!(p.at(2, 0), Cell::CurrentPlayer, "\n{:?}", p);
+        assert_eq!(p.at(3, 0), Cell::OtherPlayer,   "\n{:?}", p);
+        assert_eq!(p.at(0, 1), Cell::CurrentPlayer, "\n{:?}", p);
+        assert_eq!(p.at(0, WIDTH - 1), Cell::OtherPlayer, "\n{:?}", p);
+    }
+
+    #[test]
+    fn test_is_valid_play() {
+        let mut p = Position::new();
+        for c in 0..WIDTH {
+            for _ in 0..HEIGHT {
+                assert!(p.is_valid_play(c));
+                p = p.play(c);
+            }
+            assert!(!p.is_valid_play(c));
         }
-        Ok(())
+    }
+
+    fn assert_not_winning_play(p: Position, col_pos: u64) -> Position {
+        assert!(!p.is_winning_play(col_pos), "\n{:?}", p);
+        p.play(col_pos)
+    }
+
+    #[test]
+    fn test_is_winning_play() {
+        let mut p = Position::new();
+        p = assert_not_winning_play(p, 0);
+        p = assert_not_winning_play(p, WIDTH - 1);
+        p = assert_not_winning_play(p, 1);
+        p = assert_not_winning_play(p, WIDTH - 1);
+        p = assert_not_winning_play(p, 2);
+        p = assert_not_winning_play(p, WIDTH - 1);
+        assert!(p.is_winning_play(3), "\n{:?}", p);  // horizontal
+        p = p.play(4);
+        assert!(p.is_winning_play(WIDTH - 1), "\n{:?}", p);  // vertical
+
+        p = Position::new();
+        p = assert_not_winning_play(p, 3); // w
+        p = assert_not_winning_play(p, 2);
+        p = assert_not_winning_play(p, 2); // w
+        p = assert_not_winning_play(p, 1);
+        p = assert_not_winning_play(p, 0); // w
+        p = assert_not_winning_play(p, 1);
+        p = assert_not_winning_play(p, 1); // w
+        p = assert_not_winning_play(p, 0);
+        p = assert_not_winning_play(p, 5); // w
+        p = assert_not_winning_play(p, 0);
+        assert!(p.is_winning_play(0), "\n{:?}", p);  // diagonal
+
+        p = Position::new();
+        p = assert_not_winning_play(p, 0); // w
+        p = assert_not_winning_play(p, 1);
+        p = assert_not_winning_play(p, 1); // w
+        p = assert_not_winning_play(p, 2);
+        p = assert_not_winning_play(p, 3); // w
+        p = assert_not_winning_play(p, 2);
+        p = assert_not_winning_play(p, 2); // w
+        p = assert_not_winning_play(p, 3);
+        p = assert_not_winning_play(p, 5); // w
+        p = assert_not_winning_play(p, 3);
+        assert!(p.is_winning_play(3), "\n{:?}", p);  // anti diagonal
     }
 }
 
 fn main() {
     let mut p = Position::new();
-    // println!("{:?}", p);
-    p.play(2);
-    p.play(2);
-    p.play(1);
-    p.play(5);
-    p.play(5);
+    p = p.play(2);
+    p = p.play(2);
+    p = p.play(1);
+    p = p.play(5);
     println!("{:?}", p);
-    // p.play(4);
-    // println!("------------------------------");
-    // println!("{:?}", p);
 }