<commit_before>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use model::{Color, FieldCoord, HexCoord};

#[derive(Clone, Copy)]
pub struct Board {
    /*
                ____
               /    \
          ____/  +y  \____
         /    \      /    \
        /      \____/  +x  \
        \      /    \      /
         \____/      \____/
         /    \      /    \
        /  -x  \____/      \
        \      /    \      /
         \____/  -y  \____/
              \      /
               \____/

        The hex board uses an axial coordinate system with (0, 0) at the center. The x-axis slopes
        up to the right, and the y-axis goes up and down. The board is stored as a dense 2D array,
        because a ragged array won't work (since each row would have to be a different type).

        See http://www.redblobgames.com/grids/hexagons/#coordinates-axial for more info.
    */
    board: [[Option<Hex>; 5]; 5],
    turn: Color,
    black_pieces: u32,
    black_hexes: u32,
    white_pieces: u32,
    white_hexes: u32,
}

// Fields are numbered clockwise from the top. Even indicies are black, odd indicies are white.
type Hex = [Field; 6];

// There's no need to store the color of the piece on this field, since only white pieces can be on
// white fields, and vice versa.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Field {
    Piece,
    Empty,
}

// Public methods
impl Board {
    #[cfg_attr(rustfmt, rustfmt_skip)]
    /// Create a new board with the "Laurentius" starting position.
    pub fn new() -> Board {
        let mut board = Board {
            board: [[None; 5]; 5],
            turn: Color::White,
            black_pieces: 18,
            black_hexes: 0,
            white_pieces: 18,
            white_hexes: 0,
        };

        // (0, 0) is the only empty hex.
        board.set_hex(&HexCoord::new(0, 0), Some([Field::Empty; 6]));

        // All other hexes have exactly two pieces on them in the starting position.
        let piece_locations = [
            (-2,  2, 0, 4),
            (-2,  1, 0, 3),
            (-2,  0, 3, 5),
            (-1,  2, 1, 4),
            (-1,  1, 0, 4),
            (-1,  0, 3, 5),
            (-1, -1, 2, 5),
            ( 0,  2, 1, 5),
            ( 0,  1, 1, 5),
            ( 0, -1, 2, 4),
            ( 0, -2, 2, 4),
            ( 1,  1, 2, 5),
            ( 1,  0, 0, 2),
            ( 1, -1, 1, 3),
            ( 1, -2, 1, 4),
            ( 2,  0, 0, 2),
            ( 2, -1, 0, 3),
            ( 2, -2, 1, 3),
        ];

        for &(x, y, f1, f2) in &piece_locations {
            let mut hex = [Field::Empty; 6];
            hex[f1] = Field::Piece;
            hex[f2] = Field::Piece;
            board.set_hex(&HexCoord::new(x, y), Some(hex));
        }

        board
    }
    pub fn can_move_piece(&self, from: &FieldCoord, to: &FieldCoord) -> bool {
        from.color() == self.turn && self.is_piece_on_field(from) && !self.is_piece_on_field(to)
            && self.get_field_vertex_neighbors(from).contains(to)
    }
    pub fn move_piece(&mut self, from: &FieldCoord, to: &FieldCoord) {
        assert!(self.can_move_piece(from, to));

        self.set_field(from, Field::Empty);
        self.set_field(to, Field::Piece);

        let (capture_count, mut fields_to_check) = self.check_hexes(&from.to_hex());
        fields_to_check.append(&mut self.get_field_edge_neighbors(to));
        self.check_captures(&fields_to_check);

        match self.turn {
            Color::White => self.white_hexes += capture_count,
            Color::Black => self.black_hexes += capture_count,
        }
        self.turn = self.turn.switch();
    }
    pub fn get_available_moves(&self, field: &FieldCoord) -> Vec<FieldCoord> {
        if self.is_piece_on_field(field) {
            self.get_field_vertex_neighbors(field)
                .into_iter()
                .filter(|c| !self.is_piece_on_field(c))
                .collect()
        } else {
            vec![]
        }
    }
    pub fn can_exchange(&self) -> bool {
        2 <= match self.turn {
            Color::Black => self.black_hexes,
            Color::White => self.white_hexes,
        }
    }
    pub fn can_exchange_piece(&self, coord: &FieldCoord) -> bool {
        self.can_exchange() && coord.color() != self.turn && self.is_piece_on_field(coord)
    }
    pub fn exchange_piece(&mut self, coord: &FieldCoord) {
        assert!(self.can_exchange_piece(coord));

        self.remove_piece(coord);
        match self.turn {
            Color::White => self.white_hexes -= 2,
            Color::Black => self.black_hexes -= 2,
        }

        // Players don't collect hexes removed due to an exchange
        let (_, fields_to_check) = self.check_hexes(&coord.to_hex());
        self.check_captures(&fields_to_check);
        self.turn = self.turn.switch();
    }
    pub fn is_piece_on_field(&self, coord: &FieldCoord) -> bool {
        self.get_field(coord) == &Field::Piece
    }
    /// > extant (adj.): Still in existence; not destroyed, lost, or extinct (The Free Dictionary)
    ///
    /// Return the coordinates of the hexes that have not been removed yet.
    pub fn extant_hexes(&self) -> Vec<HexCoord> {
        let mut coords = vec![];
        for x in -2..3 {
            for y in -2..3 {
                coords.push((x, y));
            }
        }
        coords
            .iter()
            .filter_map(|&coord| self.try_hex(coord))
            .collect()
    }
    pub fn turn(&self) -> Color {
        self.turn
    }
    pub fn black_pieces(&self) -> u32 {
        self.black_pieces
    }
    pub fn black_hexes(&self) -> u32 {
        self.black_hexes
    }
    pub fn white_pieces(&self) -> u32 {
        self.white_pieces
    }
    pub fn white_hexes(&self) -> u32 {
        self.white_hexes
    }
}

