33-3-4/src/db.rs

use crate::city::House;
use rusqlite::{Connection, NO_PARAMS, params, Result};
use std::collections::HashMap;

pub struct SqliteLayoutDB {
    connection: Connection,
    memory_db: MemoryLayoutDB,
}

pub struct MemoryLayoutDB {
    layouts: Vec<SavedLayout>,
    merge_lower_bounds: HashMap<(usize, usize, bool), MergeLowerBound>,
}

pub struct MergeLowerBound {
    left_layout_id: usize,
    right_layout_id: usize,
    y_axis: bool,
    price: u32,
}

impl MergeLowerBound {
    pub fn new(left_layout_id: usize, right_layout_id: usize, y_axis: bool, price: u32) -> Self {
        MergeLowerBound { left_layout_id, right_layout_id, y_axis, price }
    }

    pub fn left_layout_id(&self) -> usize {
        self.left_layout_id
    }

    pub fn right_layout_id(&self) -> usize {
        self.right_layout_id
    }

    pub fn y_axis(&self) -> bool {
        self.y_axis
    }

    pub fn price(&self) -> u32 {
        self.price
    }
}

#[derive(Clone)]
pub struct SavedLayout {
    id: usize,
    houses: Vec<House>,
}

impl SavedLayout {
    pub fn new(id: usize, houses: Vec<House>) -> Self {
        SavedLayout { id, houses }
    }

    pub fn houses(&self) -> &Vec<House> {
        &self.houses
    }

    pub fn id(&self) -> usize {
        self.id
    }
}

impl MemoryLayoutDB {
    pub fn new() -> Self {
        MemoryLayoutDB { layouts: Vec::new(), merge_lower_bounds: HashMap::new() }
    }

    pub fn from_collections(layouts: Vec<SavedLayout>, merge_lower_bounds: HashMap<(usize, usize, bool), MergeLowerBound>) -> Self {
        MemoryLayoutDB { layouts, merge_lower_bounds }
    }

    // This exists for use by the SqliteLayoutDB
    fn add_layout_with_id(&mut self, houses: &Vec<House>, layout_id: usize, fully_optimized: bool) {
        self.layouts.push(SavedLayout { id: layout_id, houses: houses.clone() });
    }
}

pub trait LayoutDB {
    fn layouts(&self) -> &Vec<SavedLayout>;
    fn add_layout(&mut self, houses: &Vec<House>, fully_optimized: bool);
    fn get_merge_lower_bound(&self, left_layout: &SavedLayout, right_layout: &SavedLayout, y_axis: bool) -> Option<u32>;
    fn add_merge_lower_bound(&mut self, lower_bound: MergeLowerBound);
    fn add_merge_lower_bounds(&mut self, lower_bounds: Vec<MergeLowerBound>);
}

impl LayoutDB for MemoryLayoutDB {
    fn layouts(&self) -> &Vec<SavedLayout> {
        &self.layouts
    }

    fn add_layout(&mut self, houses: &Vec<House>, fully_optimized: bool) {
        let layout_id = self.layouts.iter().map(|x| x.id).max().unwrap_or_default() + 1;
        self.layouts.push(SavedLayout { id: layout_id, houses: houses.clone() });
    }

    fn get_merge_lower_bound(&self, left_layout: &SavedLayout, right_layout: &SavedLayout, y_axis: bool) -> Option<u32> {
        if let Some(bound) = self.merge_lower_bounds.get(&(left_layout.id, right_layout.id, y_axis)) {
            return Some(bound.price);
        }
        None
    }

    fn add_merge_lower_bound(&mut self, lower_bound: MergeLowerBound) {
        self.merge_lower_bounds.insert((lower_bound.left_layout_id, lower_bound.right_layout_id, lower_bound.y_axis), lower_bound);
    }

    fn add_merge_lower_bounds(&mut self, lower_bounds: Vec<MergeLowerBound>) {
        for lower_bound in lower_bounds {
            self.merge_lower_bounds.insert((lower_bound.left_layout_id, lower_bound.right_layout_id, lower_bound.y_axis), lower_bound);
        }
    }
}

impl LayoutDB for SqliteLayoutDB {
    fn layouts(&self) -> &Vec<SavedLayout> {
        self.layouts()
    }

    fn add_layout(&mut self, houses: &Vec<House>, fully_optimized: bool) {
        self.add_layout(houses, fully_optimized).unwrap();
    }

    fn get_merge_lower_bound(&self, left_layout: &SavedLayout, right_layout: &SavedLayout, y_axis: bool) -> Option<u32> {
        self.get_merge_lower_bound(left_layout, right_layout, y_axis)
    }

    fn add_merge_lower_bound(&mut self, lower_bound: MergeLowerBound) {
        self.add_merge_lower_bound(lower_bound).unwrap()
    }

