Rearranged the existing code to make it more manageable.

src/gtfs/loader.rs

@@ -0,0 +1,200 @@
+use std::{collections::{HashMap, HashSet}, fs::File, hash::Hash};
+use gtfs_structures::{Calendar, CalendarDate, RawStopTime, RawTrip, Route, Stop};
+use serde::de::DeserializeOwned;
+use zip::ZipArchive;
+
+use crate::gtfs::{structs::Gtfs, utils::{route_ids_from_numbers, stop_ids_from_codes}};
+
+trait Filter<T> {
+    fn accept(&self, v: &T) -> bool;
+}
+
+// No filter on loaded records
+struct LoadAll {}
+impl<T> Filter<T> for LoadAll {
+    fn accept(&self, _: &T) -> bool {
+        return true;
+    }
+}
+
+struct LoadRoutes<'a> {
+    routes: &'a HashSet<String>,
+}
+impl Filter<Route> for LoadRoutes<'_> {
+    fn accept(&self, r: &Route) -> bool {
+        let short_name = &r.short_name;
+        return short_name.is_some() && self.routes.contains(short_name.as_ref().unwrap());
+    }
+}
+
+struct LoadStops<'a> {
+    stops: &'a HashSet<String>,
+}
+impl Filter<Stop> for LoadStops<'_> {
+    fn accept(&self, s: &Stop) -> bool {
+        let stop_code = &s.code;
+        return stop_code.is_some() && self.stops.contains(s.code.as_ref().unwrap());
+    }
+}
+
+struct LoadTrips<'a> {
+    route_ids: &'a HashSet<String>,
+}
+impl Filter<RawTrip> for LoadTrips<'_> {
+    fn accept(&self, t: &RawTrip) -> bool {
+        let route_id = &t.route_id;
+        return self.route_ids.contains(route_id);
+    }
+}
+
+struct LoadStopTimes<'a> {
+    trip_ids: &'a HashSet<String>,
+    stop_ids: &'a HashSet<String>,
+}
+impl Filter<RawStopTime> for LoadStopTimes<'_> {
+    fn accept(&self, st: &RawStopTime) -> bool {
+        return self.stop_ids.contains(&st.stop_id) && self.trip_ids.contains(&st.trip_id);
+    }
+}
+
+// Loads a vector of the selected type
+fn load_vector<T: serde::de::DeserializeOwned>(
+    destination: &mut Vec<T>,
+    zip_reader: &mut ZipArchive<File>,
+    table_name: &str,
+) {
+    let file_reader = zip_reader.by_name(table_name).unwrap();
+    let mut rdr = csv::Reader::from_reader(file_reader);
+
+    for row in rdr.deserialize() {
+        let record: T = row.unwrap();
+        destination.push(record);
+    }
+}
+
+// Loads a HashMap of the selected type, using the provided index function as the key
+fn load_map<K, V, IndexFn, FilterT>(
+    destination: &mut HashMap<K, V>,
+    zip_reader: &mut ZipArchive<File>,
+    table_name: &str,
+    index: IndexFn,
+    filter: FilterT,
+) where
+  K: Eq + Hash,
+  V: DeserializeOwned,
+  IndexFn: Fn(&V) -> K,
+  FilterT: Filter<V>,
+ {
+    let file_reader = zip_reader.by_name(table_name).unwrap();
+    let mut rdr = csv::Reader::from_reader(file_reader);
+
+    for row in rdr.deserialize() {
+        if row.is_ok() {
+            let record: V = row.unwrap();
+            if filter.accept(&record) {
+                let idx: K = index(&record);
+                destination.insert(idx, record);
+            }
+        } else {
+            print!("Row failed to deserialize row {:#?}", row.err());
+            panic!();
+        }
+    }
+}
+
+// Loads a HashMap of a vector of the selected type, using the provided index function as the key
+// And a predicate as a filter
+fn load_vector_map<'a, K, V, IndexFn, FilterT>(
+    destination: &mut HashMap<K, Vec<V>>,
+    zip_reader: &mut ZipArchive<File>,
+    table_name: &str,
+    index: IndexFn,
+    filter: FilterT,
+) where 
+  K: Eq + Hash,
+  V: DeserializeOwned,
+  IndexFn: Fn(&V) -> K,
+  FilterT: Filter<V>,
+{
+    let file_reader = zip_reader.by_name(table_name).unwrap();
+    let mut rdr = csv::Reader::from_reader(file_reader);
+
+    for row in rdr.deserialize() {
+        let record: V = row.unwrap();
+        if filter.accept(&record) {
+            let idx = index(&record);
+            destination.entry(idx).or_insert_with(Vec::new).push(record);
+        }
+    }
+}
+
+
+pub fn load_gtfs(gtfs: &mut Gtfs, zip_reader: &mut ZipArchive<File>, route_numbers: HashSet<String>, stop_codes: HashSet<String>) {
+
+    // Agencies
+    load_vector(&mut gtfs.agencies, zip_reader, "agency.txt");
+
+    // Calendars
+    load_map(
+        &mut gtfs.calendar,
+        zip_reader,
+        "calendar.txt",
+        |c: &Calendar| String::from(&c.id),
+        LoadAll {},
+    );
+
+    // Calendar Dates
+    load_vector_map(
+        &mut gtfs.calendar_dates,
+        zip_reader,
+        "calendar_dates.txt",
+        |d: &CalendarDate| String::from(&d.service_id),
+        LoadAll {},
+    );
+
+    // Routes
+    load_map(
+        &mut gtfs.routes,
+        zip_reader,
+        "routes.txt",
+        |r: &Route| String::from(&r.id),
+        LoadRoutes {
+            routes: &route_numbers,
+        },
+    );
+
+    // Stops
+    load_map(
+        &mut gtfs.stops,
+        zip_reader,
+        "stops.txt",
+        |s: &Stop| String::from(&s.id),
+        LoadStops { stops: &stop_codes },
+    );
+
+    let route_ids = route_ids_from_numbers(&gtfs, &route_numbers);
+    // Trips
+    load_map(
+        &mut gtfs.trips,
+        zip_reader,
+        "trips.txt",
+        |t: &RawTrip| String::from(&t.id),
+        LoadTrips {
+            route_ids: &route_ids,
+        },
+    );
+
+    // Load stop times for the chosen routes and stops
+    let stop_ids = stop_ids_from_codes(&gtfs, &stop_codes);
+    let trip_ids = HashSet::<String>::from_iter(gtfs.trips.keys().cloned());
+    load_map(
+        &mut gtfs.stop_times,
+        zip_reader,
+        "stop_times.txt",
+        |st: &RawStopTime| (st.trip_id.clone(), st.stop_sequence),
+        LoadStopTimes {
+            trip_ids: &trip_ids,
+            stop_ids: &stop_ids,
+        },
+    );
+}