// Field and piece methods
impl Board {
    fn get_field(&self, coord: &FieldCoord) -> &Field {
        match *self.get_hex(&coord.to_hex()) {
            Some(ref hex) => &hex[coord.f as usize],
            None => panic!(
                "Cannot get field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    fn set_field(&mut self, coord: &FieldCoord, field: Field) {
        match *self.get_hex(&coord.to_hex()) {
            Some(mut hex) => {
                hex[coord.f as usize] = field;
                self.set_hex(&coord.to_hex(), Some(hex));
            }
            None => panic!(
                "Cannot set field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    /// Return fields that share an edge with the given field. These fields are always the opposite
    /// color of the given field. If all of a piece's edge neighbors are occupied, that piece might
    /// be capturable.
    fn get_field_edge_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let mut neighbors = vec![
            // There are always two edge neighbors on the same hex as the given field
            FieldCoord::new(coord.x, coord.y, (coord.f + 1) % 6),
            FieldCoord::new(coord.x, coord.y, (coord.f + 5) % 6),
        ];

        if let Some(hex) = self.get_hex_neighbor(&coord.to_hex(), coord.f) {
            neighbors.push(hex.to_field((coord.f + 3) % 6));
        }

        neighbors
    }
    /// Return fields that share a vertex with the given field and have the same color as the given
    /// field. Pieces can move to fields that are vertex neighbors of the field they are on.
    fn get_field_vertex_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        // A field's vertex neighbors can be defined as the edge neighbors of its edge neighbors
        let mut neighbors = vec![];
        for field in self.get_field_edge_neighbors(coord) {
            neighbors.append(&mut self.get_field_edge_neighbors(&field));
        }
        // A field is not its own neighbor
        neighbors.into_iter().filter(|n| n != coord).collect()
    }
    fn remove_piece(&mut self, coord: &FieldCoord) {
        assert!(
            self.is_piece_on_field(coord),
            "There is no piece at {:?} to remove",
            coord
        );
        self.set_field(coord, Field::Empty);
        match coord.color() {
            Color::Black => self.black_pieces -= 1,
            Color::White => self.white_pieces -= 1,
        }
    }
    fn check_captures(&mut self, fields_to_check: &[FieldCoord]) {
        for field in fields_to_check {
            if field.color() != self.turn && self.is_piece_on_field(field)
                && self.get_field_edge_neighbors(field)
                    .iter()
                    .all(|coord| self.is_piece_on_field(coord))
            {
                self.remove_piece(field);
            }
        }
    }
}

// Hex methods
impl Board {
    fn get_hex(&self, coord: &HexCoord) -> &Option<Hex> {
        &self.board[(coord.x + 2) as usize][(coord.y + 2) as usize]
    }
    fn set_hex(&mut self, coord: &HexCoord, hex: Option<Hex>) {
        self.board[(coord.x + 2) as usize][(coord.y + 2) as usize] = hex;
    }
    fn get_hex_neighbor(&self, coord: &HexCoord, direction: u32) -> Option<HexCoord> {
        assert!(direction < 6);

        let neighbors = [
            (coord.x, coord.y + 1),
            (coord.x + 1, coord.y),
            (coord.x + 1, coord.y - 1),
            (coord.x, coord.y - 1),
            (coord.x - 1, coord.y),
            (coord.x - 1, coord.y + 1),
        ];

        self.try_hex(neighbors[direction as usize])
    }
    /// A hex is removable (and must be removed) if it is empty and is "attached to the board by 3
    /// or less adjacent sides."
    fn is_hex_removable(&self, coord: &HexCoord) -> bool {
        match *self.get_hex(coord) {
            Some(hex) => if hex != [Field::Empty; 6] {
                return false;
            },
            None => return false,
        }
        // After repeatedly failing to find an efficient and elegant way to check if a hex is
        // removable, I have opted to just use a lookup table. Even if it is not elegant, at
        // least it is simple and efficient.
        let neighbor_combinations = [
            0b000001, 0b000010, 0b000100, 0b001000, 0b010000, 0b100000, 0b000011, 0b000110,
            0b001100, 0b011000, 0b110000, 0b100001, 0b000111, 0b001110, 0b011100, 0b111000,
            0b110001, 0b100011,
        ];

        let mut neighbors = 0;
        for f in 0..6 {
            if self.get_hex_neighbor(coord, f).is_some() {
                neighbors += 1 << f;
            }
        }
        assert!(
            neighbors != 0,
            "A hex at {:?} is empty and has no neighbors",
            coord
        );
        neighbor_combinations.contains(&neighbors)
    }
    fn remove_hex(&mut self, coord: &HexCoord) -> bool {
        let removable = self.is_hex_removable(coord);

        if removable {
            self.set_hex(coord, None);
        }
        removable
    }
    fn try_hex(&self, coord: (i32, i32)) -> Option<HexCoord> {
        if let Some(coord) = HexCoord::try_new(coord.0, coord.1) {
            if self.get_hex(&coord).is_some() {
                return Some(coord);
            }
        }
        None
    }
    fn check_hexes(&mut self, coord: &HexCoord) -> (u32, Vec<FieldCoord>) {
        let mut remove_count = 0;
        let mut fields = vec![];

        if self.remove_hex(coord) {
            remove_count += 1;
            for f in 0..6 {
                if let Some(neighbor) = self.get_hex_neighbor(coord, f) {
                    let (new_remove_count, mut new_fields) = self.check_hexes(&neighbor);
                    if new_remove_count == 0 {
                        fields.push(neighbor.to_field((f + 3) % 6));
                    } else {
                        remove_count += new_remove_count;
                        fields.append(&mut new_fields);
                    }
                }
            }
        }
        (remove_count, fields)
    }
}
<commit_msg>Rewrite field neighbor methods<commit_after>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use model::{Color, FieldCoord, HexCoord};

#[derive(Clone, Copy)]
pub struct Board {
    /*
                ____
               /    \
          ____/  +y  \____
         /    \      /    \
        /      \____/  +x  \
        \      /    \      /
         \____/      \____/
         /    \      /    \
        /  -x  \____/      \
        \      /    \      /
         \____/  -y  \____/
              \      /
               \____/

        The hex board uses an axial coordinate system with (0, 0) at the center. The x-axis slopes
        up to the right, and the y-axis goes up and down. The board is stored as a dense 2D array,
        because a ragged array won't work (since each row would have to be a different type).

        See http://www.redblobgames.com/grids/hexagons/#coordinates-axial for more info.
    */
    board: [[Option<Hex>; 5]; 5],
    turn: Color,
    black_pieces: u32,
    black_hexes: u32,
    white_pieces: u32,
    white_hexes: u32,
}

// Fields are numbered clockwise from the top. Even indicies are black, odd indicies are white.
type Hex = [Field; 6];

// There's no need to store the color of the piece on this field, since only white pieces can be on
// white fields, and vice versa.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Field {
    Piece,
    Empty,
}

// Public methods
impl Board {
    #[cfg_attr(rustfmt, rustfmt_skip)]
    /// Create a new board with the "Laurentius" starting position.
    pub fn new() -> Board {
        let mut board = Board {
            board: [[None; 5]; 5],
            turn: Color::White,
            black_pieces: 18,
            black_hexes: 0,
            white_pieces: 18,
            white_hexes: 0,
        };

        // (0, 0) is the only empty hex.
        board.set_hex(&HexCoord::new(0, 0), Some([Field::Empty; 6]));

        // All other hexes have exactly two pieces on them in the starting position.
        let piece_locations = [
            (-2,  2, 0, 4),
            (-2,  1, 0, 3),
            (-2,  0, 3, 5),
            (-1,  2, 1, 4),
            (-1,  1, 0, 4),
            (-1,  0, 3, 5),
            (-1, -1, 2, 5),
            ( 0,  2, 1, 5),
            ( 0,  1, 1, 5),
            ( 0, -1, 2, 4),
            ( 0, -2, 2, 4),
            ( 1,  1, 2, 5),
            ( 1,  0, 0, 2),
            ( 1, -1, 1, 3),
            ( 1, -2, 1, 4),
            ( 2,  0, 0, 2),
            ( 2, -1, 0, 3),
            ( 2, -2, 1, 3),
        ];

        for &(x, y, f1, f2) in &piece_locations {
            let mut hex = [Field::Empty; 6];
            hex[f1] = Field::Piece;
            hex[f2] = Field::Piece;
            board.set_hex(&HexCoord::new(x, y), Some(hex));
        }

        board
    }
    pub fn can_move_piece(&self, from: &FieldCoord, to: &FieldCoord) -> bool {
        from.color() == self.turn && self.is_piece_on_field(from) && !self.is_piece_on_field(to)
            && self.get_field_vertex_neighbors(from).contains(to)
    }
    pub fn move_piece(&mut self, from: &FieldCoord, to: &FieldCoord) {
        assert!(self.can_move_piece(from, to));

        self.set_field(from, Field::Empty);
        self.set_field(to, Field::Piece);

        let (capture_count, mut fields_to_check) = self.check_hexes(&from.to_hex());
        fields_to_check.append(&mut self.get_field_edge_neighbors(to));
        self.check_captures(&fields_to_check);

        match self.turn {
            Color::White => self.white_hexes += capture_count,
            Color::Black => self.black_hexes += capture_count,
        }
        self.turn = self.turn.switch();
    }
    pub fn get_available_moves(&self, field: &FieldCoord) -> Vec<FieldCoord> {
        if self.is_piece_on_field(field) {
            self.get_field_vertex_neighbors(field)
                .into_iter()
                .filter(|c| !self.is_piece_on_field(c))
                .collect()
        } else {
            vec![]
        }
    }
    pub fn can_exchange(&self) -> bool {
        2 <= match self.turn {
            Color::Black => self.black_hexes,
            Color::White => self.white_hexes,
        }
    }
    pub fn can_exchange_piece(&self, coord: &FieldCoord) -> bool {
        self.can_exchange() && coord.color() != self.turn && self.is_piece_on_field(coord)
    }
    pub fn exchange_piece(&mut self, coord: &FieldCoord) {
        assert!(self.can_exchange_piece(coord));

        self.remove_piece(coord);
        match self.turn {
            Color::White => self.white_hexes -= 2,
            Color::Black => self.black_hexes -= 2,
        }

        // Players don't collect hexes removed due to an exchange
        let (_, fields_to_check) = self.check_hexes(&coord.to_hex());
        self.check_captures(&fields_to_check);
        self.turn = self.turn.switch();
    }
    pub fn is_piece_on_field(&self, coord: &FieldCoord) -> bool {
        self.get_field(coord) == &Field::Piece
    }
    /// > extant (adj.): Still in existence; not destroyed, lost, or extinct (The Free Dictionary)
    ///
    /// Return the coordinates of the hexes that have not been removed yet.
    pub fn extant_hexes(&self) -> Vec<HexCoord> {
        let mut coords = vec![];
        for x in -2..3 {
            for y in -2..3 {
                coords.push((x, y));
            }
        }
        coords
            .iter()
            .filter_map(|&coord| self.try_hex(coord))
            .collect()
    }
    pub fn turn(&self) -> Color {
        self.turn
    }
    pub fn black_pieces(&self) -> u32 {
        self.black_pieces
    }
    pub fn black_hexes(&self) -> u32 {
        self.black_hexes
    }
    pub fn white_pieces(&self) -> u32 {
        self.white_pieces
    }
    pub fn white_hexes(&self) -> u32 {
        self.white_hexes
    }
}

// Field and piece methods
impl Board {
    fn get_field(&self, coord: &FieldCoord) -> &Field {
        match *self.get_hex(&coord.to_hex()) {
            Some(ref hex) => &hex[coord.f as usize],
            None => panic!(
                "Cannot get field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    fn set_field(&mut self, coord: &FieldCoord, field: Field) {
        match *self.get_hex(&coord.to_hex()) {
            Some(mut hex) => {
                hex[coord.f as usize] = field;
                self.set_hex(&coord.to_hex(), Some(hex));
            }
            None => panic!(
                "Cannot set field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    /// Return fields that share an edge with the given field. These fields are always the opposite
    /// color of the given field. If all of a piece's edge neighbors are occupied, that piece might
    /// be capturable.
    fn get_field_edge_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let hex = coord.to_hex();
        let mut neighbors = vec![
            // There are always two edge neighbors on the same hex as the given field
            hex.to_field((coord.f + 1) % 6),
            hex.to_field((coord.f + 5) % 6),
        ];

        if let Some(hex) = self.get_hex_neighbor(&hex, coord.f) {
            neighbors.push(hex.to_field((coord.f + 3) % 6));
        }

        neighbors
    }
    /// Return fields that share a vertex with the given field and have the same color as the given
    /// field. Pieces can move to fields that are vertex neighbors of the field they are on.
    fn get_field_vertex_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let hex = coord.to_hex();
        let mut neighbors = vec![
            // There are always two vertex neighbors on the same hex as the given field
            hex.to_field((coord.f + 2) % 6),
            hex.to_field((coord.f + 4) % 6),
        ];
        if let Some(hex) = self.get_hex_neighbor(&hex, coord.f) {
            neighbors.push(hex.to_field((coord.f + 2) % 6));
            neighbors.push(hex.to_field((coord.f + 4) % 6));
        }
        if let Some(hex) = self.get_hex_neighbor(&hex, (coord.f + 1) % 6) {
            neighbors.push(hex.to_field((coord.f + 4) % 6));
        }
        if let Some(hex) = self.get_hex_neighbor(&hex, (coord.f + 5) % 6) {
            neighbors.push(hex.to_field((coord.f + 2) % 6));
        }
        neighbors
    }
    fn remove_piece(&mut self, coord: &FieldCoord) {
        assert!(
            self.is_piece_on_field(coord),
            "There is no piece at {:?} to remove",
            coord
        );
        self.set_field(coord, Field::Empty);
        match coord.color() {
            Color::Black => self.black_pieces -= 1,
            Color::White => self.white_pieces -= 1,
        }
    }
    fn check_captures(&mut self, fields_to_check: &[FieldCoord]) {
        for field in fields_to_check {
            if field.color() != self.turn && self.is_piece_on_field(field)
                && self.get_field_edge_neighbors(field)
                    .iter()
                    .all(|coord| self.is_piece_on_field(coord))
            {
                self.remove_piece(field);
            }
        }
    }
}

// Hex methods
impl Board {
    fn get_hex(&self, coord: &HexCoord) -> &Option<Hex> {
        &self.board[(coord.x + 2) as usize][(coord.y + 2) as usize]
    }
    fn set_hex(&mut self, coord: &HexCoord, hex: Option<Hex>) {
        self.board[(coord.x + 2) as usize][(coord.y + 2) as usize] = hex;
    }
    fn get_hex_neighbor(&self, coord: &HexCoord, direction: u32) -> Option<HexCoord> {
        assert!(direction < 6);

        let neighbors = [
            (coord.x, coord.y + 1),
            (coord.x + 1, coord.y),
            (coord.x + 1, coord.y - 1),
            (coord.x, coord.y - 1),
            (coord.x - 1, coord.y),
            (coord.x - 1, coord.y + 1),
        ];

        self.try_hex(neighbors[direction as usize])
    }
    /// A hex is removable (and must be removed) if it is empty and is "attached to the board by 3
    /// or less adjacent sides."
    fn is_hex_removable(&self, coord: &HexCoord) -> bool {
        match *self.get_hex(coord) {
            Some(hex) => if hex != [Field::Empty; 6] {
                return false;
            },
            None => return false,
        }
        // After repeatedly failing to find an efficient and elegant way to check if a hex is
        // removable, I have opted to just use a lookup table. Even if it is not elegant, at
        // least it is simple and efficient.
        let neighbor_combinations = [
            0b000001, 0b000010, 0b000100, 0b001000, 0b010000, 0b100000, 0b000011, 0b000110,
            0b001100, 0b011000, 0b110000, 0b100001, 0b000111, 0b001110, 0b011100, 0b111000,
            0b110001, 0b100011,
        ];

        let mut neighbors = 0;
        for f in 0..6 {
            if self.get_hex_neighbor(coord, f).is_some() {
                neighbors += 1 << f;
            }
        }
        assert!(
            neighbors != 0,
            "A hex at {:?} is empty and has no neighbors",
            coord
        );
        neighbor_combinations.contains(&neighbors)
    }
    fn remove_hex(&mut self, coord: &HexCoord) -> bool {
        let removable = self.is_hex_removable(coord);

        if removable {
            self.set_hex(coord, None);
        }
        removable
    }
    fn try_hex(&self, coord: (i32, i32)) -> Option<HexCoord> {
        if let Some(coord) = HexCoord::try_new(coord.0, coord.1) {
            if self.get_hex(&coord).is_some() {
                return Some(coord);
            }
        }
        None
    }
    fn check_hexes(&mut self, coord: &HexCoord) -> (u32, Vec<FieldCoord>) {
        let mut remove_count = 0;
        let mut fields = vec![];

        if self.remove_hex(coord) {
            remove_count += 1;
            for f in 0..6 {
                if let Some(neighbor) = self.get_hex_neighbor(coord, f) {
                    let (new_remove_count, mut new_fields) = self.check_hexes(&neighbor);
                    if new_remove_count == 0 {
                        fields.push(neighbor.to_field((f + 3) % 6));
                    } else {
                        remove_count += new_remove_count;
                        fields.append(&mut new_fields);
                    }
                }
            }
        }
        (remove_count, fields)
    }
}
<|endoftext|>


<commit_before>use postgres;
use uuid::Uuid;
use error::Error;
use chrono::{NaiveDateTime, Utc};
use super::{conn, Model};
use model::enclosure::Enclosure;
use scraper;
use Track;
use Playlist;
use Album;
use model::PaginatedCollection;
use serde_json::Value;
use feed_rs;

static PROPS: [&'static str; 20]  = ["id",
                                     "url",
                                     "title",
                                     "description",
                                     "visual_url",
                                     "locale",

                                     "summary",
                                     "content",
                                     "author",
                                     "crawled",
                                     "published",
                                     "updated",
                                     "fingerprint",
                                     "origin_id",
                                     "alternate",
                                     "keywords",
                                     "enclosure",

                                     "feed_id",

                                     "created_at",
                                     "updated_at"];

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Entry {
    pub id:          Uuid,
    pub url:         String,
    pub title:       Option<String>,
    pub description: Option<String>,
    pub visual_url:  Option<String>,
    pub locale:      Option<String>,
    pub summary:     Option<String>,
    pub content:     Option<String>,
    pub author:      Option<String>,
    pub crawled:     NaiveDateTime,
    pub published:   NaiveDateTime,
    pub updated:     Option<NaiveDateTime>,
    pub fingerprint: String,
    pub origin_id:   String,
    pub alternate:   Value,
    pub keywords:    Value,
    pub enclosure:   Value,
    pub feed_id:     Option<Uuid>,
    pub created_at:  NaiveDateTime,
    pub updated_at:  NaiveDateTime,
    pub tracks:      Vec<Track>,
    pub playlists:   Vec<Playlist>,
    pub albums:      Vec<Album>,
}

impl<'a> Model<'a> for Entry {
    fn table_name() -> String {
        "entries".to_string()
    }
    fn props_str(prefix: &str) -> String {
        PROPS
            .iter()
            .map(|&p| format!("{}{}", prefix, p))
            .collect::<Vec<String>>().join(",")
    }
    fn rows_to_items(rows: postgres::rows::Rows) -> Vec<Entry> {
        let mut entries = Vec::new();
        for row in rows.iter() {
            entries.push(Entry {
                id:          row.get(0),
                url:         row.get(1),
                title:       row.get(2),
                description: row.get(3),
                visual_url:  row.get(4),
                locale:      row.get(5),
                summary:     row.get(6),
                content:     row.get(7),
                author:      row.get(8),
                crawled:     row.get(9),
                published:   row.get(10),
                updated:     row.get(11),
                fingerprint: row.get(12),
                origin_id:   row.get(13),
                alternate:   row.get(14),
                keywords:    row.get(15),
                enclosure:   row.get(16),
                feed_id:     row.get(17),
                created_at:  row.get(18),
                updated_at:  row.get(19),
                tracks:      Track::find_by_entry_id(row.get(0)),
                playlists:   Playlist::find_by_entry_id(row.get(0)),
                albums:      Album::find_by_entry_id(row.get(0)),
            })
        }
        entries
    }
    fn create(&self) -> Result<Entry, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO entries (url, published) VALUES ($1, $2) RETURNING id"));
        let rows = try!(stmt.query(&[&self.url, &NaiveDateTime::from_timestamp(0, 0)]));
        let mut entry = self.clone();
        for row in rows.iter() {
            entry.id = row.get(0);
        }
        Ok(entry)
    }
    fn save(&mut self) -> Result<(), Error> {
        self.updated_at = Utc::now().naive_utc();
        let conn = try!(conn());
        let stmt = try!(conn.prepare("UPDATE entries SET
                                   url         = $2,
                                   title       = $3,
                                   description = $4,
                                   visual_url  = $5,
                                   locale      = $6,
                                   summary     = $7,
                                   content     = $8,
                                   author      = $9,
                                   crawled     = $10,
                                   published   = $11,
                                   updated     = $12,
                                   fingerprint = $13,
                                   origin_id   = $14,
                                   alternate   = $15,
                                   keywords    = $16,
                                   enclosure   = $17,
                                   feed_id     = $18,
                                   created_at  = $19,
                                   updated_at  = $20
                                 WHERE id = $1"));
        let result = stmt.query(&[&self.id,
                                  &self.url,
                                  &self.title,
                                  &self.description,
                                  &self.visual_url,
                                  &self.locale,
                                  &self.summary,
                                  &self.content,
                                  &self.author,
                                  &self.crawled,
                                  &self.published,
                                  &self.updated,
                                  &self.fingerprint,
                                  &self.origin_id,
                                  &self.alternate,
                                  &self.keywords,
                                  &self.enclosure,
                                  &self.feed_id,
                                  &self.created_at,
                                  &self.updated_at]);
        try!(result);
        Ok(())
    }
}

impl Entry {
    pub fn new(url: String) -> Result<Entry, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO entries (url) VALUES ($1) RETURNING id"));
        let rows = try!(stmt.query(&[&url]));
        for row in rows.iter() {
            let entry = Entry {
                id:          row.get(0),
                url:         url,
                title:       None,
                description: None,
                visual_url:  None,
                locale:      None,

                summary:     None,
                content:     None,
                author:      None,
                crawled:     NaiveDateTime::from_timestamp(0, 0),
                published:   NaiveDateTime::from_timestamp(0, 0), // exclude from api response
                updated:     None,
                fingerprint: "".to_string(),
                origin_id:   "".to_string(),
                alternate:   Value::Null,
                keywords:    Value::Null,
                enclosure:   Value::Null,
                feed_id:     None,

                created_at:  Utc::now().naive_utc(),
                updated_at:  Utc::now().naive_utc(),

                tracks:      Vec::new(),
                playlists:   Vec::new(),
                albums:      Vec::new(),
            };
            return Ok(entry);
        }
        Err(Error::Unexpected)
    }

    pub fn find_by_url(url: &str) -> Result<Entry, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare(
            &format!("SELECT {} FROM entries
                        WHERE url = $1", Self::props_str(""))));
        let rows = try!(stmt.query(&[&url]));
        let entries = Entry::rows_to_items(rows);
        if entries.len() > 0 {
            return Ok(entries[0].clone());
        }
        return Err(Error::NotFound)
    }

    pub fn find_or_create_by_url(url: String) -> Result<Entry, Error> {
        match Entry::find_by_url(&url) {
            Ok(entry) => Ok(entry),
            Err(Error::NotFound) => Entry::create_by_url(url),
            Err(e) => Err(e),
        }
    }

    pub fn find_by_feed_id(feed_id: Uuid, newer_than: Option<NaiveDateTime>, page: i64, per_page: i64) -> PaginatedCollection<Entry> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM entries
                        WHERE entries.feed_id = $1 AND entries.published >= $2
                        ORDER BY entries.published DESC
                        LIMIT $4 OFFSET $3",
                     Entry::props_str(""))).unwrap();
        let offset = page * per_page;
        let published = newer_than.unwrap_or(NaiveDateTime::from_timestamp(1000, 0)); // ignore 0 timestamp
        let rows   = stmt.query(&[&feed_id, &published, &offset, &per_page]).unwrap();
        let items  = Self::rows_to_items(rows);
        let mut total: i64 = 0;
        let sql = "SELECT COUNT(*) FROM entries WHERE entries.feed_id = $1";
        for row in conn.query(&sql, &[&feed_id]).unwrap().iter() {
            total = row.get(0);
        }
        PaginatedCollection {
            page:     page,
            per_page: per_page,
            total:    total,
            items:    items,
        }
    }

    pub fn create_by_url(url: String) -> Result<Entry, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO entries (url, published) VALUES ($1, $2) RETURNING id"));
        let rows = try!(stmt.query(&[&url, &NaiveDateTime::from_timestamp(0, 0)]));
        for row in rows.iter() {
            let entry = Entry {
                id:          row.get(0),
                url:         url,
                title:       None,
                description: None,
                visual_url:  None,
                locale:      None,

                summary:     None,
                content:     None,
                author:      None,
                crawled:     Utc::now().naive_utc(),
                published:   Utc::now().naive_utc(),
                updated:     None,
                fingerprint: "".to_string(),
                origin_id:   "".to_string(),
                alternate:   Value::Null,
                keywords:    Value::Null,
                enclosure:   Value::Null,
                feed_id:     None,

                created_at:  Utc::now().naive_utc(),
                updated_at:  Utc::now().naive_utc(),

                tracks:      Vec::new(),
                playlists:   Vec::new(),
                albums:      Vec::new(),
            };
            return Ok(entry);
        }
        Err(Error::Unexpected)
    }

    pub fn is_valid(&self) -> bool {
        self.published.timestamp() >= 1000 && self.feed_id != None
    }

    pub fn find_or_create_by_url_if_invalid(url: String) -> Result<Entry, Error> {
        match Entry::find_by_url(&url) {
            Ok(entry) => {
                if !entry.is_valid() {
                    Ok(entry)
                } else {
                    Err(Error::NotFound)
                }
            },
            Err(_) => {
                Entry::create_by_url(url)
            }
        }
    }

    pub fn update_with_feed_entry(&mut self, entry: &feed_rs::Entry) {
        self.title       = entry.title.clone().map(|s| s.trim().to_string());
        self.summary     = entry.summary.clone().map(|s| s.trim().to_string());
        self.content     = entry.content.clone().map(|s| s.trim().to_string());
        self.author      = entry.author.clone();
        self.crawled     = Utc::now().naive_utc();
        self.published   = entry.published;
        self.updated     = entry.updated;
        self.fingerprint = entry.fingerprint.clone();
        self.origin_id   = entry.id.clone();
        self.alternate   = json!(entry.alternate);
        self.keywords    = json!(entry.keywords);
        self.enclosure   = json!(entry.enclosure);

        self.updated_at  = Utc::now().naive_utc();
    }

    pub fn has_title(&self) -> bool {
        let title = &self.title.clone().unwrap_or("".to_string());
        !title.is_empty()
    }

    pub fn playlistify(&mut self) -> Result<(), Error> {
        let product = try!(scraper::extract(&self.url));
        match product.og_obj {
            Some(og_obj) => {
                if !self.has_title() {
                    self.title = Some(og_obj.title);
                }
                self.description = og_obj.description;
                self.locale      = og_obj.locale;
                self.visual_url  = og_obj.images.first().map(|i| i.url.clone());
            },
            None => (),
        }
        for t in product.tracks {
            let new_track = try!(Track::find_or_create(t.provider, t.identifier.to_string()));
            let mut track = t.clone();
            track.id      = new_track.id;
            try!(track.fetch_detail().save());
            match self.tracks.iter().find(|&t| t.id == track.id) {
                Some(_) => (),
                None    => try!(self.add_track(track.clone())),
            }
        }
        for p in product.playlists {
            let new_playlist = try!(Playlist::find_or_create(p.provider, p.identifier.to_string()));
            let mut playlist = p.clone();
            playlist.id      = new_playlist.id;
            try!(playlist.save());
            match self.playlists.iter().find(|&p| p.id == playlist.id) {
                Some(_) => (),
                None    => try!(self.add_playlist(playlist.clone())),
            }
        }
        for a in product.albums {
            let new_album = try!(Album::find_or_create(a.provider, a.identifier.to_string()));
            let mut album = a.clone();
            album.id      = new_album.id;
            try!(album.save());
            match self.albums.iter().find(|&a| a.id == album.id) {
                Some(_) => (),
                None    => try!(self.add_album(album.clone())),
            }
        }

        Ok(())
    }

    pub fn add_track(&mut self, track: Track) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO track_entries (track_id, entry_id)
                                 VALUES ($1, $2)"));
        try!(stmt.query(&[&track.id, &self.id]));
        self.tracks.push(track);
        Ok(())
    }

    pub fn add_playlist(&mut self, playlist: Playlist) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO playlist_entries (playlist_id, entry_id)
                                 VALUES ($1, $2)"));
        try!(stmt.query(&[&playlist.id, &self.id]));
        self.playlists.push(playlist);
        Ok(())
    }

