src/test/bench/sudoku.rs

   1 // Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 // ignore-pretty very bad with line comments
  12
  13 #![allow(non_snake_case)]
  14
  15 use std::io::BufferedReader;
  16 use std::io::stdio::StdReader;
  17 use std::io;
  18 use std::iter::repeat;
  19 use std::num::Int;
  20 use std::os;
  21
  22 // Computes a single solution to a given 9x9 sudoku
  23 //
  24 // Call with "-" to read input sudoku from stdin
  25 //
  26 // The expected line-based format is:
  27 //
  28 // 9,9
  29 // <row>,<column>,<color>
  30 // ...
  31 //
  32 // Row and column are 0-based (i.e. <= 8) and color is 1-based (>=1,<=9).
  33 // A color of 0 indicates an empty field.
  34 //
  35 // If called without arguments, sudoku solves a built-in example sudoku
  36 //
  37
  38 // internal type of sudoku grids
  39 type grid = Vec<Vec<u8> > ;
  40
  41 struct Sudoku {
  42     grid: grid
  43 }
  44
  45 impl Sudoku {
  46     pub fn new(g: grid) -> Sudoku {
  47         return Sudoku { grid: g }
  48     }
  49
  50     pub fn from_vec(vec: &[[u8;9];9]) -> Sudoku {
  51         let g = range(0, 9u).map(|i| {
  52             range(0, 9u).map(|j| { vec[i][j] }).collect()
  53         }).collect();
  54         return Sudoku::new(g)
  55     }
  56
  57     pub fn equal(&self, other: &Sudoku) -> bool {
  58         for row in range(0u8, 9u8) {
  59             for col in range(0u8, 9u8) {
  60                 if self.grid[row as uint][col as uint] !=
  61                         other.grid[row as uint][col as uint] {
  62                     return false;
  63                 }
  64             }
  65         }
  66         return true;
  67     }
  68
  69     pub fn read(mut reader: &mut BufferedReader<StdReader>) -> Sudoku {
  70         /* assert first line is exactly "9,9" */
  71         assert!(reader.read_line().unwrap() == "9,9".to_string());
  72
  73         let mut g = repeat(vec![0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8])
  74                           .take(10).collect::<Vec<_>>();
  75         for line in reader.lines() {
  76             let line = line.unwrap();
  77             let comps: Vec<&str> = line.as_slice()
  78                                        .trim()
  79                                        .split(',')
  80                                        .collect();
  81
  82             if comps.len() == 3u {
  83                 let row = comps[0].parse::<u8>().unwrap();
  84                 let col = comps[1].parse::<u8>().unwrap();
  85                 g[row as uint][col as uint] = comps[2].parse().unwrap();
  86             }
  87             else {
  88                 panic!("Invalid sudoku file");
  89             }
  90         }
  91         return Sudoku::new(g)
  92     }
  93
  94     pub fn write(&self, writer: &mut io::Writer) {
  95         for row in range(0u8, 9u8) {
  96             write!(writer, "{}", self.grid[row as uint][0]);
  97             for col in range(1u8, 9u8) {
  98                 write!(writer, " {}", self.grid[row as uint][col as uint]);
  99             }
 100             write!(writer, "\n");
 101          }
 102     }
 103
 104     // solve sudoku grid
 105     pub fn solve(&mut self) {
 106         let mut work: Vec<(u8, u8)> = Vec::new(); /* queue of uncolored fields */
 107         for row in range(0u8, 9u8) {
 108             for col in range(0u8, 9u8) {
 109                 let color = self.grid[row as uint][col as uint];
 110                 if color == 0u8 {
 111                     work.push((row, col));
 112                 }
 113             }
 114         }
 115
 116         let mut ptr = 0u;
 117         let end = work.len();
 118         while ptr < end {
 119             let (row, col) = work[ptr];
 120             // is there another color to try?
 121             let the_color = self.grid[row as uint][col as uint] +
 122                                 (1 as u8);
 123             if self.next_color(row, col, the_color) {
 124                 //  yes: advance work list
 125                 ptr = ptr + 1u;
 126             } else {
 127                 // no: redo this field aft recoloring pred; unless there is none
 128                 if ptr == 0u { panic!("No solution found for this sudoku"); }
 129                 ptr = ptr - 1u;
 130             }
 131         }
 132     }
 133
 134     fn next_color(&mut self, row: u8, col: u8, start_color: u8) -> bool {
 135         if start_color < 10u8 {
 136             // colors not yet used
 137             let mut avail = box Colors::new(start_color);
 138
 139             // drop colors already in use in neighbourhood
 140             self.drop_colors(&mut *avail, row, col);
 141
 142             // find first remaining color that is available
 143             let next = avail.next();
 144             self.grid[row as uint][col as uint] = next;
 145             return 0u8 != next;
 146         }
 147         self.grid[row as uint][col as uint] = 0u8;
 148         return false;
 149     }
 150
 151     // find colors available in neighbourhood of (row, col)
 152     fn drop_colors(&mut self, avail: &mut Colors, row: u8, col: u8) {
 153         for idx in range(0u8, 9u8) {
 154             /* check same column fields */
 155             avail.remove(self.grid[idx as uint][col as uint]);
 156             /* check same row fields */
 157             avail.remove(self.grid[row as uint][idx as uint]);
 158         }
 159
 160         // check same block fields
 161         let row0 = (row / 3u8) * 3u8;
 162         let col0 = (col / 3u8) * 3u8;
 163         for alt_row in range(row0, row0 + 3u8) {
 164             for alt_col in range(col0, col0 + 3u8) {
 165                 avail.remove(self.grid[alt_row as uint][alt_col as uint]);
 166             }
 167         }
 168     }
 169 }
 170
 171 // Stores available colors as simple bitfield, bit 0 is always unset
 172 struct Colors(u16);
 173
 174 static HEADS: u16 = (1u16 << 10) - 1; /* bits 9..0 */
 175
 176 impl Colors {
 177     fn new(start_color: u8) -> Colors {
 178         // Sets bits 9..start_color
 179         let tails = !0u16 << start_color as uint;
 180         return Colors(HEADS & tails);
 181     }
 182
 183     fn next(&self) -> u8 {
 184         let Colors(c) = *self;
 185         let val = c & HEADS;
 186         if (0u16 == val) {
 187             return 0u8;
 188         } else {
 189             return val.trailing_zeros() as u8
 190         }
 191     }
 192
 193     fn remove(&mut self, color: u8) {
 194         if color != 0u8 {
 195             let Colors(val) = *self;
 196             let mask = !(1u16 << color as uint);
 197             *self    = Colors(val & mask);
 198         }
 199     }
 200 }
 201
 202 static DEFAULT_SUDOKU: [[u8;9];9] = [
 203          /* 0    1    2    3    4    5    6    7    8    */
 204   /* 0 */  [0u8, 4u8, 0u8, 6u8, 0u8, 0u8, 0u8, 3u8, 2u8],
 205   /* 1 */  [0u8, 0u8, 8u8, 0u8, 2u8, 0u8, 0u8, 0u8, 0u8],
 206   /* 2 */  [7u8, 0u8, 0u8, 8u8, 0u8, 0u8, 0u8, 0u8, 0u8],
 207   /* 3 */  [0u8, 0u8, 0u8, 5u8, 0u8, 0u8, 0u8, 0u8, 0u8],
 208   /* 4 */  [0u8, 5u8, 0u8, 0u8, 0u8, 3u8, 6u8, 0u8, 0u8],
 209   /* 5 */  [6u8, 8u8, 0u8, 0u8, 0u8, 0u8, 0u8, 9u8, 0u8],
 210   /* 6 */  [0u8, 9u8, 5u8, 0u8, 0u8, 6u8, 0u8, 7u8, 0u8],
 211   /* 7 */  [0u8, 0u8, 0u8, 0u8, 4u8, 0u8, 0u8, 6u8, 0u8],
 212   /* 8 */  [4u8, 0u8, 0u8, 0u8, 0u8, 7u8, 2u8, 0u8, 3u8]
 213 ];
 214
 215 #[cfg(test)]
 216 static DEFAULT_SOLUTION: [[u8;9];9] = [
 217          /* 0    1    2    3    4    5    6    7    8    */
 218   /* 0 */  [1u8, 4u8, 9u8, 6u8, 7u8, 5u8, 8u8, 3u8, 2u8],
 219   /* 1 */  [5u8, 3u8, 8u8, 1u8, 2u8, 9u8, 7u8, 4u8, 6u8],
 220   /* 2 */  [7u8, 2u8, 6u8, 8u8, 3u8, 4u8, 1u8, 5u8, 9u8],
 221   /* 3 */  [9u8, 1u8, 4u8, 5u8, 6u8, 8u8, 3u8, 2u8, 7u8],
 222   /* 4 */  [2u8, 5u8, 7u8, 4u8, 9u8, 3u8, 6u8, 1u8, 8u8],
 223   /* 5 */  [6u8, 8u8, 3u8, 7u8, 1u8, 2u8, 5u8, 9u8, 4u8],
 224   /* 6 */  [3u8, 9u8, 5u8, 2u8, 8u8, 6u8, 4u8, 7u8, 1u8],
 225   /* 7 */  [8u8, 7u8, 2u8, 3u8, 4u8, 1u8, 9u8, 6u8, 5u8],
 226   /* 8 */  [4u8, 6u8, 1u8, 9u8, 5u8, 7u8, 2u8, 8u8, 3u8]
 227 ];
 228
 229 #[test]
 230 fn colors_new_works() {
 231     assert_eq!(*Colors::new(1), 1022u16);
 232     assert_eq!(*Colors::new(2), 1020u16);
 233     assert_eq!(*Colors::new(3), 1016u16);
 234     assert_eq!(*Colors::new(4), 1008u16);
 235     assert_eq!(*Colors::new(5), 992u16);
 236     assert_eq!(*Colors::new(6), 960u16);
 237     assert_eq!(*Colors::new(7), 896u16);
 238     assert_eq!(*Colors::new(8), 768u16);
 239     assert_eq!(*Colors::new(9), 512u16);
 240 }
 241
 242 #[test]
 243 fn colors_next_works() {
 244     assert_eq!(Colors(0).next(), 0u8);
 245     assert_eq!(Colors(2).next(), 1u8);
 246     assert_eq!(Colors(4).next(), 2u8);
 247     assert_eq!(Colors(8).next(), 3u8);
 248     assert_eq!(Colors(16).next(), 4u8);
 249     assert_eq!(Colors(32).next(), 5u8);
 250     assert_eq!(Colors(64).next(), 6u8);
 251     assert_eq!(Colors(128).next(), 7u8);
 252     assert_eq!(Colors(256).next(), 8u8);
 253     assert_eq!(Colors(512).next(), 9u8);
 254     assert_eq!(Colors(1024).next(), 0u8);
 255 }
 256
 257 #[test]
 258 fn colors_remove_works() {
 259     // GIVEN
 260     let mut colors = Colors::new(1);
 261
 262     // WHEN
 263     colors.remove(1);
 264
 265     // THEN
 266     assert_eq!(colors.next(), 2u8);
 267 }
 268
 269 #[test]
 270 fn check_DEFAULT_SUDOKU_solution() {
 271     // GIVEN
 272     let mut sudoku = Sudoku::from_vec(&DEFAULT_SUDOKU);
 273     let solution   = Sudoku::from_vec(&DEFAULT_SOLUTION);
 274
 275     // WHEN
 276     sudoku.solve();
 277
 278     // THEN
 279     assert!(sudoku.equal(&solution));
 280 }
 281
 282 fn main() {
 283     let args        = os::args();
 284     let use_default = args.len() == 1u;
 285     let mut sudoku = if use_default {
 286         Sudoku::from_vec(&DEFAULT_SUDOKU)
 287     } else {
 288         Sudoku::read(&mut *io::stdin().lock())
 289     };
 290     sudoku.solve();
 291     sudoku.write(&mut io::stdout());
 292 }