use crate::city::House; use rusqlite::{Connection, NO_PARAMS, params, Result}; use std::collections::HashMap; pub struct LayoutDB { connection: Connection, 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 LayoutDB { pub fn from_file(filename: &str) -> 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(); { 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(); Ok(LayoutDB { connection, layouts, merge_lower_bounds: merges }) } pub fn layouts(&self) -> &Vec { &self.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.layouts.push(SavedLayout { id: layout_id as usize, houses: houses.clone() }); Ok(()) } pub 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 } 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.merge_lower_bounds.insert((lower_bound.left_layout_id, lower_bound.right_layout_id, lower_bound.y_axis), lower_bound); Ok(()) } pub fn add_merge_lower_bounds(&mut self, lower_bounds: Vec) -> 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) 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 ], )?; self.merge_lower_bounds.insert((lower_bound.left_layout_id, lower_bound.right_layout_id, lower_bound.y_axis), lower_bound); } transaction.commit()?; Ok(()) } }