    pub fn add_album(&mut self, album: Album) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO album_entries (album_id, entry_id)
                                 VALUES ($1, $2)"));
        try!(stmt.query(&[&album.id, &self.id]));
        self.albums.push(album);
        Ok(())
    }
}
<commit_msg>If feed_rs::Entry#origin_id is empty string, then use Entry#url for Entry#origin_id<commit_after>use postgres;
use uuid::Uuid;
use error::Error;
use chrono::{NaiveDateTime, Utc};
use super::{conn, Model};
use model::enclosure::Enclosure;
use scraper;
use Track;
use Playlist;
use Album;
use model::PaginatedCollection;
use serde_json::Value;
use feed_rs;

static PROPS: [&'static str; 20]  = ["id",
                                     "url",
                                     "title",
                                     "description",
                                     "visual_url",
                                     "locale",

                                     "summary",
                                     "content",
                                     "author",
                                     "crawled",
                                     "published",
                                     "updated",
                                     "fingerprint",
                                     "origin_id",
                                     "alternate",
                                     "keywords",
                                     "enclosure",

                                     "feed_id",

                                     "created_at",
                                     "updated_at"];

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Entry {
    pub id:          Uuid,
    pub url:         String,
    pub title:       Option<String>,
    pub description: Option<String>,
    pub visual_url:  Option<String>,
    pub locale:      Option<String>,
    pub summary:     Option<String>,
    pub content:     Option<String>,
    pub author:      Option<String>,
    pub crawled:     NaiveDateTime,
    pub published:   NaiveDateTime,
    pub updated:     Option<NaiveDateTime>,
    pub fingerprint: String,
    pub origin_id:   String,
    pub alternate:   Value,
    pub keywords:    Value,
    pub enclosure:   Value,
    pub feed_id:     Option<Uuid>,
    pub created_at:  NaiveDateTime,
    pub updated_at:  NaiveDateTime,
    pub tracks:      Vec<Track>,
    pub playlists:   Vec<Playlist>,
    pub albums:      Vec<Album>,
}

