From b1b715cc4c11f91ab13b5f2f6c78e6c2e1d2a346 Mon Sep 17 00:00:00 2001 From: "cloverzero@gmail.com" Date: Thu, 10 Jul 2025 17:19:35 +0800 Subject: [PATCH] refactor: reduce duplicated code --- pbf-craft/src/models/mod.rs | 15 +- pbf-craft/src/readers/indexed_reader.rs | 211 ++++++++++++++---------- 2 files changed, 137 insertions(+), 89 deletions(-) diff --git a/pbf-craft/src/models/mod.rs b/pbf-craft/src/models/mod.rs index f7e9cfc..9b9cd88 100644 --- a/pbf-craft/src/models/mod.rs +++ b/pbf-craft/src/models/mod.rs @@ -199,7 +199,8 @@ pub struct RelationMember { pub role: String, } -pub trait BasicElement { +pub trait BasicElement: Clone { + fn get_element_type() -> ElementType; fn get_id(&self) -> i64; fn get_version(&self) -> i32; fn get_timestamp(&self) -> Option>; @@ -210,6 +211,10 @@ pub trait BasicElement { } impl BasicElement for Node { + fn get_element_type() -> ElementType { + ElementType::Node + } + fn get_id(&self) -> i64 { self.id } @@ -240,6 +245,10 @@ impl BasicElement for Node { } impl BasicElement for Way { + fn get_element_type() -> ElementType { + ElementType::Way + } + fn get_id(&self) -> i64 { self.id } @@ -270,6 +279,10 @@ impl BasicElement for Way { } impl BasicElement for Relation { + fn get_element_type() -> ElementType { + ElementType::Relation + } + fn get_id(&self) -> i64 { self.id } diff --git a/pbf-craft/src/readers/indexed_reader.rs b/pbf-craft/src/readers/indexed_reader.rs index b01a827..6163cb4 100644 --- a/pbf-craft/src/readers/indexed_reader.rs +++ b/pbf-craft/src/readers/indexed_reader.rs @@ -10,7 +10,8 @@ use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}; use super::cached_reader::CachedReader; use super::raw_reader::PbfReader; use super::traits::PbfRandomRead; -use crate::models::{Element, ElementType, Node, Relation, Way}; +use crate::models::{BasicElement, Element, ElementType, Node, Relation, Way}; +use crate::readers::traits::BlobData; use crate::utils::file; fn get_index_path_from_pbf_path(pbf_path: &str) -> String { @@ -248,59 +249,75 @@ impl IndexedReader { } impl IndexedReader { - /// Finds an node by its ID. - pub fn find_node(&mut self, node_id: i64) -> anyhow::Result> { - let has_offset = self.pbf_index.get_offset(&ElementType::Node, node_id); + fn find_element( + &mut self, + element_type: &ElementType, + element_id: i64, + get_vec: F, + ) -> anyhow::Result> + where + E: BasicElement, + F: FnOnce(&BlobData) -> &Vec, + { + let has_offset = self.pbf_index.get_offset(element_type, element_id); if has_offset.is_none() { return Ok(None); } let offset = has_offset.unwrap(); let blob_data = self.pbf_reader.read_blob_by_offset(offset)?; - let node = blob_data.nodes.iter().find(|node| node.id == node_id); - match node { - Some(n) => Ok(Some(n.clone())), + let elem = get_vec(&blob_data) + .iter() + .find(|e| e.get_id() == element_id); + match elem { + Some(e) => Ok(Some(e.clone())), None => Ok(None), } } - /// Finds nodes by their IDs. - /// - /// `find_nodes` is more efficient than calling `find_node` multiple times when you have a batch of node IDs. - /// - pub fn find_nodes(&mut self, node_ids: &[i64]) -> anyhow::Result> { - let offsets: HashSet = node_ids + fn find_elements( + &mut self, + element_type: &ElementType, + element_ids: &[i64], + get_vec: F, + ) -> anyhow::Result> + where + E: BasicElement, + F: Fn(&BlobData) -> &Vec, + { + let offsets: HashSet = element_ids .into_iter() - .filter_map(|id| self.pbf_index.get_offset(&ElementType::Node, *id)) + .filter_map(|id| self.pbf_index.get_offset(element_type, *id)) .collect(); - let result: Vec = offsets + let result: Vec = offsets .into_iter() .flat_map(|offset| { let blob_data = self.pbf_reader.read_blob_by_offset(offset).unwrap(); - let nodes: Vec = blob_data - .nodes + get_vec(&blob_data) .iter() - .filter(|node| node_ids.contains(&node.id)) - .map(|node| node.clone()) - .collect(); - nodes + .filter(|e| element_ids.contains(&e.get_id())) + .cloned() + .collect::>() }) .collect(); Ok(result) } + /// Finds an node by its ID. + pub fn find_node(&mut self, node_id: i64) -> anyhow::Result> { + self.find_element(&ElementType::Node, node_id, |blob_data| &blob_data.nodes) + } + + /// Finds nodes by their IDs. + /// + /// `find_nodes` is more efficient than calling `find_node` multiple times when you have a batch of node IDs. + /// + pub fn find_nodes(&mut self, node_ids: &[i64]) -> anyhow::Result> { + self.find_elements(&ElementType::Node, node_ids, |blob_data| &blob_data.nodes) + } + /// Finds a way by its ID. pub fn find_way(&mut self, way_id: i64) -> anyhow::Result> { - let has_offset = self.pbf_index.get_offset(&ElementType::Way, way_id); - if has_offset.is_none() { - return Ok(None); - } - let offset = has_offset.unwrap(); - let blob_data = self.pbf_reader.read_blob_by_offset(offset)?; - let way = blob_data.ways.iter().find(|way| way.id == way_id); - match way { - Some(w) => Ok(Some(w.clone())), - None => Ok(None), - } + self.find_element(&ElementType::Way, way_id, |blob_data| &blob_data.ways) } /// Finds ways by their IDs. @@ -308,44 +325,14 @@ impl IndexedReader { /// `find_ways` is more efficient than calling `find_way` multiple times when you have a batch of way IDs. /// pub fn find_ways(&mut self, way_ids: &[i64]) -> anyhow::Result> { - let offsets: HashSet = way_ids - .into_iter() - .filter_map(|id| self.pbf_index.get_offset(&ElementType::Way, *id)) - .collect(); - let result: Vec = offsets - .into_iter() - .flat_map(|offset| { - let blob_data = self.pbf_reader.read_blob_by_offset(offset).unwrap(); - let ways: Vec = blob_data - .ways - .iter() - .filter(|way| way_ids.contains(&way.id)) - .map(|way| way.clone()) - .collect(); - ways - }) - .collect(); - Ok(result) + self.find_elements(&ElementType::Way, way_ids, |blob_data| &blob_data.ways) } /// Finds a relation by its ID. pub fn find_relation(&mut self, relation_id: i64) -> anyhow::Result> { - let has_offset = self - .pbf_index - .get_offset(&ElementType::Relation, relation_id); - if has_offset.is_none() { - return Ok(None); - } - let offset = has_offset.unwrap(); - let blob_data = self.pbf_reader.read_blob_by_offset(offset)?; - let rel = blob_data - .relations - .iter() - .find(|relation| relation.id == relation_id); - match rel { - Some(r) => Ok(Some(r.clone())), - None => Ok(None), - } + self.find_element(&ElementType::Relation, relation_id, |blob_data| { + &blob_data.relations + }) } /// Finds relations by their IDs. @@ -353,24 +340,9 @@ impl IndexedReader { /// `find_relations` is more efficient than calling `find_relation` multiple times when you have a batch of relation IDs. /// pub fn find_relations(&mut self, relation_ids: &[i64]) -> anyhow::Result> { - let offsets: HashSet = relation_ids - .into_iter() - .filter_map(|id| self.pbf_index.get_offset(&ElementType::Relation, *id)) - .collect(); - let result: Vec = offsets - .into_iter() - .flat_map(|offset| { - let blob_data = self.pbf_reader.read_blob_by_offset(offset).unwrap(); - let relations: Vec = blob_data - .relations - .iter() - .filter(|relation| relation_ids.contains(&relation.id)) - .map(|relation| relation.clone()) - .collect(); - relations - }) - .collect(); - Ok(result) + self.find_elements(&ElementType::Relation, relation_ids, |blob_data| { + &blob_data.relations + }) } /// Finds an element by its type and ID. @@ -552,7 +524,7 @@ mod tests { if let Element::Node(node) = target { assert_eq!(node.id, 4254529698); } else { - assert!(false); + panic!("Expected Node element"); } let target_op = indexed_reader.find(&ElementType::Way, 1055523837).unwrap(); @@ -560,17 +532,70 @@ mod tests { if let Element::Way(way) = target { assert_eq!(way.id, 1055523837); } else { - assert!(false); + panic!("Expected Way element"); } } + #[test] + fn test_find_nonexistent_element() { + let pbf_file = "./resources/andorra-latest.osm.pbf"; + let mut indexed_reader = IndexedReader::from_path(pbf_file).unwrap(); + + // Test non-existent node + assert!(indexed_reader.find_node(-1).unwrap().is_none()); + + // Test non-existent way + assert!(indexed_reader.find_way(-1).unwrap().is_none()); + + // Test non-existent relation + assert!(indexed_reader.find_relation(-1).unwrap().is_none()); + } + + #[test] + fn test_batch_operations() { + let pbf_file = "./resources/andorra-latest.osm.pbf"; + let mut indexed_reader = IndexedReader::from_path(pbf_file).unwrap(); + + // Test find_nodes with mixed existing/non-existing IDs + let node_ids = vec![4254529698, 4254529699, -1]; // Last ID doesn't exist + let nodes = indexed_reader.find_nodes(&node_ids).unwrap(); + assert_eq!(nodes.len(), 2); + assert!(nodes.iter().any(|n| n.id == 4254529698)); + assert!(nodes.iter().any(|n| n.id == 4254529699)); + + // Test empty input + assert!(indexed_reader.find_nodes(&[]).unwrap().is_empty()); + } + + #[test] + fn test_dependency_resolution() { + let pbf_file = "./resources/andorra-latest.osm.pbf"; + let mut indexed_reader = IndexedReader::from_path(pbf_file).unwrap(); + + // Test getting a way with its nodes + let way_with_deps = indexed_reader + .get_with_deps(&ElementType::Way, 1055523837) + .unwrap(); + assert!(way_with_deps.len() > 1); + assert!(way_with_deps.iter().any(|e| matches!(e, Element::Way(_)))); + assert!(way_with_deps.iter().any(|e| matches!(e, Element::Node(_)))); + } + + #[test] + fn test_invalid_file_handling() { + // Test with non-existent PBF file + assert!(PbfIndex::new("nonexistent.pbf").is_err()); + + // Test with invalid PIF file + assert!(PbfIndex::load_from_file("nonexistent.pif").is_err()); + } + #[bench] fn bench_find_without_cache(b: &mut Bencher) { let pbf_file = "./resources/andorra-latest.osm.pbf"; let mut indexed_reader = IndexedReader::from_path(pbf_file).unwrap(); b.iter(|| { - // Inner closure, the actual test for _ in 1..30 { let target_op = indexed_reader.find(&ElementType::Node, 4254529698).unwrap(); target_op.unwrap(); @@ -584,11 +609,21 @@ mod tests { let mut indexed_reader = IndexedReader::from_path_with_cache(pbf_file, 10000).unwrap(); b.iter(|| { - // Inner closure, the actual test for _ in 1..30 { let target_op = indexed_reader.find(&ElementType::Node, 4254529698).unwrap(); target_op.unwrap(); } }); } + + #[bench] + fn bench_batch_operations(b: &mut Bencher) { + let pbf_file = "./resources/andorra-latest.osm.pbf"; + let mut indexed_reader = IndexedReader::from_path(pbf_file).unwrap(); + let node_ids = vec![4254529698, 4254529699, 4254529700]; + + b.iter(|| { + indexed_reader.find_nodes(&node_ids).unwrap(); + }); + } }