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::intrinsics::cttz16;
  20 use std::slice;
  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 = slice::from_fn(9u, |i| {
  52             slice::from_fn(9u, |j| { vec[i][j] })
  53         });
  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][col] != other.grid[row][col] {
  61                     return false;
  62                 }
  63             }
  64         }
  65         return true;
  66     }
  67
  68     pub fn read(mut reader: BufferedReader<StdReader>) -> Sudoku {
  69         assert!(reader.read_line().unwrap() == ~"9,9"); /* assert first line is exactly "9,9" */
  70
  71         let mut g = slice::from_fn(10u, { |_i| vec!(0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8) });
  72         for line in reader.lines() {
  73             let comps: Vec<&str> = line.trim().split(',').collect();
  74
  75             if comps.len() == 3u {
  76                 let row     = from_str::<uint>(comps[0]).unwrap() as u8;
  77                 let col     = from_str::<uint>(comps[1]).unwrap() as u8;
  78                 g[row][col] = from_str::<uint>(comps[2]).unwrap() as u8;
  79             }
  80             else {
  81                 fail!("Invalid sudoku file");
  82             }
  83         }
  84         return Sudoku::new(g)
  85     }
  86
  87     pub fn write(&self, writer: &mut io::Writer) {
  88         for row in range(0u8, 9u8) {
  89             write!(writer, "{}", self.grid[row][0]);
  90             for col in range(1u8, 9u8) {
  91                 write!(writer, " {}", self.grid[row][col]);
  92             }
  93             write!(writer, "\n");
  94          }
  95     }
  96
  97     // solve sudoku grid
  98     pub fn solve(&mut self) {
  99         let mut work: Vec<(u8, u8)> = Vec::new(); /* queue of uncolored fields */
 100         for row in range(0u8, 9u8) {
 101             for col in range(0u8, 9u8) {
 102                 let color = self.grid[row][col];
 103                 if color == 0u8 {
 104                     work.push((row, col));
 105                 }
 106             }
 107         }
 108
 109         let mut ptr = 0u;
 110         let end = work.len();
 111         while ptr < end {
 112             let (row, col) = work[ptr];
 113             // is there another color to try?
 114             if self.next_color(row, col, self.grid[row][col] + (1 as u8)) {
 115                 //  yes: advance work list
 116                 ptr = ptr + 1u;
 117             } else {
 118                 // no: redo this field aft recoloring pred; unless there is none
 119                 if ptr == 0u { fail!("No solution found for this sudoku"); }
 120                 ptr = ptr - 1u;
 121             }
 122         }
 123     }
 124
 125     fn next_color(&mut self, row: u8, col: u8, start_color: u8) -> bool {
 126         if start_color < 10u8 {
 127             // colors not yet used
 128             let mut avail = ~Colors::new(start_color);
 129
 130             // drop colors already in use in neighbourhood
 131             self.drop_colors(avail, row, col);
 132
 133             // find first remaining color that is available
 134             let next = avail.next();
 135             self.grid[row][col] = next;
 136             return 0u8 != next;
 137         }
 138         self.grid[row][col] = 0u8;
 139         return false;
 140     }
 141
 142     // find colors available in neighbourhood of (row, col)
 143     fn drop_colors(&mut self, avail: &mut Colors, row: u8, col: u8) {
 144         for idx in range(0u8, 9u8) {
 145             avail.remove(self.grid[idx][col]); /* check same column fields */
 146             avail.remove(self.grid[row][idx]); /* check same row fields */
 147         }
 148
 149         // check same block fields
 150         let row0 = (row / 3u8) * 3u8;
 151         let col0 = (col / 3u8) * 3u8;
 152         for alt_row in range(row0, row0 + 3u8) {
 153             for alt_col in range(col0, col0 + 3u8) {
 154                 avail.remove(self.grid[alt_row][alt_col]);
 155             }
 156         }
 157     }
 158 }
 159
 160 // Stores available colors as simple bitfield, bit 0 is always unset
 161 struct Colors(u16);
 162
 163 static HEADS: u16 = (1u16 << 10) - 1; /* bits 9..0 */
 164
 165 impl Colors {
 166     fn new(start_color: u8) -> Colors {
 167         // Sets bits 9..start_color
 168         let tails = !0u16 << start_color;
 169         return Colors(HEADS & tails);
 170     }
 171
 172     fn next(&self) -> u8 {
 173         let Colors(c) = *self;
 174         let val = c & HEADS;
 175         if (0u16 == val) {
 176             return 0u8;
 177         } else {
 178             unsafe {
 179                 return cttz16(val as i16) as u8;
 180             }
 181         }
 182     }
 183
 184     fn remove(&mut self, color: u8) {
 185         if color != 0u8 {
 186             let Colors(val) = *self;
 187             let mask = !(1u16 << color);
 188             *self    = Colors(val & mask);
 189         }
 190     }
 191 }
 192
 193 static DEFAULT_SUDOKU: [[u8, ..9], ..9] = [
 194          /* 0    1    2    3    4    5    6    7    8    */
 195   /* 0 */  [0u8, 4u8, 0u8, 6u8, 0u8, 0u8, 0u8, 3u8, 2u8],
 196   /* 1 */  [0u8, 0u8, 8u8, 0u8, 2u8, 0u8, 0u8, 0u8, 0u8],
 197   /* 2 */  [7u8, 0u8, 0u8, 8u8, 0u8, 0u8, 0u8, 0u8, 0u8],
 198   /* 3 */  [0u8, 0u8, 0u8, 5u8, 0u8, 0u8, 0u8, 0u8, 0u8],
 199   /* 4 */  [0u8, 5u8, 0u8, 0u8, 0u8, 3u8, 6u8, 0u8, 0u8],
 200   /* 5 */  [6u8, 8u8, 0u8, 0u8, 0u8, 0u8, 0u8, 9u8, 0u8],
 201   /* 6 */  [0u8, 9u8, 5u8, 0u8, 0u8, 6u8, 0u8, 7u8, 0u8],
 202   /* 7 */  [0u8, 0u8, 0u8, 0u8, 4u8, 0u8, 0u8, 6u8, 0u8],
 203   /* 8 */  [4u8, 0u8, 0u8, 0u8, 0u8, 7u8, 2u8, 0u8, 3u8]
 204 ];
 205
 206 #[cfg(test)]
 207 static DEFAULT_SOLUTION: [[u8, ..9], ..9] = [
 208          /* 0    1    2    3    4    5    6    7    8    */
 209   /* 0 */  [1u8, 4u8, 9u8, 6u8, 7u8, 5u8, 8u8, 3u8, 2u8],
 210   /* 1 */  [5u8, 3u8, 8u8, 1u8, 2u8, 9u8, 7u8, 4u8, 6u8],
 211   /* 2 */  [7u8, 2u8, 6u8, 8u8, 3u8, 4u8, 1u8, 5u8, 9u8],
 212   /* 3 */  [9u8, 1u8, 4u8, 5u8, 6u8, 8u8, 3u8, 2u8, 7u8],
 213   /* 4 */  [2u8, 5u8, 7u8, 4u8, 9u8, 3u8, 6u8, 1u8, 8u8],
 214   /* 5 */  [6u8, 8u8, 3u8, 7u8, 1u8, 2u8, 5u8, 9u8, 4u8],
 215   /* 6 */  [3u8, 9u8, 5u8, 2u8, 8u8, 6u8, 4u8, 7u8, 1u8],
 216   /* 7 */  [8u8, 7u8, 2u8, 3u8, 4u8, 1u8, 9u8, 6u8, 5u8],
 217   /* 8 */  [4u8, 6u8, 1u8, 9u8, 5u8, 7u8, 2u8, 8u8, 3u8]
 218 ];
 219
 220 #[test]
 221 fn colors_new_works() {
 222     assert_eq!(*Colors::new(1), 1022u16);
 223     assert_eq!(*Colors::new(2), 1020u16);
 224     assert_eq!(*Colors::new(3), 1016u16);
 225     assert_eq!(*Colors::new(4), 1008u16);
 226     assert_eq!(*Colors::new(5), 992u16);
 227     assert_eq!(*Colors::new(6), 960u16);
 228     assert_eq!(*Colors::new(7), 896u16);
 229     assert_eq!(*Colors::new(8), 768u16);
 230     assert_eq!(*Colors::new(9), 512u16);
 231 }
 232
 233 #[test]
 234 fn colors_next_works() {
 235     assert_eq!(Colors(0).next(), 0u8);
 236     assert_eq!(Colors(2).next(), 1u8);
 237     assert_eq!(Colors(4).next(), 2u8);
 238     assert_eq!(Colors(8).next(), 3u8);
 239     assert_eq!(Colors(16).next(), 4u8);
 240     assert_eq!(Colors(32).next(), 5u8);
 241     assert_eq!(Colors(64).next(), 6u8);
 242     assert_eq!(Colors(128).next(), 7u8);
 243     assert_eq!(Colors(256).next(), 8u8);
 244     assert_eq!(Colors(512).next(), 9u8);
 245     assert_eq!(Colors(1024).next(), 0u8);
 246 }
 247
 248 #[test]
 249 fn colors_remove_works() {
 250     // GIVEN
 251     let mut colors = Colors::new(1);
 252
 253     // WHEN
 254     colors.remove(1);
 255
 256     // THEN
 257     assert_eq!(colors.next(), 2u8);
 258 }
 259
 260 #[test]
 261 fn check_DEFAULT_SUDOKU_solution() {
 262     // GIVEN
 263     let mut sudoku = Sudoku::from_vec(&DEFAULT_SUDOKU);
 264     let solution   = Sudoku::from_vec(&DEFAULT_SOLUTION);
 265
 266     // WHEN
 267     sudoku.solve();
 268
 269     // THEN
 270     assert!(sudoku.equal(&solution));
 271 }
 272
 273 fn main() {
 274     let args        = os::args();
 275     let use_default = args.len() == 1u;
 276     let mut sudoku = if use_default {
 277         Sudoku::from_vec(&DEFAULT_SUDOKU)
 278     } else {
 279         Sudoku::read(io::stdin())
 280     };
 281     sudoku.solve();
 282     sudoku.write(&mut io::stdout());
 283 }