impl<'a> Model<'a> for Entry {
    fn table_name() -> String {
        "entries".to_string()
    }
    fn props_str(prefix: &str) -> String {
        PROPS
            .iter()
            .map(|&p| format!("{}{}", prefix, p))
            .collect::<Vec<String>>().join(",")
    }
    fn rows_to_items(rows: postgres::rows::Rows) -> Vec<Entry> {
        let mut entries = Vec::new();
        for row in rows.iter() {
            entries.push(Entry {
                id:          row.get(0),
                url:         row.get(1),
                title:       row.get(2),
                description: row.get(3),
                visual_url:  row.get(4),
                locale:      row.get(5),
                summary:     row.get(6),
                content:     row.get(7),
                author:      row.get(8),
                crawled:     row.get(9),
                published:   row.get(10),
                updated:     row.get(11),
                fingerprint: row.get(12),
                origin_id:   row.get(13),
                alternate:   row.get(14),
                keywords:    row.get(15),
                enclosure:   row.get(16),
                feed_id:     row.get(17),
                created_at:  row.get(18),
                updated_at:  row.get(19),
                tracks:      Track::find_by_entry_id(row.get(0)),
                playlists:   Playlist::find_by_entry_id(row.get(0)),
                albums:      Album::find_by_entry_id(row.get(0)),
            })
        }
        entries
    }
    fn create(&self) -> Result<Entry, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO entries (url, published) VALUES ($1, $2) RETURNING id"));
        let rows = try!(stmt.query(&[&self.url, &NaiveDateTime::from_timestamp(0, 0)]));
        let mut entry = self.clone();
        for row in rows.iter() {
            entry.id = row.get(0);
        }
        Ok(entry)
    }
    fn save(&mut self) -> Result<(), Error> {
        self.updated_at = Utc::now().naive_utc();
        let conn = try!(conn());
        let stmt = try!(conn.prepare("UPDATE entries SET
                                   url         = $2,
                                   title       = $3,
                                   description = $4,
                                   visual_url  = $5,
                                   locale      = $6,
                                   summary     = $7,
                                   content     = $8,
                                   author      = $9,
                                   crawled     = $10,
                                   published   = $11,
                                   updated     = $12,
                                   fingerprint = $13,
                                   origin_id   = $14,
                                   alternate   = $15,
                                   keywords    = $16,
                                   enclosure   = $17,
                                   feed_id     = $18,
                                   created_at  = $19,
                                   updated_at  = $20
                                 WHERE id = $1"));
        let result = stmt.query(&[&self.id,
                                  &self.url,
                                  &self.title,
                                  &self.description,
                                  &self.visual_url,
                                  &self.locale,
                                  &self.summary,
                                  &self.content,
                                  &self.author,
                                  &self.crawled,
                                  &self.published,
                                  &self.updated,
                                  &self.fingerprint,
                                  &self.origin_id,
                                  &self.alternate,
                                  &self.keywords,
                                  &self.enclosure,
                                  &self.feed_id,
                                  &self.created_at,
                                  &self.updated_at]);
        try!(result);
        Ok(())
    }
}

impl Entry {
    pub fn new(url: String) -> Result<Entry, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO entries (url) VALUES ($1) RETURNING id"));
        let rows = try!(stmt.query(&[&url]));
        for row in rows.iter() {
            let entry = Entry {
                id:          row.get(0),
                url:         url,
                title:       None,
                description: None,
                visual_url:  None,
                locale:      None,

                summary:     None,
                content:     None,
                author:      None,
                crawled:     NaiveDateTime::from_timestamp(0, 0),
                published:   NaiveDateTime::from_timestamp(0, 0), // exclude from api response
                updated:     None,
                fingerprint: "".to_string(),
                origin_id:   "".to_string(),
                alternate:   Value::Null,
                keywords:    Value::Null,
                enclosure:   Value::Null,
                feed_id:     None,

                created_at:  Utc::now().naive_utc(),
                updated_at:  Utc::now().naive_utc(),

                tracks:      Vec::new(),
                playlists:   Vec::new(),
                albums:      Vec::new(),
            };
            return Ok(entry);
        }
        Err(Error::Unexpected)
    }

    pub fn find_by_url(url: &str) -> Result<Entry, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare(
            &format!("SELECT {} FROM entries
                        WHERE url = $1", Self::props_str(""))));
        let rows = try!(stmt.query(&[&url]));
        let entries = Entry::rows_to_items(rows);
        if entries.len() > 0 {
            return Ok(entries[0].clone());
        }
        return Err(Error::NotFound)
    }

    pub fn find_or_create_by_url(url: String) -> Result<Entry, Error> {
        match Entry::find_by_url(&url) {
            Ok(entry) => Ok(entry),
            Err(Error::NotFound) => Entry::create_by_url(url),
            Err(e) => Err(e),
        }
    }

    pub fn find_by_feed_id(feed_id: Uuid, newer_than: Option<NaiveDateTime>, page: i64, per_page: i64) -> PaginatedCollection<Entry> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM entries
                        WHERE entries.feed_id = $1 AND entries.published >= $2
                        ORDER BY entries.published DESC
                        LIMIT $4 OFFSET $3",
                     Entry::props_str(""))).unwrap();
        let offset = page * per_page;
        let published = newer_than.unwrap_or(NaiveDateTime::from_timestamp(1000, 0)); // ignore 0 timestamp
        let rows   = stmt.query(&[&feed_id, &published, &offset, &per_page]).unwrap();
        let items  = Self::rows_to_items(rows);
        let mut total: i64 = 0;
        let sql = "SELECT COUNT(*) FROM entries WHERE entries.feed_id = $1";
        for row in conn.query(&sql, &[&feed_id]).unwrap().iter() {
            total = row.get(0);
        }
        PaginatedCollection {
            page:     page,
            per_page: per_page,
            total:    total,
            items:    items,
        }
    }

    pub fn create_by_url(url: String) -> Result<Entry, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO entries (url, published) VALUES ($1, $2) RETURNING id"));
        let rows = try!(stmt.query(&[&url, &NaiveDateTime::from_timestamp(0, 0)]));
        for row in rows.iter() {
            let entry = Entry {
                id:          row.get(0),
                url:         url,
                title:       None,
                description: None,
                visual_url:  None,
                locale:      None,

                summary:     None,
                content:     None,
                author:      None,
                crawled:     Utc::now().naive_utc(),
                published:   Utc::now().naive_utc(),
                updated:     None,
                fingerprint: "".to_string(),
                origin_id:   "".to_string(),
                alternate:   Value::Null,
                keywords:    Value::Null,
                enclosure:   Value::Null,
                feed_id:     None,

                created_at:  Utc::now().naive_utc(),
                updated_at:  Utc::now().naive_utc(),

                tracks:      Vec::new(),
                playlists:   Vec::new(),
                albums:      Vec::new(),
            };
            return Ok(entry);
        }
        Err(Error::Unexpected)
    }

    pub fn is_valid(&self) -> bool {
        self.published.timestamp() >= 1000 && self.feed_id != None
    }

    pub fn find_or_create_by_url_if_invalid(url: String) -> Result<Entry, Error> {
        match Entry::find_by_url(&url) {
            Ok(entry) => {
                if !entry.is_valid() {
                    Ok(entry)
                } else {
                    Err(Error::NotFound)
                }
            },
            Err(_) => {
                Entry::create_by_url(url)
            }
        }
    }

    pub fn update_with_feed_entry(&mut self, entry: &feed_rs::Entry) {
        self.title       = entry.title.clone().map(|s| s.trim().to_string());
        self.summary     = entry.summary.clone().map(|s| s.trim().to_string());
        self.content     = entry.content.clone().map(|s| s.trim().to_string());
        self.author      = entry.author.clone();
        self.crawled     = Utc::now().naive_utc();
        self.published   = entry.published;
        self.updated     = entry.updated;
        self.fingerprint = entry.fingerprint.clone();
        self.alternate   = json!(entry.alternate);
        self.keywords    = json!(entry.keywords);
        self.enclosure   = json!(entry.enclosure);

        self.updated_at  = Utc::now().naive_utc();
        if !entry.id.is_empty() {
            self.origin_id = entry.id.clone()
        } else {
            self.origin_id = self.url.clone()
        }
    }

    pub fn has_title(&self) -> bool {
        let title = &self.title.clone().unwrap_or("".to_string());
        !title.is_empty()
    }

    pub fn playlistify(&mut self) -> Result<(), Error> {
        let product = try!(scraper::extract(&self.url));
        match product.og_obj {
            Some(og_obj) => {
                if !self.has_title() {
                    self.title = Some(og_obj.title);
                }
                self.description = og_obj.description;
                self.locale      = og_obj.locale;
                self.visual_url  = og_obj.images.first().map(|i| i.url.clone());
            },
            None => (),
        }
        for t in product.tracks {
            let new_track = try!(Track::find_or_create(t.provider, t.identifier.to_string()));
            let mut track = t.clone();
            track.id      = new_track.id;
            try!(track.fetch_detail().save());
            match self.tracks.iter().find(|&t| t.id == track.id) {
                Some(_) => (),
                None    => try!(self.add_track(track.clone())),
            }
        }
        for p in product.playlists {
            let new_playlist = try!(Playlist::find_or_create(p.provider, p.identifier.to_string()));
            let mut playlist = p.clone();
            playlist.id      = new_playlist.id;
            try!(playlist.save());
            match self.playlists.iter().find(|&p| p.id == playlist.id) {
                Some(_) => (),
                None    => try!(self.add_playlist(playlist.clone())),
            }
        }
        for a in product.albums {
            let new_album = try!(Album::find_or_create(a.provider, a.identifier.to_string()));
            let mut album = a.clone();
            album.id      = new_album.id;
            try!(album.save());
            match self.albums.iter().find(|&a| a.id == album.id) {
                Some(_) => (),
                None    => try!(self.add_album(album.clone())),
            }
        }

        Ok(())
    }

    pub fn add_track(&mut self, track: Track) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO track_entries (track_id, entry_id)
                                 VALUES ($1, $2)"));
        try!(stmt.query(&[&track.id, &self.id]));
        self.tracks.push(track);
        Ok(())
    }

    pub fn add_playlist(&mut self, playlist: Playlist) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO playlist_entries (playlist_id, entry_id)
                                 VALUES ($1, $2)"));
        try!(stmt.query(&[&playlist.id, &self.id]));
        self.playlists.push(playlist);
        Ok(())
    }

    pub fn add_album(&mut self, album: Album) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO album_entries (album_id, entry_id)
                                 VALUES ($1, $2)"));
        try!(stmt.query(&[&album.id, &self.id]));
        self.albums.push(album);
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Matrix rooms.

