ubongo_trigo_solver/src/lib.rs

use std::collections::HashMap;
use std::io::Write;


#[derive(Clone,Debug,PartialEq)]
pub struct Coord {
    pub a: i8,
    pub b: i8,
    pub c: i8,
}

impl Coord {

    fn from(a : i8, b: i8, c: i8) -> Coord {
        Coord {
            a: a,
            b: b,
            c: c
        }
    }

    fn from_ab(a : i8, b : i8, upright: bool) -> Coord {
        Coord {
            a: a,
            b: b,
            c: if upright { 2 - a - b} else { 1 - a - b },
        }

    }

    /**
     * Rotates coordinate around the origin by 120°
     */
    fn rotate(&self, rotations: u8) -> Coord {
        match rotations % 3 {
            0 => self.clone(),
            1 => Coord::from(self.b, self.c, self.a),
            _ => Coord::from(self.c, self.a, self.b),
        }
    }

    pub fn rotate60(&self) -> Coord {
        Coord {
            a: 1 - self.c,
            b: 1 - self.a,
            c: 1 - self.b
        }
    }

    /**
     * Flip coordinate at the y-axis
     */
    fn flip(&self) -> Coord {
        Coord::from(self.c, self.b, self.a)
    }

    fn upright(&self) -> bool {
        self.a + self.b + self.c == 2
    }

    pub fn to_cartesian(&self) -> (i8, i8) {
        (self.a - self.c + 1, 1 - self.b)
    }

    pub fn translate_x(&self, x: i8) -> Coord {
        Coord::from_ab(self.a + x, self.b, self.upright())
    }

    pub fn translate_y(&self, y: i8, origin_upright: bool) -> Coord {
        if origin_upright {
            Coord::from_ab(self.a + (y + 1) / 2, self.b - y, self.upright())
        } else {
            Coord::from_ab(self.a + y / 2, self.b - y, self.upright())
        }
    }
}

impl std::ops::Add<Coord> for Coord {
    type Output = Coord;

    fn add(self, _rhs: Coord) -> Coord {
        return Coord {
            a: self.a + _rhs.a,
            b: self.b + _rhs.b,
            c: self.c + _rhs.c
        }
    }
}

/* Numbers 1..7 are used to identify a given part */
pub type PartID = u8;

pub type Part = Vec<Coord>;

/** Allow part definition with tuples for brevity
 */
fn generate_part(arr: &[(i8, i8, i8)]) -> Part {
    let v: Vec<Coord> = arr.iter()
        .map(|(a,b,c)| Coord::from(*a,*b,*c))
        .collect();

    return v;
}

pub fn generate_parts() -> Vec<(PartID, Part)> {
    vec![
        (1, generate_part(&[(0,1,1), (0,0,1), (1,0,1), (1,0,0), (2,0,0), (1,-1,1)])),
        (2, generate_part(&[(0,1,1), (0,1,0), (0,0,1), (1,0,1), (1,0,0), (2,0,0)])),
        (3, generate_part(&[(0,1,1), (0,0,1), (1,0,1), (1,0,0), (2,0,0), (2, -1, 0)])),
        (4, generate_part(&[(0,1,1), (0,1,0), (1,1,0), (1,0,0), (2,0,0), (2,-1,0)])),
        (5, generate_part(&[(0,1,1), (0,1,0), (1,1,0), (1,1,-1), (0,0,1), (0,0,2)])),
        (6, generate_part(&[(0,1,1), (0,0,1), (1,0,1), (1,-1,1), (2,-1,1), (1,-1,2)])),
        (7, generate_part(&[(0,1,1), (0,0,1), (1,0,1), (1,0,0), (1,1,0), (2,0,0)])),
    ]
}

/*
 * For in-memory representation however, we simply use a two-dimensional array of cells.
 * Each cell can be a Barrier (meaning no part may be placed on the cell), Empty (meaning at the
 * current time, it is not occupied by a part) or Occupied.
 */
#[derive(PartialEq, Debug, Copy, Clone)]
pub enum Cell {
    Barrier,
    Empty,
    Occupied(PartID)
}

