Sudoku Solver using Z3

sudoku_solver.rs

extern crate z3;

fn main() {
    let mut initial_s: Vec<&str> = vec![
        " 9 4     ",
        "   7 6 1 ",
        "4   58   ",
        "9 6   85 ",
        "8       3",
        " 51   2 9",
        "   64   2",
        " 3 2 9   ",
        "     1 4 ",
    ];
    let initial = initial_s
        .drain(0..9)
        .collect::<String>()
        .char_indices()
        .filter(|(_, c)| *c != ' ')
        .map(|(i, c)| (i, c.to_digit(10).expect("Bad integer") as i64))
        .collect::<Vec<(usize, i64)>>();

    sudoku(&initial)
}

fn create_board(context: &z3::Context) -> Vec<z3::Ast> {
    (0..81)
        .map(|_| context.fresh_int_const("board"))
        .collect::<Vec<z3::Ast>>()
}

fn add_group_constraints(context: &z3::Context, solver: &z3::Solver, cells: &[&z3::Ast]) {
    for i in 1..=9 {
        let mut n_of_i: Vec<z3::Ast> = Vec::with_capacity(9);
        for cell in cells {
            let is_i = cell
                ._eq(&context.from_i64(i))
                .ite(&context.from_i64(1), &context.from_i64(0));
            n_of_i.push(is_i);
        }
        let n_of_i_refs = n_of_i.iter().collect::<Vec<&z3::Ast>>();
        solver.assert(
            &context
                .from_i64(0)
                .add(&n_of_i_refs)
                ._eq(&context.from_i64(1)),
        );
    }
}

fn add_row_constraints(context: &z3::Context, solver: &z3::Solver, board: &[&z3::Ast]) {
    for i in (0..81).step_by(9) {
        add_group_constraints(&context, &solver, &board[i..(i + 9)]);
    }
}

fn add_col_constraints(context: &z3::Context, solver: &z3::Solver, board: &[&z3::Ast]) {
    for i in 0..9 {
        let mut col_items: Vec<&z3::Ast> = Vec::with_capacity(9);
        for j in (0..81).step_by(9) {
            col_items.push(board[i + j]);
        }
        add_group_constraints(&context, &solver, &col_items);
    }
}

fn add_square_constraints(context: &z3::Context, solver: &z3::Solver, board: &[&z3::Ast]) {
    let corners: Vec<usize> = vec![0, 3, 6, 27, 30, 33, 54, 57, 60];
    let offsets: Vec<usize> = vec![0, 1, 2, 9, 10, 11, 18, 19, 20];
    for c in &corners {
        let mut square_items: Vec<&z3::Ast> = Vec::with_capacity(9);
        for o in &offsets {
            square_items.push(board[c + o]);
        }
        add_group_constraints(&context, &solver, &square_items);
    }
}

fn sudoku(initial: &[(usize, i64)]) {
    let mut config = z3::Config::new();
    config.set_model_generation(true);
    config.set_param_value("model_compress", "false");
    let context = z3::Context::new(&config);
    let solver = z3::Solver::new(&context);

    let board = create_board(&context);
    let board_refs = board.iter().collect::<Vec<&z3::Ast>>();

    add_row_constraints(&context, &solver, &board_refs);
    add_col_constraints(&context, &solver, &board_refs);
    add_square_constraints(&context, &solver, &board_refs);

    let assumptions = initial
        .iter()
        .map(|(i, v)| context.from_i64(*v)._eq(&board[*i]))
        .collect::<Vec<z3::Ast>>();

    println!("Solving...");
    if solver.check_assumptions(&assumptions[..]) {
        let model = solver.get_model();

        for i in (0..81).step_by(9) {
            for j in 0..9 {
                let val = model.eval(&board[i + j]).unwrap().as_i64().unwrap();
                print!("{} ", val);
            }
            println!();
        }
    } else {
        println!("No solution found!");
    }
}