use std::collections::{HashMap, HashSet};
use std::convert::{TryFrom, TryInto};

use diesel::{Connection, ExecuteDsl, ExpressionMethods, FilterDsl, FindDsl, LoadDsl, OrderDsl, insert};
use diesel::pg::PgConnection;
use diesel::pg::data_types::PgTimestamp;
use diesel::pg::expression::dsl::any;
use diesel::result::Error as DieselError;
use ruma_events::EventType;
use ruma_events::room::create::{CreateEvent, CreateEventContent};
use ruma_events::room::history_visibility::{
    HistoryVisibility,
    HistoryVisibilityEvent,
    HistoryVisibilityEventContent,
};
use ruma_events::room::join_rules::{
    JoinRule,
    JoinRulesEvent,
    JoinRulesEventContent,
};
use ruma_events::room::name::{NameEvent, NameEventContent};
use ruma_events::room::power_levels::{PowerLevelsEvent, PowerLevelsEventContent};
use ruma_events::room::topic::{TopicEvent, TopicEventContent};
use ruma_identifiers::{EventId, RoomAliasId, RoomId, UserId};

use error::ApiError;
use models::event::{Event, NewEvent};
use models::room_alias::{NewRoomAlias, RoomAlias};
use models::room_membership::{RoomMembership, RoomMembershipOptions};
use models::user::User;
use schema::{events, rooms, users};