    fn add_merge_lower_bounds(&mut self, lower_bounds: Vec<MergeLowerBound>) {
        self.add_merge_lower_bounds(lower_bounds).unwrap()
    }
}

impl SqliteLayoutDB {
    pub fn from_file(filename: &str) -> Result<Self> {
        let connection = Connection::open(filename)?;
        let mut layouts: HashMap<u32, Vec<(u32, u32)>> = HashMap::new();

        // We need to dispose the statement so we can move the connection later
        {
            let mut stmt = connection.prepare("SELECT layout_id, x, y FROM houses")?;
            let mut rows = stmt.query(NO_PARAMS)?;

            while let Some(row) = rows.next()? {
                let id: u32 = row.get(0)?;
                let x: u32 = row.get(1)?;
                let y: u32 = row.get(2)?;
                layouts.entry(id).or_default().push((x, y));
            }
        }

        let mut merges = HashMap::new();
        {
            let mut stmt = connection.prepare("SELECT left_layout_id, right_layout_id, axis, price FROM merge_lower_bounds")?;
            let mut rows = stmt.query(NO_PARAMS)?;

            while let Some(row) = rows.next()? {
                let left_id: u32 = row.get(0)?;
                let right_id: u32 = row.get(1)?;
                let axis: u32 = row.get(2)?;
                let price: u32 = row.get(3)?;
                merges.insert((left_id as usize, right_id as usize, axis == 1), MergeLowerBound {
                    left_layout_id: left_id as usize,
                    right_layout_id: right_id as usize,
                    y_axis: axis == 1,
                    price,
                });
            }
        }

        let layouts = layouts.into_iter().map(|(id, xy_pairs)|
            SavedLayout
            {
                id: id as usize,
                houses: xy_pairs.into_iter().map(|(x, y)| House { x: x as usize, y: y as usize }).collect(),
            }
        ).collect();

        let memory_db = MemoryLayoutDB::from_collections(layouts, merges);
        Ok(SqliteLayoutDB { connection, memory_db})
    }

    pub fn memory_db(&self) -> &MemoryLayoutDB {
        &self.memory_db
    }

    pub fn layouts(&self) -> &Vec<SavedLayout> {
        self.memory_db.layouts()
    }

    pub fn add_layout(&mut self, houses: &Vec<House>, fully_optimized: bool) -> Result<()> {
        let transaction = self.connection.transaction()?;
        transaction.execute("INSERT INTO layouts (is_fully_optimized) VALUES (?1)",
                            params![fully_optimized])?;
        let layout_id = transaction.last_insert_rowid();
        for house in houses {
            transaction.execute("INSERT INTO houses (layout_id, x, y) VALUES (?1, ?2, ?3)",
                                params![layout_id, house.x as u32, house.y as u32])?;
        }
        transaction.commit()?;
        self.memory_db.add_layout_with_id(houses, layout_id as usize, fully_optimized);
        Ok(())
    }

    pub fn get_merge_lower_bound(&self, left_layout: &SavedLayout, right_layout: &SavedLayout, y_axis: bool) -> Option<u32> {
        self.memory_db.get_merge_lower_bound(left_layout, right_layout, y_axis)
    }

    pub fn add_merge_lower_bound(&mut self, lower_bound: MergeLowerBound) -> Result<()> {
        let transaction = self.connection.transaction()?;
        transaction.execute("INSERT INTO merge_lower_bounds (left_layout_id, right_layout_id, axis, price) VALUES (?1, ?2, ?3, ?4) ON CONFLICT(left_layout_id, right_layout_id) DO UPDATE SET price = ?4",
                            params![lower_bound.left_layout_id as u32,
                                lower_bound.right_layout_id as u32,
                                if lower_bound.y_axis {1u32} else {0u32},
                                lower_bound.price
                            ],
        )?;
        transaction.commit()?;
        self.memory_db.add_merge_lower_bound(lower_bound);
        Ok(())
    }

    pub fn add_merge_lower_bounds(&mut self, lower_bounds: Vec<MergeLowerBound>) -> Result<()> {
        let transaction = self.connection.transaction()?;
        for lower_bound in &lower_bounds {
            transaction.execute("INSERT INTO merge_lower_bounds (left_layout_id, right_layout_id, axis, price) VALUES (?1, ?2, ?3, ?4) ON CONFLICT(left_layout_id, right_layout_id, axis) DO UPDATE SET price = ?4",
                                params![lower_bound.left_layout_id as u32,
                                lower_bound.right_layout_id as u32,
                                if lower_bound.y_axis {1u32} else {0u32},
                                lower_bound.price
                            ],
            )?;
        }
        transaction.commit()?;
        self.memory_db.add_merge_lower_bounds(lower_bounds);
        Ok(())
    }
}