src/test/bench/sudoku.rs

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