/// Options provided by the user to customize the room upon creation.
pub struct CreationOptions {
    /// An initial alias for the room.
    pub alias: Option<String>,
    /// Whehter or not the room should be federated.
    pub federate: bool,
    /// A list of users to invite to the room.
    pub invite_list: Option<Vec<String>>,
    /// An initial name for the room.
    pub name: Option<String>,
    /// A convenience parameter for setting a few default state events.
    pub preset: RoomPreset,
    /// An initial topic for the room.
    pub topic: Option<String>,
}

/// A new Matrix room, not yet saved.
#[derive(Debug, Insertable)]
#[table_name = "rooms"]
pub struct NewRoom {
    /// The room's unique ID.
    pub id: RoomId,
    /// The ID of the user creating the room.
    pub user_id: UserId,
    /// Whether or not the room is visible in the directory.
    pub public: bool,
}

/// A Matrix room.
#[derive(Debug, Queryable)]
pub struct Room {
    /// The room's unique ID.
    pub id: RoomId,
    /// The ID of the user who created the room.
    pub user_id: UserId,
    /// Whether or not the room is visible in the directory.
    pub public: bool,
    /// The time the room was created.
    pub created_at: PgTimestamp,
}

/// A convenience parameter for setting a few default state events.
#[derive(Clone, Copy, Debug, Deserialize)]
pub enum RoomPreset {
    /// `join_rules` is set to `invite` and `history_visibility` is set to `shared`.
    PrivateChat,
    /// `join_rules` is set to `public` and `history_visibility` is set to `shared`.
    PublicChat,
    /// Same as `PrivateChat`, but all initial invitees get the same power level as the creator.
    TrustedPrivateChat,
}

