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