Refactor graph scaping code. a-b-street/utils#2

Noting that the tests caught a bug -- I forgot to transform barrier points at first!
a-b-street · Apr 26, 2024 · 64e2c2b · 64e2c2b
1 parent b29c783
commit 64e2c2b
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 116 deletions.
diff --git a/backend/Cargo.lock b/backend/Cargo.lock
diff --git a/backend/Cargo.toml b/backend/Cargo.toml
@@ -7,7 +7,7 @@ edition = "2021"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-anyhow = "1.0.75"
+anyhow = "1.0.82"
 bincode = "1.3.3"
 console_error_panic_hook = "0.1.6"
 console_log = "1.0.0"

diff --git a/backend/src/scrape.rs b/backend/src/scrape.rs
@@ -1,25 +1,22 @@
 use std::collections::{BTreeMap, BTreeSet, HashMap};
 
 use anyhow::Result;
-use geo::{ConvexHull, Coord, Geometry, GeometryCollection, LineString, Point};
-use osm_reader::{Element, NodeID, WayID};
-use utils::{Mercator, Tags};
+use geo::Coord;
+use osm_reader::{Element, NodeID};
+use utils::Tags;
 
 use crate::{Direction, FilterKind, Intersection, IntersectionID, MapModel, Road, RoadID, Router};
 
-struct Way {
-    id: WayID,
-    node_ids: Vec<NodeID>,
-    tags: Tags,
-}
-
 pub fn scrape_osm(input_bytes: &[u8], study_area_name: Option<String>) -> Result<MapModel> {
     info!("Parsing {} bytes of OSM data", input_bytes.len());
+    // This doesn't use osm2graph's helper, because it needs to scrape more things from OSM
     let mut node_mapping = HashMap::new();
     let mut highways = Vec::new();
     let mut all_barriers: BTreeSet<NodeID> = BTreeSet::new();
     osm_reader::parse(input_bytes, |elem| match elem {
-        Element::Node { id, lon, lat, tags } => {
+        Element::Node {
+            id, lon, lat, tags, ..
+        } => {
             let pt = Coord { x: lon, y: lat };
             node_mapping.insert(id, pt);
 
@@ -36,6 +33,7 @@ pub fn scrape_osm(input_bytes: &[u8], study_area_name: Option<String>) -> Result
             id,
             mut node_ids,
             tags,
+            ..
         } => {
             let tags = tags.into();
             if is_road(&tags) {
@@ -46,11 +44,12 @@ pub fn scrape_osm(input_bytes: &[u8], study_area_name: Option<String>) -> Result
                     warn!("{id} refers to nodes outside the imported area");
                 }
                 if node_ids.len() >= 2 {
-                    highways.push(Way { id, node_ids, tags });
+                    highways.push(utils::osm2graph::Way { id, node_ids, tags });
                 }
             }
         }
         Element::Relation { .. } => {}
+        Element::Bounds { .. } => {}
     })?;
 
     // There'll be many barrier nodes on non-driveable paths we don't consider roads. Filter for
@@ -65,44 +64,49 @@ pub fn scrape_osm(input_bytes: &[u8], study_area_name: Option<String>) -> Result
     }
 
     info!("Splitting {} ways into edges", highways.len());
-    let (mut roads, mut intersections) = split_edges(&node_mapping, highways);
+    let graph = utils::osm2graph::Graph::from_scraped_osm(node_mapping, highways);
+    // Copy all the fields
+    let intersections: Vec<Intersection> = graph
+        .intersections
+        .into_iter()
+        .map(|i| Intersection {
+            id: IntersectionID(i.id.0),
+            point: i.point,
+            node: i.osm_node,
+            roads: i.edges.into_iter().map(|e| RoadID(e.0)).collect(),
+        })
+        .collect();
+
+    // Add in a bit
+    let roads: Vec<Road> = graph
+        .edges
+        .into_iter()
+        .map(|e| Road {
+            id: RoadID(e.id.0),
+            src_i: IntersectionID(e.src.0),
+            dst_i: IntersectionID(e.dst.0),
+            way: e.osm_way,
+            node1: e.osm_node1,
+            node2: e.osm_node2,
+            linestring: e.linestring,
+            tags: e.osm_tags,
+        })
+        .collect();
+    for coord in &mut barrier_pts {
+        *coord = graph.mercator.pt_to_mercator(*coord);
+    }
     info!("Finalizing the map model");
 
     let mut directions = BTreeMap::new();
     for r in &roads {
         directions.insert(r.id, Direction::from_osm(&r.tags));
     }
 