impl Room {
    /// Creates a new room in the database.
    pub fn create(
        connection: &PgConnection,
        new_room: &NewRoom,
        homeserver_domain: &str,
        creation_options: &CreationOptions,
    ) -> Result<Room, ApiError> {
        connection.transaction::<Room, ApiError, _>(|| {
            let room: Room = insert(new_room)
                .into(rooms::table)
                .get_result(connection)
                .map_err(ApiError::from)?;

            if let Some(ref alias) = creation_options.alias {
                let new_room_alias = NewRoomAlias {
                    alias: RoomAliasId::try_from(&format!("#{}:{}", alias, homeserver_domain))?,
                    room_id: room.id.clone(),
                    user_id: new_room.user_id.clone(),
                    servers: vec![homeserver_domain.to_string()],
                };

                RoomAlias::create(connection, homeserver_domain, &new_room_alias)?;
            }

            let mut new_events = Vec::new();

            let new_create_event: NewEvent = CreateEvent {
                content: CreateEventContent {
                    creator: new_room.user_id.clone(),
                    federate: creation_options.federate,
                },
                event_id: EventId::new(homeserver_domain)?,
                event_type: EventType::RoomCreate,
                prev_content: None,
                room_id: room.id.clone(),
                state_key: "".to_string(),
                unsigned: None,
                user_id: new_room.user_id.clone(),
            }.try_into()?;

            new_events.push(new_create_event);

            if let Some(ref name) = creation_options.name {
                let new_name_event: NewEvent = NameEvent {
                    content: NameEventContent {
                        name: name.to_string(),
                    },
                    event_id: EventId::new(homeserver_domain)?,
                    event_type: EventType::RoomName,
                    prev_content: None,
                    room_id: room.id.clone(),
                    state_key: "".to_string(),
                    unsigned: None,
                    user_id: new_room.user_id.clone(),
                }.try_into()?;

                new_events.push(new_name_event);
            }

            if let Some(ref topic) = creation_options.topic {
                let new_topic_event: NewEvent = TopicEvent {
                    content: TopicEventContent {
                        topic: topic.to_string(),
                    },
                    event_id: EventId::new(homeserver_domain)?,
                    event_type: EventType::RoomTopic,
                    prev_content: None,
                    room_id: room.id.clone(),
                    state_key: "".to_string(),
                    unsigned: None,
                    user_id: new_room.user_id.clone(),
                }.try_into()?;

                new_events.push(new_topic_event);
            }

            let new_history_visibility_event: NewEvent = HistoryVisibilityEvent {
                content: HistoryVisibilityEventContent {
                    history_visibility: HistoryVisibility::Shared,
                },
                event_id: EventId::new(homeserver_domain)?,
                event_type: EventType::RoomHistoryVisibility,
                prev_content: None,
                room_id: room.id.clone(),
                state_key: "".to_string(),
                unsigned: None,
                user_id: new_room.user_id.clone(),
            }.try_into()?;

            new_events.push(new_history_visibility_event);

            match creation_options.preset {
                RoomPreset::PrivateChat => {
                    let new_join_rules_event: NewEvent = JoinRulesEvent {
                        content: JoinRulesEventContent { join_rule: JoinRule::Invite },
                        event_id: EventId::new(homeserver_domain)?,
                        event_type: EventType::RoomJoinRules,
                        prev_content: None,
                        room_id: room.id.clone(),
                        state_key: "".to_string(),
                        unsigned: None,
                        user_id: new_room.user_id.clone(),
                    }.try_into()?;

                    new_events.push(new_join_rules_event);
                }
                RoomPreset::PublicChat => {
                    let new_join_rules_event: NewEvent = JoinRulesEvent {
                        content: JoinRulesEventContent { join_rule: JoinRule::Public },
                        event_id: EventId::new(homeserver_domain)?,
                        event_type: EventType::RoomJoinRules,
                        prev_content: None,
                        room_id: room.id.clone(),
                        state_key: "".to_string(),
                        unsigned: None,
                        user_id: new_room.user_id.clone(),
                    }.try_into()?;

                    new_events.push(new_join_rules_event);
                }
                RoomPreset::TrustedPrivateChat => {
                    let new_join_rules_event: NewEvent = JoinRulesEvent {
                        content: JoinRulesEventContent { join_rule: JoinRule::Invite },
                        event_id: EventId::new(homeserver_domain)?,
                        event_type: EventType::RoomJoinRules,
                        prev_content: None,
                        room_id: room.id.clone(),
                        state_key: "".to_string(),
                        unsigned: None,
                        user_id: new_room.user_id.clone(),
                    }.try_into()?;

                    new_events.push(new_join_rules_event);
                }
            }

            insert(&new_events)
                .into(events::table)
                .execute(connection)
                .map_err(ApiError::from)?;

            if let Some(ref invite_list) = creation_options.invite_list {
                let mut user_ids = HashSet::with_capacity(invite_list.len());

                for invitee in invite_list {
                    let user_id = UserId::try_from(invitee)?;

                    if user_id.hostname().to_string() != homeserver_domain {
                        return Err(
                            ApiError::unimplemented("Federation is not yet supported.".to_string())
                        );
                    }

                    user_ids.insert(user_id);
                }

                let users: Vec<User> = users::table
                    .filter(users::id.eq(any(
                        user_ids.iter().cloned().collect::<Vec<UserId>>()))
                    )
                    .get_results(connection)
                    .map_err(ApiError::from)?;

                let loaded_user_ids: HashSet<UserId> = users
                    .iter()
                    .map(|user| user.id.clone())
                    .collect();

                let missing_user_ids: Vec<UserId> = user_ids
                    .difference(&loaded_user_ids)
                    .cloned()
                    .collect();

                if missing_user_ids.len() > 0 {
                    return Err(
                        ApiError::bad_json(format!(
                            "Unknown users in invite list: {}",
                            &missing_user_ids
                                .iter()
                                .map(|user_id| user_id.to_string())
                                .collect::<Vec<String>>()
                                .join(", ")
                        ))
                    )
                }

                for user in users {
                    let options = RoomMembershipOptions {
                        room_id: room.id.clone(),
                        user_id: user.id.clone(),
                        sender: room.user_id.clone(),
                        membership: "invite".to_string(),
                    };

                    RoomMembership::create(connection, &homeserver_domain, options)?;
                }
            }

            Ok(room)
        }).map_err(ApiError::from)
    }

    /// Looks up the most recent power levels event for the room.
    ///
    /// If the room does not have a power levels event, a default one is created according to the
    /// specification.
    pub fn current_power_levels(&self, connection: &PgConnection)
    -> Result<PowerLevelsEventContent, ApiError> {
        match events::table
            .filter(events::room_id.eq(self.id.clone()))
            .filter(events::event_type.eq(EventType::RoomPowerLevels.to_string()))
            .order(events::ordering.desc())
            .first::<Event>(connection)
        {
            Ok(event) => {
                let power_levels_event: PowerLevelsEvent = event.try_into()?;

                Ok(power_levels_event.content)
            }
            Err(error) => match error {
                DieselError::NotFound => Ok(PowerLevelsEventContent {
                    ban: 50,
                    events: HashMap::new(),
                    events_default: 0,
                    invite: 50,
                    kick: 50,
                    redact: 50,
                    state_default: 0,
                    users: HashMap::new(),
                    users_default: 0,
                }),
                _ => Err(error.into()),
            },
        }
    }

    /// Look up a `Room` given the `RoomId`.
    pub fn find(connection: &PgConnection, room_id: &RoomId)
    -> Result<Room, ApiError> {
        rooms::table
            .find(room_id)
            .get_result(connection)
            .map(Room::from)
            .map_err(|err| {
                match err {
                    DieselError::NotFound => ApiError::not_found(
                        "The room was not found on this server".to_string()
                    ),
                    _ => ApiError::from(err)
                }
            })
    }
}
<commit_msg>Move the membership creation process to its own method<commit_after>//! Matrix rooms.

use std::collections::{HashMap, HashSet};
use std::convert::{TryFrom, TryInto};

use diesel::{Connection, ExecuteDsl, ExpressionMethods, FilterDsl, FindDsl, LoadDsl, OrderDsl, insert};
use diesel::pg::PgConnection;
use diesel::pg::data_types::PgTimestamp;
use diesel::pg::expression::dsl::any;
use diesel::result::Error as DieselError;
use ruma_events::EventType;
use ruma_events::room::create::{CreateEvent, CreateEventContent};
use ruma_events::room::history_visibility::{
    HistoryVisibility,
    HistoryVisibilityEvent,
    HistoryVisibilityEventContent,
};
use ruma_events::room::join_rules::{
    JoinRule,
    JoinRulesEvent,
    JoinRulesEventContent,
};
use ruma_events::room::name::{NameEvent, NameEventContent};
use ruma_events::room::power_levels::{PowerLevelsEvent, PowerLevelsEventContent};
use ruma_events::room::topic::{TopicEvent, TopicEventContent};
use ruma_identifiers::{EventId, RoomAliasId, RoomId, UserId};

use error::ApiError;
use models::event::{Event, NewEvent};
use models::room_alias::{NewRoomAlias, RoomAlias};
use models::room_membership::{RoomMembership, RoomMembershipOptions};
use models::user::User;
use schema::{events, rooms, users};

/// Options provided by the user to customize the room upon creation.
pub struct CreationOptions {
    /// An initial alias for the room.
    pub alias: Option<String>,
    /// Whether or not the room should be federated.
    pub federate: bool,
    /// A list of users to invite to the room.
    pub invite_list: Option<Vec<String>>,
    /// An initial name for the room.
    pub name: Option<String>,
    /// A convenience parameter for setting a few default state events.
    pub preset: RoomPreset,
    /// An initial topic for the room.
    pub topic: Option<String>,
}

/// A new Matrix room, not yet saved.
#[derive(Debug, Insertable)]
#[table_name = "rooms"]
pub struct NewRoom {
    /// The room's unique ID.
    pub id: RoomId,
    /// The ID of the user creating the room.
    pub user_id: UserId,
    /// Whether or not the room is visible in the directory.
    pub public: bool,
}

/// A Matrix room.
#[derive(Debug, Queryable)]
pub struct Room {
    /// The room's unique ID.
    pub id: RoomId,
    /// The ID of the user who created the room.
    pub user_id: UserId,
    /// Whether or not the room is visible in the directory.
    pub public: bool,
    /// The time the room was created.
    pub created_at: PgTimestamp,
}