impl Cell {
    fn is_empty(&self) -> bool {
        match self {
            Cell::Empty => true,
            Cell::Barrier => false,
            Cell::Occupied(_) => false,
        }
    }
}

impl std::fmt::Display for Cell {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Cell::Barrier => write!(f, "B"),
            Cell::Empty => write!(f, " "),
            Cell::Occupied(id) => write!(f, "{}", id),
        }
    }

}

/*
 * A triangular grid looks like this:
 *       __  __
 *     /\  /\  /\
 *    /__\/__\/__\
 *    \  /\  /\  /
 *     \/__\/__\/
 *     /\  /\  /\
 *    /__\/__\/__\
 * 
 * We require by definition that the left upper-most triangle is upright.
 * A map is then represented by a two-dimensional row-major (meaning the outer array represents the
 * different rows) two-dimensional array with a maximum
 * size of 16x16.
 * Since there are two sides, we store both of them in a tuple.
 */
pub const MAP_WIDTH : i8 = 16;
pub const MAP_HEIGHT : i8 = 5;
pub type Map = [[Cell; MAP_WIDTH as usize]; MAP_HEIGHT as usize];


/**
 * Helper function that returns true if the coordinates are within map bounds
 */
pub fn in_bound(x: i8, y: i8) -> bool {
    x >= 0 && y >= 0 && x < MAP_WIDTH && y < MAP_HEIGHT
}


/**
 * Call f for each triangle of the part at least. If f returns false, iteration will stop and the
 * function will return false. Only returns true if all calls of r returned true.
 */
pub fn for_each_triangle<F>(part: &Part, x: i8, y: i8, rotations: u8, flipped: bool, mut f: F) -> bool
where F: FnMut(i8, i8) -> bool
{
    for coordinate in part {
        let mut coord : Coord = if flipped { coordinate.flip() } else { coordinate.clone() };
        coord = coord.rotate(rotations);


        let origin_upright = (x + y) % 2 == 0;
        if !origin_upright {
            coord = coord.rotate60();
        }


        coord = coord.translate_y(y, origin_upright).translate_x(x / 2);

        let (wx, wy) = coord.to_cartesian();

        if in_bound(wx, wy) == false {
            return false;
        }

        if f(wx, wy) == false {
            return false;
        }
    }


    true
}


/*
 * Parts can be inverted (flipped on the other side) and also rotated.
 * On a triangular grid, three rotations are possible.
 * 
 * The task of our core logic function below is to figure out whether a given part can fit on a
 * specified location on the map.
 *
 * Returns true if the part can fit at the specified location on the map.
 */
pub fn check_part(map: &Map, part: &Part, x: i8, y: i8, rotations: u8, flipped: bool) -> bool {
    /* Make sure the start triangle can actually be placed */
    debug_assert!(in_bound(x, y));

    //println!("Testing {},{},{},{}", x, y, rotation.to_int(), flipped);
    for_each_triangle(part, x, y, rotations, flipped, |x, y| {
        if !in_bound(x,y) {
            return false;
        }

        let prev = map[y as usize][x as usize];
        if !(prev.is_empty()) {
            return false;
        }


        return true;
    })
}

///* Once we have found a valid position with `check_part`, we can place the part on the map.
// */
pub fn place_part(map: &mut Map, id: PartID, part: &Part, x: i8, y: i8, rotations: u8, flipped: bool) {
    //println!("Placing {},{},{},{}", x, y, rotation.to_int(), flipped);
    for_each_triangle(part, x, y, rotations, flipped, |x, y| {
        if !in_bound(x,y) {
            return false;
        }

        map[y as usize][x as usize] = Cell::Occupied(id);

        return true;
    });
}

/* If a part is not placed right, we might need to remove it from the map again
 */
fn erase_part(map: &mut Map, part: &Part, x: i8, y: i8, rotation: u8, flipped: bool) {
    for_each_triangle(part, x, y, rotation, flipped, |x, y| {
        debug_assert!(in_bound(x, y));

        map[y as usize][x as usize] = Cell::Empty;
        true
    });
}

