compact data structures for strings

catenarytransit · Sep 18, 2024 · b41efd9 · b41efd9
1 parent 9defeb5
commit b41efd9
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 15 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -134,6 +134,7 @@ via-rail-gtfsrt = "0.1.1"
 wyhash2 = "0.2.1"
 console-subscriber = "0.4.0"
 strumbra = "0.4.0"
+compact_str = { version = "0.8.0", features = ["serde"] }
 [[bin]]
 name = "maple"
 path = "src/maple/main.rs"

diff --git a/src/aspen/import_alpenrose.rs b/src/aspen/import_alpenrose.rs
@@ -12,6 +12,7 @@ use diesel::SelectableHelper;
 use diesel_async::RunQueryDsl;
 use gtfs_realtime::FeedMessage;
 use scc::HashMap as SccHashMap;
+use compact_str::CompactString;
 use serde::Deserialize;
 use serde::Serialize;
 use std::collections::HashMap;
@@ -96,7 +97,7 @@ pub async fn new_rt_data(
     let mut gtfs_vehicle_labels_to_ids: AHashMap<String, String> = AHashMap::new();
     let mut vehicle_routes_cache: AHashMap<String, AspenisedVehicleRouteCache> = AHashMap::new();
     let mut trip_updates: AHashMap<String, AspenisedTripUpdate> = AHashMap::new();
-    let mut trip_updates_lookup_by_trip_id_to_trip_update_ids: AHashMap<String, Vec<String>> =
+    let mut trip_updates_lookup_by_trip_id_to_trip_update_ids: AHashMap<CompactString, Vec<CompactString>> =
         AHashMap::new();
 
     //let alerts hashmap
@@ -390,7 +391,7 @@ pub async fn new_rt_data(
                                 .iter()
                                 .map(|stu| AspenisedStopTimeUpdate {
                                     stop_sequence: stu.stop_sequence,
-                                    stop_id: stu.stop_id.clone(),
+                                    stop_id: stu.stop_id.as_ref().map(|x| x.into()),
                                     arrival: stu.arrival.clone().map(|arrival| {
                                         AspenStopTimeEvent {
                                             delay: arrival.delay,
@@ -421,8 +422,8 @@ pub async fn new_rt_data(
 
                         if trip_id.is_some() {
                             trip_updates_lookup_by_trip_id_to_trip_update_ids
-                                .entry(trip_id.as_ref().unwrap().clone())
-                                .or_insert(vec![trip_update_entity.id.clone()]);
+                                .entry(trip_id.as_ref().unwrap().into())
+                                .or_insert(vec![trip_update_entity.id.clone().into()]);
                         }
 
                         trip_updates.insert(trip_update_entity.id.clone(), trip_update);

diff --git a/src/aspen/main.rs b/src/aspen/main.rs
@@ -45,6 +45,7 @@ use tarpc::{
     tokio_serde::formats::Bincode,
 };
 use tokio::sync::Mutex;
+use compact_str::CompactString;
 use tokio::time;
 use uuid::Uuid;
 mod leader_thread;
@@ -125,16 +126,16 @@ impl AspenRpc for AspenServer {
                 for trip_id in trip_id_list {
                     if let Some(trip_update_id_list) = authoritative_data
                         .trip_updates_lookup_by_trip_id_to_trip_update_ids
-                        .get(&trip_id)
+                        .get(trip_id.as_str())
                     {
                         trip_id_to_trip_update_ids
-                            .insert(trip_id.clone(), trip_update_id_list.clone());
+                            .insert(trip_id.clone(), trip_update_id_list.iter().map(|x| x.to_string()).collect::<Vec<String>>());
 
                         for trip_update_id in trip_update_id_list {
                             if let Some(trip_update) =
-                                authoritative_data.trip_updates.get(trip_update_id)
+                                authoritative_data.trip_updates.get(trip_update_id.as_str())
                             {
-                                trip_updates.insert(trip_update_id.clone(), trip_update.clone());
+                                trip_updates.insert(trip_update_id.to_string(), trip_update.clone());
                             }
                         }
                     }
@@ -552,14 +553,14 @@ impl AspenRpc for AspenServer {
 
                 let trip_updates_id_list = aspenised_data
                     .trip_updates_lookup_by_trip_id_to_trip_update_ids
-                    .get(&trip_id);
+                    .get(trip_id.as_str());
 
                 match trip_updates_id_list {
                     Some(trip_updates_id_list) => {
                         let mut trip_updates = Vec::new();
 
                         for trip_update_id in trip_updates_id_list {
-                            let trip_update = aspenised_data.trip_updates.get(trip_update_id);
+                            let trip_update = aspenised_data.trip_updates.get(trip_update_id.as_str());
 
                             match trip_update {
                                 Some(trip_update) => {

diff --git a/src/birch/get_vehicle_trip_information.rs b/src/birch/get_vehicle_trip_information.rs
@@ -221,7 +221,7 @@ struct GtfsRtRefreshData {
 
 #[derive(Deserialize, Serialize, Clone, Debug)]
 struct StopTimeRefresh {
-    pub stop_id: Option<String>,
+    pub stop_id: Option<compact_str::CompactString>,
     pub rt_arrival: Option<AspenStopTimeEvent>,
     pub rt_departure: Option<AspenStopTimeEvent>,
     pub schedule_relationship: Option<i32>,

diff --git a/src/birch/nearby_departures.rs b/src/birch/nearby_departures.rs
@@ -1096,7 +1096,9 @@ pub async fn nearby_from_coords(
                                         } else {
                                             let relevant_stop_time_update =
                                                 trip_update.stop_time_update.iter().find(|x| {
-                                                    x.stop_id.as_ref()
+                                                    x.stop_id
+                                                        .as_ref()
+                                                        .map(|compare| compare.as_str())
                                                         == Some(&trip.itinerary_options[0].stop_id)
                                                 });
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -183,7 +183,7 @@ pub mod tailscale {
 
 pub mod aspen_dataset {
     use ahash::AHashMap;
-
+    use compact_str::CompactString;
     use std::hash::Hash;
 
     #[derive(Clone, Serialize, Deserialize)]
@@ -193,7 +193,7 @@ pub mod aspen_dataset {
         pub vehicle_label_to_gtfs_id: AHashMap<String, String>,
         //id to trip update
         pub trip_updates: AHashMap<String, AspenisedTripUpdate>,
-        pub trip_updates_lookup_by_trip_id_to_trip_update_ids: AHashMap<String, Vec<String>>,
+        pub trip_updates_lookup_by_trip_id_to_trip_update_ids: AHashMap<CompactString, Vec<CompactString>>,
         //        pub raw_alerts: AHashMap<String, gtfs_realtime::Alert>,
         pub aspenised_alerts: AHashMap<String, AspenisedAlert>,
         pub impacted_routes_alerts: AHashMap<String, Vec<String>>,
@@ -412,7 +412,7 @@ pub mod aspen_dataset {
     #[derive(Clone, Debug, Serialize, Deserialize)]
     pub struct AspenisedStopTimeUpdate {
         pub stop_sequence: Option<u32>,
-        pub stop_id: Option<String>,
+        pub stop_id: Option<compact_str::CompactString>,
         pub arrival: Option<AspenStopTimeEvent>,
         pub departure: Option<AspenStopTimeEvent>,
         pub departure_occupancy_status: Option<i32>,