-    // TODO expensive
-    let mut collection: GeometryCollection = roads
-        .iter()
-        .map(|r| Geometry::LineString(r.linestring.clone()))
-        .chain(
-            intersections
-                .iter()
-                .map(|i| Geometry::Point(i.point.clone())),
-        )
-        .collect::<Vec<_>>()
-        .into();
-    let mercator = Mercator::from(collection.clone()).unwrap();
-    for r in &mut roads {
-        mercator.to_mercator_in_place(&mut r.linestring);
-    }
-    for i in &mut intersections {
-        mercator.to_mercator_in_place(&mut i.point);
-    }
-    for coord in &mut barrier_pts {
-        *coord = mercator.pt_to_mercator(*coord);
-    }
-
-    mercator.to_mercator_in_place(&mut collection);
-    let boundary_polygon = collection.convex_hull();
-
     let mut map = MapModel {
         roads,
         intersections,
-        mercator,
-        boundary_polygon,
+        mercator: graph.mercator,
+        boundary_polygon: graph.boundary_polygon,
         study_area_name,
 
         router_original: None,
@@ -142,77 +146,6 @@ pub fn scrape_osm(input_bytes: &[u8], study_area_name: Option<String>) -> Result
     Ok(map)
 }
 
-fn split_edges(
-    node_mapping: &HashMap<NodeID, Coord>,
-    ways: Vec<Way>,
-) -> (Vec<Road>, Vec<Intersection>) {
-    // Count how many ways reference each node
-    let mut node_counter: HashMap<NodeID, usize> = HashMap::new();
-    for way in &ways {
-        for node in &way.node_ids {
-            *node_counter.entry(*node).or_insert(0) += 1;
-        }
-    }
-
-    // Split each way into edges
-    let mut node_to_intersection: HashMap<NodeID, IntersectionID> = HashMap::new();
-    let mut intersections = Vec::new();
-    let mut roads = Vec::new();
-    for way in ways {
-        let mut node1 = way.node_ids[0];
-        let mut pts = Vec::new();
-
-        let num_nodes = way.node_ids.len();
-        for (idx, node) in way.node_ids.into_iter().enumerate() {
-            pts.push(node_mapping[&node]);
-            // Edges start/end at intersections between two ways. The endpoints of the way also
-            // count as intersections.
-            let is_endpoint =
-                idx == 0 || idx == num_nodes - 1 || *node_counter.get(&node).unwrap() > 1;
-            if is_endpoint && pts.len() > 1 {
-                let road_id = RoadID(roads.len());
-
-                let mut i_ids = Vec::new();
-                for (n, point) in [(node1, pts[0]), (node, *pts.last().unwrap())] {
-                    let intersection = if let Some(i) = node_to_intersection.get(&n) {
-                        &mut intersections[i.0]
-                    } else {
-                        let i = IntersectionID(intersections.len());
-                        intersections.push(Intersection {
-                            id: i,
-                            node: n,
-                            point: Point(point),
-                            roads: Vec::new(),
-                        });
-                        node_to_intersection.insert(n, i);
-                        &mut intersections[i.0]
-                    };
-
-                    intersection.roads.push(road_id);
-                    i_ids.push(intersection.id);
-                }
-
-                roads.push(Road {
-                    id: road_id,
-                    src_i: i_ids[0],
-                    dst_i: i_ids[1],
-                    way: way.id,
-                    node1,
-                    node2: node,
-                    linestring: LineString::new(std::mem::take(&mut pts)),
-                    tags: way.tags.clone(),
-                });
-
-                // Start the next edge
-                node1 = node;
-                pts.push(node_mapping[&node]);
-            }
-        }
-    }
-
-    (roads, intersections)
-}
-
 fn is_road(tags: &Tags) -> bool {
     if !tags.has("highway") || tags.is("area", "yes") {
         return false;