use crate::city::House; use rusqlite::{Connection, NO_PARAMS, params, Result}; use std::collections::HashMap; pub struct SqliteLayoutDB { connection: Connection, memory_db: MemoryLayoutDB, bounds_accessible: bool } pub struct MemoryLayoutDB { layouts: Vec, 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, } impl SavedLayout { pub fn new(id: usize, houses: Vec) -> Self { SavedLayout { id, houses } } pub fn houses(&self) -> &Vec { &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, 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, layout_id: usize, fully_optimized: bool) { self.layouts.push(SavedLayout { id: layout_id, houses: houses.clone() }); } } pub trait LayoutDB { fn layouts(&self) -> &Vec; fn add_layout(&mut self, houses: &Vec, fully_optimized: bool); fn get_merge_lower_bound(&self, left_layout: &SavedLayout, right_layout: &SavedLayout, y_axis: bool) -> Option; fn add_merge_lower_bound(&mut self, lower_bound: MergeLowerBound); fn add_merge_lower_bounds(&mut self, lower_bounds: Vec); } impl LayoutDB for MemoryLayoutDB { fn layouts(&self) -> &Vec { &self.layouts } fn add_layout(&mut self, houses: &Vec, 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 { 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) { 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 { self.layouts() } fn add_layout(&mut self, houses: &Vec, 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 { 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) { self.add_merge_lower_bounds(lower_bounds).unwrap() } } impl SqliteLayoutDB { pub fn from_file(filename: &str, read_bounds: bool) -> Result { let connection = Connection::open(filename)?; let mut layouts: HashMap> = 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(); if read_bounds { 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, }); } } else { } 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, bounds_accessible: read_bounds}) } pub fn memory_db(&self) -> &MemoryLayoutDB { &self.memory_db } pub fn layouts(&self) -> &Vec { self.memory_db.layouts() } pub fn add_layout(&mut self, houses: &Vec, 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 { if !self.bounds_accessible { panic!("Bounds are not loaded!") } 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<()> { if !self.bounds_accessible { panic!("Bounds are not loaded!") } 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) -> Result<()> { if !self.bounds_accessible { panic!("Bounds are not loaded!") } 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(()) } }