pub fn solve(map: &mut Map, parts: &[(PartID, Part)]) -> Option<()>
{
    if parts.len() == 0 {
        return Some(());
    }

    let (id, part) = &parts[0];

    // TODO: restrict to list of non-barrier position to tighten the for loops
    for y in 0..MAP_HEIGHT {
        for x in 0..MAP_WIDTH {
            if !map[y as usize][x as usize].is_empty() {
                continue;
            }

            for rotation in [0, 1, 2] {
                for flipped in [false, true] {
                    let valid_position = check_part(&map, &part, x, y, rotation, flipped);

                    if valid_position {
                        // Place our part and try the others
                        place_part(map, *id, &part, x, y, rotation, flipped);

                        //t.as_ref().map(|&&&mut Term| print_map(term, map));
                        
                        if let Some(()) = solve(map, &parts[1..]) {
                            // We have found a valid solution we can pass back!
                            return Some(());
                        } else {
                            erase_part(map, &part, x, y, rotation, flipped);
                        }
                    }
                }
            }
        }
    }

    None
}

pub fn foreground_color(idx: u8) {
    print!("\x1B[38;5;{}m", idx);
}

pub fn background_color(idx: u8) {
    print!("\x1B[48;5;{}m", idx);
}

fn color_for_cell(cell: Cell) -> console::Color {
    match cell {
        Cell::Barrier => console::Color::Color256(0),
        Cell::Empty => console::Color::Color256(15),
        Cell::Occupied(1) => console::Color::Color256(94),
        Cell::Occupied(2) => console::Color::Color256(89),
        Cell::Occupied(3) => console::Color::Color256(202),
        Cell::Occupied(4) => console::Color::Color256(220),
        Cell::Occupied(5) => console::Color::Color256(22),
        Cell::Occupied(6) => console::Color::Color256(196),
        Cell::Occupied(7) => console::Color::Color256(27),
        _ => console::Color::Black
    }
}

//  |O###O###O###
//  OOO#OOO#OOO#
//  ###O###O###O
//  |#OOO#OOO#OOO
use console::style;
pub fn print_map(t: &mut console::Term, m: &Map) {
    for y in 0..MAP_HEIGHT {
        if y % 2 == 0 {
            write!(t, "{}", style("█").fg(color_for_cell(Cell::Barrier)));
        }
        for x in 0..MAP_WIDTH {
            let c = if (x + y) % 2 == 0 {
                "∧"
            } else {
                "\\#/"
            };

            write!(t, "{}", style(c).fg(color_for_cell(m[y as usize][x as usize])));
        }
        
        write!(t, "{}", style("\n").fg(color_for_cell(Cell::Barrier)));

        if y % 2 == 1 {
            write!(t, "{}", style("█").fg(color_for_cell(Cell::Barrier)));
        }
        for x in 0..MAP_WIDTH {
            let c = if (x + y) % 2 == 1 {
                 "v"
            } else {
                "/_\\"
            };

            write!(t, "{}", style(c).fg(color_for_cell(m[y as usize][x as usize])));
        }
        write!(t, "{}", style("\n").fg(color_for_cell(Cell::Barrier)));

    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_to_cartesian_coords() {
        assert_eq!((0,0), Coord::from(0,1,1).to_cartesian());
        assert_eq!((1,0), Coord::from(0,1,0).to_cartesian());
        assert_eq!((2,0), Coord::from(1,1,0).to_cartesian());
        assert_eq!((0,1), Coord::from(0,0,1).to_cartesian());

        assert_eq!((0,-1), Coord::from(-1, 2, 0).to_cartesian());

        assert_eq!((2,1), Coord::from(1,0,0).to_cartesian());
        assert_eq!((3,1), Coord::from(2,0,0).to_cartesian());

    }



    #[test]
    fn test_translation() {
        assert_eq!(Coord::from(1,1,0), Coord::from(-1,1,2).translate_x(2));
        assert_eq!(Coord::from(-1,0,3), Coord::from(2,0,0).translate_x(-3));
        assert_eq!(Coord::from(0,1,1), Coord::from(-1,3,0).translate_y(2, true));
        assert_eq!(Coord::from(3,0,-1), Coord::from(3,-1,0).translate_y(-1, true));

    }
}