water_sort_puzzle_solver.rs

use std::cmp::min;
use std::collections::HashSet;
use std::fmt::{Debug, Formatter, Result as FmtResult};
use std::vec::Vec;

#[derive(Clone, Copy, Eq, Hash, PartialEq)]
enum Water {
    Absent,
    Present(char)
}

#[derive(Clone, Copy, Eq, Hash, PartialEq)]
struct Cup<const COUNT: usize>(
    [Water; COUNT]
);

impl<const COUNT: usize> Debug for Cup<COUNT> {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        for i in self.0.iter() {
            match i {
                Water::Absent => {
                    write!(f, "_")?;
                }
                Water::Present(c) => {
                    write!(f, "{}", c)?;
                }
            }
        }
        Ok(())
    }
}

impl<const COUNT: usize> Cup<COUNT> {

    fn new(s: &str) -> Cup<COUNT> {
        if s.len() > COUNT {
            panic!("Invalid input {}", s);
        }

        let mut arr = [Water::Absent; COUNT];
        for (i, c) in s.chars().enumerate() {
            arr[i] = Water::Present(c);
        }

        Cup(arr)
    }

    fn get_empty_slots(&self) -> usize {
        let mut slots = 0;
        for i in (0..COUNT).rev() {
            if let Water::Absent = self.0[i] {
                slots += 1;
            } else {
                break;
            }
        }

        slots
    }

    fn get_last(&self) -> Option<(char, usize, usize)> {
        let mut last: Option<char> = None;
        let mut size = 0;
        let mut idx = 0;
        for i in (0..COUNT).rev() {
            if let Water::Present(v) = self.0[i] {
                if let Some(t) = last {
                    if t != v {
                        return Some((t, idx, size));
                    } else {
                        idx = i;
                        size += 1;
                    }
                } else {
                    last = Some(v);
                    idx = i;
                    size = 1;
                }
            }
        }

        if let Some(v) = last {
            Some((v, idx, size))
        } else {
            None
        }
    }

    fn is_complete_or_empty(&self) -> bool {
        let mut cur = None;
        let mut count = 0;
        for x in self.0.iter() {
            if let Water::Present(p) = x {
                if cur.is_none() {
                    cur = Some(p);
                    count = 1;
                } else if cur.unwrap() != p {
                    return false
                } else {
                    count += 1;
                }
            }
        }

        count == 0 || count == COUNT
    }
}

#[derive(Clone, Eq, PartialEq)]
struct Board<const CUPSIZE: usize, const COUNT: usize>(
    [Cup<CUPSIZE>; COUNT]
);

impl<const CUPSIZE: usize, const COUNT: usize> Debug for Board<CUPSIZE, COUNT> {
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        write!(f, "[")?;
        for (idx, item) in self.0.iter().enumerate() {
            match idx {
                0 => {
                    write!(f, "{:?}", item)?;
                }
                _ => {
                    write!(f, " {:?}", item)?;
                }
            }
        }
        write!(f, "]")?;
        Ok(())
    }
}

impl<const CUPSIZE: usize, const COUNT: usize> Board<CUPSIZE, COUNT> {
    fn new(arr: Vec<&str>) -> Board<CUPSIZE, COUNT> {
        let mut c = [Cup::<CUPSIZE>::new(""); COUNT];

        for (idx, x) in arr.iter().enumerate() {
            c[idx] = Cup::<CUPSIZE>::new(x);
        }
        
        Board(c)
    }

    fn test(&self) -> bool {
        for x in self.0.iter() {
            if !x.is_complete_or_empty() {
                return false;
            }
        }

        true
    }

    fn can_move_to(&self, a: usize, b: usize) -> bool {
        let from = &self.0[a];
        let to = &self.0[b];
        
        let dest_slots = to.get_empty_slots();

        if dest_slots == 0 {
            false
        } else {
            if let Some((sc, _, ss)) = from.get_last() {
                if let Some((dc, _, _)) = to.get_last() {
                    if dc != sc {
                        false
                    } else {
                        true
                    }
                } else {
                    dest_slots != CUPSIZE || ss != CUPSIZE
                }
            } else {
                false
            }
        }
    }

    fn do_move(&mut self, a: usize, b: usize) -> bool {
        let from = &self.0[a];
        let to = &self.0[b];
        
        let start = if let Some((_, i, p)) = to.get_last() {
            if i + p > CUPSIZE {
                return false;
            }
            i + p
        } else {
            0
        };
        let empty_slots = to.get_empty_slots();
        if start >= CUPSIZE || empty_slots == 0 {
            false
        } else {
            if let Some((sc, si, sp)) = from.get_last() {
                let count = min(empty_slots, sp);
                for (idx, i) in (start..start + count).enumerate() {
                    self.0[b].0[i] = Water::Present(sc);
                    self.0[a].0[si + sp - 1 - idx] = Water::Absent;
                }
                true
            } else {
                false
            }
        }
    }

    fn get_candidates(&self,
                      stack: &Vec<(usize, usize, Board<CUPSIZE, COUNT>)>
    ) -> Vec<(usize, usize, Board<CUPSIZE, COUNT>)> {
        let mut out = Vec::new();

        for i in 0..self.0.len() {
            let mut set = HashSet::new();
            for j in 0..self.0.len() {
                if i == j {
                    continue;
                }
                if set.contains(&self.0[j]) {
                    continue;
                }
                set.insert(self.0[j].clone());

                if self.can_move_to(i, j) {
                    let mut moved = self.clone();
                    if moved.do_move(i, j) {
                        if stack.iter().rev().any(|e| e.2 == moved) {
                            break;
                        }

                        out.push((i, j, moved));
                    }
                }
            }
        }

        out
    }
}

fn recurse<const CUPSIZE: usize, const COUNT: usize>(
    iteration: usize,
    cur: &Board<CUPSIZE, COUNT>,
    stack: &mut Vec<(usize, usize, Board<CUPSIZE, COUNT>)>
) -> bool {
    for (i, j, moved) in cur.get_candidates(&stack) {
        stack.push((i, j, moved.clone()));
        
        if moved.test() {
            return true;
        }
        
        if recurse(iteration + 1, &moved, stack) {
            return true;
        }
        
        stack.pop();
    }
    false
}

fn resolve<const CUPSIZE: usize, const COUNT: usize>(
    cur: &Board<CUPSIZE, COUNT>
) -> Option<Vec<(usize, usize, Board<CUPSIZE, COUNT>)>> {
    let mut stack = vec![(0, 0, cur.clone())];
    
    if recurse(0, cur, &mut stack) {
        Some(stack)
    } else {
        None
    }
}

fn main() {
    let board = Board::<4, 14>::new(vec![
        "1234", "5678", "7496", "a891", "7ab6", "1736", "c532",
        "95c8", "45bb", "29ac", "2143", "bca8"
    ]);

    match resolve(&board) {
        None => {
            println!("No solution.");
        }
        Some(e) => {
            for i in e {
                println!("{:?} from {} to {}", i.2, i.0, i.1);
            }
        }
    }
}