/// A convenience parameter for setting a few default state events.
#[derive(Clone, Copy, Debug, Deserialize)]
pub enum RoomPreset {
    /// `join_rules` is set to `invite` and `history_visibility` is set to `shared`.
    PrivateChat,
    /// `join_rules` is set to `public` and `history_visibility` is set to `shared`.
    PublicChat,
    /// Same as `PrivateChat`, but all initial invitees get the same power level as the creator.
    TrustedPrivateChat,
}

impl Room {
    /// Creates a new room in the database.
    pub fn create(
        connection: &PgConnection,
        new_room: &NewRoom,
        homeserver_domain: &str,
        creation_options: &CreationOptions,
    ) -> Result<Room, ApiError> {
        connection.transaction::<Room, ApiError, _>(|| {
            let room: Room = insert(new_room)
                .into(rooms::table)
                .get_result(connection)
                .map_err(ApiError::from)?;

            if let Some(ref alias) = creation_options.alias {
                let new_room_alias = NewRoomAlias {
                    alias: RoomAliasId::try_from(&format!("#{}:{}", alias, homeserver_domain))?,
                    room_id: room.id.clone(),
                    user_id: new_room.user_id.clone(),
                    servers: vec![homeserver_domain.to_string()],
                };

                RoomAlias::create(connection, homeserver_domain, &new_room_alias)?;
            }

            let mut new_events = Vec::new();

            let new_create_event: NewEvent = CreateEvent {
                content: CreateEventContent {
                    creator: new_room.user_id.clone(),
                    federate: creation_options.federate,
                },
                event_id: EventId::new(homeserver_domain)?,
                event_type: EventType::RoomCreate,
                prev_content: None,
                room_id: room.id.clone(),
                state_key: "".to_string(),
                unsigned: None,
                user_id: new_room.user_id.clone(),
            }.try_into()?;

            new_events.push(new_create_event);

            if let Some(ref name) = creation_options.name {
                let new_name_event: NewEvent = NameEvent {
                    content: NameEventContent {
                        name: name.to_string(),
                    },
                    event_id: EventId::new(homeserver_domain)?,
                    event_type: EventType::RoomName,
                    prev_content: None,
                    room_id: room.id.clone(),
                    state_key: "".to_string(),
                    unsigned: None,
                    user_id: new_room.user_id.clone(),
                }.try_into()?;

                new_events.push(new_name_event);
            }

            if let Some(ref topic) = creation_options.topic {
                let new_topic_event: NewEvent = TopicEvent {
                    content: TopicEventContent {
                        topic: topic.to_string(),
                    },
                    event_id: EventId::new(homeserver_domain)?,
                    event_type: EventType::RoomTopic,
                    prev_content: None,
                    room_id: room.id.clone(),
                    state_key: "".to_string(),
                    unsigned: None,
                    user_id: new_room.user_id.clone(),
                }.try_into()?;

                new_events.push(new_topic_event);
            }

            let new_history_visibility_event: NewEvent = HistoryVisibilityEvent {
                content: HistoryVisibilityEventContent {
                    history_visibility: HistoryVisibility::Shared,
                },
                event_id: EventId::new(homeserver_domain)?,
                event_type: EventType::RoomHistoryVisibility,
                prev_content: None,
                room_id: room.id.clone(),
                state_key: "".to_string(),
                unsigned: None,
                user_id: new_room.user_id.clone(),
            }.try_into()?;

            new_events.push(new_history_visibility_event);

            match creation_options.preset {
                RoomPreset::PrivateChat => {
                    let new_join_rules_event: NewEvent = JoinRulesEvent {
                        content: JoinRulesEventContent { join_rule: JoinRule::Invite },
                        event_id: EventId::new(homeserver_domain)?,
                        event_type: EventType::RoomJoinRules,
                        prev_content: None,
                        room_id: room.id.clone(),
                        state_key: "".to_string(),
                        unsigned: None,
                        user_id: new_room.user_id.clone(),
                    }.try_into()?;

                    new_events.push(new_join_rules_event);
                }
                RoomPreset::PublicChat => {
                    let new_join_rules_event: NewEvent = JoinRulesEvent {
                        content: JoinRulesEventContent { join_rule: JoinRule::Public },
                        event_id: EventId::new(homeserver_domain)?,
                        event_type: EventType::RoomJoinRules,
                        prev_content: None,
                        room_id: room.id.clone(),
                        state_key: "".to_string(),
                        unsigned: None,
                        user_id: new_room.user_id.clone(),
                    }.try_into()?;

                    new_events.push(new_join_rules_event);
                }
                RoomPreset::TrustedPrivateChat => {
                    let new_join_rules_event: NewEvent = JoinRulesEvent {
                        content: JoinRulesEventContent { join_rule: JoinRule::Invite },
                        event_id: EventId::new(homeserver_domain)?,
                        event_type: EventType::RoomJoinRules,
                        prev_content: None,
                        room_id: room.id.clone(),
                        state_key: "".to_string(),
                        unsigned: None,
                        user_id: new_room.user_id.clone(),
                    }.try_into()?;

                    new_events.push(new_join_rules_event);
                }
            }

            insert(&new_events)
                .into(events::table)
                .execute(connection)
                .map_err(ApiError::from)?;

            if let Some(ref invite_list) = creation_options.invite_list {
                room.create_memberships(connection, &invite_list, homeserver_domain)?;
            }

            Ok(room)
        }).map_err(ApiError::from)
    }

    /// Looks up the most recent power levels event for the room.
    ///
    /// If the room does not have a power levels event, a default one is created according to the
    /// specification.
    pub fn current_power_levels(&self, connection: &PgConnection)
    -> Result<PowerLevelsEventContent, ApiError> {
        match events::table
            .filter(events::room_id.eq(self.id.clone()))
            .filter(events::event_type.eq(EventType::RoomPowerLevels.to_string()))
            .order(events::ordering.desc())
            .first::<Event>(connection)
        {
            Ok(event) => {
                let power_levels_event: PowerLevelsEvent = event.try_into()?;

                Ok(power_levels_event.content)
            }
            Err(error) => match error {
                DieselError::NotFound => Ok(PowerLevelsEventContent {
                    ban: 50,
                    events: HashMap::new(),
                    events_default: 0,
                    invite: 50,
                    kick: 50,
                    redact: 50,
                    state_default: 0,
                    users: HashMap::new(),
                    users_default: 0,
                }),
                _ => Err(error.into()),
            },
        }
    }

    /// Look up a `Room` given the `RoomId`.
    pub fn find(connection: &PgConnection, room_id: &RoomId)
    -> Result<Room, ApiError> {
        rooms::table
            .find(room_id)
            .get_result(connection)
            .map(Room::from)
            .map_err(|err| {
                match err {
                    DieselError::NotFound => ApiError::not_found(
                        "The room was not found on this server".to_string()
                    ),
                    _ => ApiError::from(err)
                }
            })
    }

    /// Given a list of invited users create the appropriate `m.room.member` events.
    fn create_memberships(&self, connection: &PgConnection, invite_list: &Vec<String>, homeserver_domain: &str)
    -> Result<(), ApiError> {
        let mut user_ids = HashSet::with_capacity(invite_list.len());

        for invitee in invite_list {
            let user_id = UserId::try_from(invitee)?;

            if user_id.hostname().to_string() != homeserver_domain {
                return Err(
                    ApiError::unimplemented("Federation is not yet supported.".to_string())
                );
            }

            user_ids.insert(user_id);
        }

        let users: Vec<User> = users::table
            .filter(users::id.eq(any(
                user_ids.iter().cloned().collect::<Vec<UserId>>()))
            )
            .get_results(connection)
            .map_err(ApiError::from)?;

        let loaded_user_ids: HashSet<UserId> = users
            .iter()
            .map(|user| user.id.clone())
            .collect();

        let missing_user_ids: Vec<UserId> = user_ids
            .difference(&loaded_user_ids)
            .cloned()
            .collect();

        if missing_user_ids.len() > 0 {
            return Err(
                ApiError::bad_json(format!(
                    "Unknown users in invite list: {}",
                    &missing_user_ids
                        .iter()
                        .map(|user_id| user_id.to_string())
                        .collect::<Vec<String>>()
                        .join(", ")
                ))
            )
        }

        for user in users {
            let options = RoomMembershipOptions {
                room_id: self.id.clone(),
                user_id: user.id.clone(),
                sender: self.user_id.clone(),
                membership: "invite".to_string(),
            };

            RoomMembership::create(connection, &homeserver_domain, options)?;
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>// xfail-stage0
// xfail-stage1
// xfail-stage2
// -*- rust -*-

tag mlist {
  cons(int,mutable @mlist);
  nil;
}

fn main() {
  cons(10, @cons(11, @cons(12, @nil)));
}
<commit_msg>Un-xfail test/run-pass/mlist<commit_after>// -*- rust -*-

tag mlist {
  cons(int,mutable @mlist);
  nil;
}

fn main() {
  cons(10, @cons(11, @cons(12, @nil)));
}
<|endoftext|>