search: Extract function for parallel blob iteration

src/search.rs

@@ -1,7 +1,7 @@
-use scoped_threadpool::Pool;
 use coord::LatLon;
-use osmpbf::{Blob, BlobDecode, BlobReader};
+use osmpbf::{Blob, BlobDecode, BlobReader, PrimitiveBlock};
 use regex::Regex;
+use scoped_threadpool::Pool;
 use std::path::{Path, PathBuf};
 use std::sync::mpsc::sync_channel;
 use std::thread;
@@ -55,62 +55,99 @@ pub fn par_search_blocking<P, F>(
 ) -> Result<(), String>
 where P: AsRef<Path>,
       F: Fn(Vec<LatLon>) -> ControlFlow + Send + 'static,
+{
+    let re = Regex::new(search_pattern)
+        .map_err(|e| format!("{}", e))?;
+    let re = &re;
+
+    let search = move |block: &PrimitiveBlock, _: &()| {
+        let mut matches = vec![];
+
+        for node in block.groups().flat_map(|g| g.nodes()) {
+            for (_key, val) in node.tags() {
+                if re.is_match(val) {
+                    let pos = LatLon::new(node.lat(), node.lon());
+                    matches.push(pos);
+                    break;
+                }
+            }
+        }
+
+        for node in block.groups().flat_map(|g| g.dense_nodes()) {
+            for (_key, val) in node.tags() {
+                if re.is_match(val) {
+                    let pos = LatLon::new(node.lat(), node.lon());
+                    matches.push(pos);
+                    break;
+                }
+            }
+        }
+
+        matches
+    };
+
+    par_iter_blobs(
+        pbf_path,
+        || {},
+        search,
+        found_func,
+    )
+}
+
+fn par_iter_blobs<P, D, R, IF, CF, RF>(
+    pbf_path: P,
+    init_func: IF,
+    compute_func: CF,
+    result_func: RF,
+) -> Result<(), String>
+where P: AsRef<Path>,
+      IF: Fn() -> D,
+      CF: Fn(&PrimitiveBlock, &D) -> R + Send + Sync,
+      RF: Fn(R) -> ControlFlow + Send + 'static,
+      R: Send,
+      D: Send,
 {
     let num_threads = ::num_cpus::get();
     let mut pool = Pool::new(num_threads as u32);
 
     pool.scoped(|scope| {
-        let re = Regex::new(search_pattern)
-            .map_err(|e| format!("{}", e))?;
         let mut reader = BlobReader::from_path(&pbf_path)
             .map_err(|e| format!("{}", e))?;
 
         let mut chans = Vec::with_capacity(num_threads);
-        let (result_tx, result_rx) = sync_channel::<(usize, Result<Vec<LatLon>, String>)>(0);
+        let (result_tx, result_rx) = sync_channel::<(usize, Result<Option<R>, String>)>(0);
 
         for thread_id in 0..num_threads {
-            let re = re.clone();
+            let thread_data = init_func();
             let result_tx = result_tx.clone();
 
             let (request_tx, request_rx) = sync_channel::<WorkerMessage>(0);
             chans.push(request_tx);
 
+            let compute = &compute_func;
+
             scope.execute(move || {
                 for request in request_rx.iter() {
                     match request {
                         WorkerMessage::PleaseStop => return,
                         WorkerMessage::DoBlob(blob) => {
-                            let mut matches = vec![];
-                            let block = match blob.decode() {
-                                Ok(b) => b,
+                            match blob.decode() {
+                                Ok(BlobDecode::OsmData(block)) => {
+                                    let result = compute(&block, &thread_data);
+                                    if result_tx.send((thread_id, Ok(Some(result)))).is_err() {
+                                        return;
+                                    }
+                                },
+                                //TODO also include other blob types in compute function
+                                Ok(_) => {
+                                    if result_tx.send((thread_id, Ok(None))).is_err() {
+                                        return;
+                                    }
+                                },
                                 Err(err) => {
                                     let _ = result_tx.send((thread_id, Err(format!("{}", err))));
                                     return;
-                                }
-                            };
-                            if let BlobDecode::OsmData(block) = block {
-                                for node in block.groups().flat_map(|g| g.nodes()) {
-                                    for (_key, val) in node.tags() {
-                                        if re.is_match(val) {
-                                            let pos = LatLon::new(node.lat(), node.lon());
-                                            matches.push(pos);
-                                            break;
-                                        }
-                                    }
-                                }
-
-                                for node in block.groups().flat_map(|g| g.dense_nodes()) {
-                                    for (_key, val) in node.tags() {
-                                        if re.is_match(val) {
-                                            let pos = LatLon::new(node.lat(), node.lon());
-                                            matches.push(pos);
-                                            break;
-                                        }
-                                    }
-                                }
-                            }
-                            if result_tx.send((thread_id, Ok(matches))).is_err() {
-                                return;
+                                },
                             }
                         }
                     };
@@ -144,10 +181,14 @@ where P: AsRef<Path>,
         }
 
         for (thread_id, matches) in result_rx.iter() {
-            let matches = matches?;
-
-            if found_func(matches) == ControlFlow::Break {
-                break;
+            match matches {
+                Err(err) => return Err(err),
+                Ok(Some(matches)) => {
+                    if result_func(matches) == ControlFlow::Break {
+                        break;
+                    }
+                },
+                _ => {},
             }
 
             match reader.next() {