hrdf_parser/
storage.rs

1use std::error::Error;
2
3use chrono::{Days, NaiveDate};
4use rustc_hash::{FxHashMap, FxHashSet};
5use serde::{Deserialize, Serialize};
6
7use crate::{
8    models::{
9        Attribute, BitField, Direction, ExchangeTimeAdministration, ExchangeTimeJourney,
10        ExchangeTimeLine, Holiday, InformationText, Journey, JourneyPlatform, Line, Model,
11        Platform, Stop, StopConnection, ThroughService, TimetableMetadataEntry, TransportCompany,
12        TransportType, Version,
13    },
14    parsing,
15    utils::{count_days_between_two_dates, timetable_end_date, timetable_start_date},
16};
17
18// ------------------------------------------------------------------------------------------------
19// --- DataStorage
20// ------------------------------------------------------------------------------------------------
21
22#[derive(Debug, Serialize, Deserialize)]
23pub struct DataStorage {
24    // Time-relevant data.
25    bit_fields: ResourceStorage<BitField>,
26    holidays: ResourceStorage<Holiday>,
27    timetable_metadata: ResourceStorage<TimetableMetadataEntry>,
28
29    // Basic data.
30    attributes: ResourceStorage<Attribute>,
31    information_texts: ResourceStorage<InformationText>,
32    directions: ResourceStorage<Direction>,
33    lines: ResourceStorage<Line>,
34    transport_companies: ResourceStorage<TransportCompany>,
35    transport_types: ResourceStorage<TransportType>,
36
37    // Stop data
38    stops: ResourceStorage<Stop>,
39    stop_connections: ResourceStorage<StopConnection>,
40
41    // Timetable data
42    journeys: ResourceStorage<Journey>,
43    journey_platform: ResourceStorage<JourneyPlatform>,
44    platforms: ResourceStorage<Platform>,
45    through_service: ResourceStorage<ThroughService>,
46
47    // Exchange times
48    exchange_times_administration: ResourceStorage<ExchangeTimeAdministration>,
49    exchange_times_journey: ResourceStorage<ExchangeTimeJourney>,
50    exchange_times_line: ResourceStorage<ExchangeTimeLine>,
51
52    // Maps
53    bit_fields_by_day: FxHashMap<NaiveDate, FxHashSet<i32>>,
54    bit_fields_by_stop_id: FxHashMap<i32, FxHashSet<i32>>,
55    journeys_by_stop_id_and_bit_field_id: FxHashMap<(i32, i32), Vec<i32>>,
56    stop_connections_by_stop_id: FxHashMap<i32, FxHashSet<i32>>,
57    exchange_times_administration_map: FxHashMap<(Option<i32>, String, String), i32>,
58    exchange_times_journey_map: FxHashMap<(i32, i32, i32), FxHashSet<i32>>,
59
60    // Additional global data
61    default_exchange_time: (i16, i16), // (InterCity exchange time, Exchange time for all other journey types)
62}
63
64#[allow(unused)]
65impl DataStorage {
66    pub fn new(version: Version, path: &str) -> Result<Self, Box<dyn Error>> {
67        // Time-relevant data
68        let bit_fields = parsing::load_bit_fields(path)?;
69        let holidays = parsing::load_holidays(path)?;
70        let timetable_metadata = parsing::load_timetable_metadata(path)?;
71
72        // Basic data
73        let (attributes, attributes_pk_type_converter) = parsing::load_attributes(path)?;
74        let (directions, directions_pk_type_converter) = parsing::load_directions(path)?;
75        let information_texts = parsing::load_information_texts(path)?;
76        let lines = parsing::load_lines(path)?;
77        let transport_companies = parsing::load_transport_companies(path)?;
78        let (transport_types, transport_types_pk_type_converter) =
79            parsing::load_transport_types(version, path)?;
80
81        // Stop data
82        let stop_connections = parsing::load_stop_connections(path, &attributes_pk_type_converter)?;
83        let (stops, default_exchange_time) = parsing::load_stops(version, path)?;
84
85        // Timetable data
86        let (journeys, journeys_pk_type_converter) = parsing::load_journeys(
87            path,
88            &transport_types_pk_type_converter,
89            &attributes_pk_type_converter,
90            &directions_pk_type_converter,
91        )?;
92        let (journey_platform, platforms) =
93            parsing::load_platforms(version, path, &journeys_pk_type_converter)?;
94        let through_service = parsing::load_through_service(path, &journeys_pk_type_converter)?;
95
96        // Exchange times
97        let exchange_times_administration = parsing::load_exchange_times_administration(path)?;
98        let exchange_times_journey =
99            parsing::load_exchange_times_journey(path, &journeys_pk_type_converter)?;
100        let exchange_times_line =
101            parsing::load_exchange_times_line(path, &transport_types_pk_type_converter)?;
102
103        log::info!("Building bit_fields_by_day...");
104        let bit_fields_by_day = create_bit_fields_by_day(&bit_fields, &timetable_metadata)?;
105        log::info!("Building bit_fields_by_stop_id...");
106        let bit_fields_by_stop_id = create_bit_fields_by_stop_id(&journeys);
107        log::info!("Building journeys_by_stop_id_and_bit_field_id...");
108        let journeys_by_stop_id_and_bit_field_id =
109            create_journeys_by_stop_id_and_bit_field_id(&journeys);
110        log::info!("Building stop_connections_by_stop_id...");
111        let stop_connections_by_stop_id = create_stop_connections_by_stop_id(&stop_connections);
112        log::info!("Building exchange_times_administration_map...");
113        let exchange_times_administration_map =
114            create_exchange_times_administration_map(&exchange_times_administration);
115        log::info!("Building exchange_times_journey_map...");
116        let exchange_times_journey_map = create_exchange_times_journey_map(&exchange_times_journey);
117
118        let mut data_storage = Self {
119            // Time-relevant data
120            bit_fields,
121            holidays,
122            timetable_metadata,
123            // Basic data
124            attributes,
125            information_texts,
126            directions,
127            lines,
128            transport_companies,
129            transport_types,
130            // Stop data
131            stop_connections,
132            stops,
133            // Timetable data
134            journeys,
135            journey_platform,
136            platforms,
137            through_service,
138            // Exchange times
139            exchange_times_administration,
140            exchange_times_journey,
141            exchange_times_line,
142            // Maps
143            bit_fields_by_day,
144            bit_fields_by_stop_id,
145            journeys_by_stop_id_and_bit_field_id,
146            stop_connections_by_stop_id,
147            exchange_times_administration_map,
148            exchange_times_journey_map,
149            // Additional global data
150            default_exchange_time,
151        };
152
153        Ok(data_storage)
154    }
155
156    // Getters/Setters
157
158    pub fn bit_fields(&self) -> &ResourceStorage<BitField> {
159        &self.bit_fields
160    }
161
162    pub fn journeys(&self) -> &ResourceStorage<Journey> {
163        &self.journeys
164    }
165
166    pub fn lines(&self) -> &ResourceStorage<Line> {
167        &self.lines
168    }
169
170    pub fn platforms(&self) -> &ResourceStorage<Platform> {
171        &self.platforms
172    }
173
174    pub fn stop_connections(&self) -> &ResourceStorage<StopConnection> {
175        &self.stop_connections
176    }
177
178    pub fn stops(&self) -> &ResourceStorage<Stop> {
179        &self.stops
180    }
181
182    pub fn transport_types(&self) -> &ResourceStorage<TransportType> {
183        &self.transport_types
184    }
185
186    pub fn timetable_metadata(&self) -> &ResourceStorage<TimetableMetadataEntry> {
187        &self.timetable_metadata
188    }
189
190    pub fn exchange_times_administration(&self) -> &ResourceStorage<ExchangeTimeAdministration> {
191        &self.exchange_times_administration
192    }
193
194    pub fn exchange_times_journey(&self) -> &ResourceStorage<ExchangeTimeJourney> {
195        &self.exchange_times_journey
196    }
197
198    pub fn exchange_times_line(&self) -> &ResourceStorage<ExchangeTimeLine> {
199        &self.exchange_times_line
200    }
201
202    pub fn bit_fields_by_day(&self) -> &FxHashMap<NaiveDate, FxHashSet<i32>> {
203        &self.bit_fields_by_day
204    }
205
206    pub fn bit_fields_by_stop_id(&self) -> &FxHashMap<i32, FxHashSet<i32>> {
207        &self.bit_fields_by_stop_id
208    }
209
210    pub fn journeys_by_stop_id_and_bit_field_id(&self) -> &FxHashMap<(i32, i32), Vec<i32>> {
211        &self.journeys_by_stop_id_and_bit_field_id
212    }
213
214    pub fn stop_connections_by_stop_id(&self) -> &FxHashMap<i32, FxHashSet<i32>> {
215        &self.stop_connections_by_stop_id
216    }
217
218    pub fn exchange_times_administration_map(
219        &self,
220    ) -> &FxHashMap<(Option<i32>, String, String), i32> {
221        &self.exchange_times_administration_map
222    }
223
224    pub fn exchange_times_journey_map(&self) -> &FxHashMap<(i32, i32, i32), FxHashSet<i32>> {
225        &self.exchange_times_journey_map
226    }
227
228    pub fn default_exchange_time(&self) -> (i16, i16) {
229        self.default_exchange_time
230    }
231}
232
233// ------------------------------------------------------------------------------------------------
234// --- ResourceStorage
235// ------------------------------------------------------------------------------------------------
236
237#[derive(Debug, Serialize, Deserialize)]
238pub struct ResourceStorage<M: Model<M>> {
239    data: FxHashMap<M::K, M>,
240}
241
242impl<M: Model<M>> ResourceStorage<M> {
243    pub fn new(data: FxHashMap<M::K, M>) -> Self {
244        Self { data }
245    }
246
247    pub fn data(&self) -> &FxHashMap<M::K, M> {
248        &self.data
249    }
250
251    pub fn find(&self, k: M::K) -> Option<&M> {
252        // TODO: there might be a problem when k is not in data so we can't unwrap here
253        self.data().get(&k)
254    }
255
256    pub fn entries(&self) -> Vec<&M> {
257        self.data.values().collect()
258    }
259
260    pub fn resolve_ids(&self, ids: &FxHashSet<M::K>) -> Option<Vec<&M>> {
261        ids.iter().map(|&id| self.find(id)).collect()
262    }
263}
264
265// ------------------------------------------------------------------------------------------------
266// --- Maps
267// ------------------------------------------------------------------------------------------------
268
269fn create_bit_fields_by_day(
270    bit_fields: &ResourceStorage<BitField>,
271    timetable_metadata: &ResourceStorage<TimetableMetadataEntry>,
272) -> Result<FxHashMap<NaiveDate, FxHashSet<i32>>, Box<dyn Error>> {
273    let start_date = timetable_start_date(timetable_metadata)?;
274    let num_days =
275        count_days_between_two_dates(start_date, timetable_end_date(timetable_metadata)?);
276
277    let dates: Vec<NaiveDate> = (0..num_days)
278        .map(|i| {
279            start_date
280                // unwrap: Converting i from usize to u64 will never fail.
281                .checked_add_days(Days::new(i.try_into().unwrap()))
282                // unwrap: Adding days will never fail.
283                .unwrap()
284        })
285        .collect();
286
287    let mut map = FxHashMap::default();
288    dates.iter().for_each(|date| {
289        map.entry(*date).or_insert(FxHashSet::default()).insert(0);
290    });
291
292    let result = bit_fields.data().keys().fold(map, |mut acc, bit_field_id| {
293        let bit_field = bit_fields
294            .find(*bit_field_id)
295            .unwrap_or_else(|| panic!("Bitfield id {:?} not found.", bit_field_id));
296        let indexes: Vec<usize> = bit_field
297            .bits()
298            .iter()
299            // The first two bits must be ignored.
300            .skip(2)
301            .enumerate()
302            .filter(|&(ref i, &x)| *i < num_days && x == 1)
303            .map(|(i, _)| i)
304            .collect();
305
306        indexes.iter().for_each(|&i| {
307            acc.entry(dates[i]).or_default().insert(bit_field.id());
308        });
309
310        acc
311    });
312    Ok(result)
313}
314
315fn create_bit_fields_by_stop_id(
316    journeys: &ResourceStorage<Journey>,
317) -> FxHashMap<i32, FxHashSet<i32>> {
318    journeys
319        .entries()
320        .into_iter()
321        .fold(FxHashMap::default(), |mut acc, journey| {
322            journey.route().iter().for_each(|route_entry| {
323                acc.entry(route_entry.stop_id())
324                    .or_default()
325                    // If the journey has no bit_field_id, the default value is 0. A value of 0 means that the journey operates every day.
326                    .insert(journey.bit_field_id().unwrap_or(0));
327            });
328            acc
329        })
330}
331
332fn create_journeys_by_stop_id_and_bit_field_id(
333    journeys: &ResourceStorage<Journey>,
334) -> FxHashMap<(i32, i32), Vec<i32>> {
335    journeys
336        .entries()
337        .into_iter()
338        .fold(FxHashMap::default(), |mut acc, journey| {
339            journey.route().iter().for_each(|route_entry| {
340                // If the journey has no bit_field_id, the default value is 0. A value of 0 means that the journey operates every day.
341                acc.entry((route_entry.stop_id(), journey.bit_field_id().unwrap_or(0)))
342                    .or_default()
343                    .push(journey.id());
344            });
345            acc
346        })
347}
348
349fn create_stop_connections_by_stop_id(
350    stop_connections: &ResourceStorage<StopConnection>,
351) -> FxHashMap<i32, FxHashSet<i32>> {
352    stop_connections
353        .entries()
354        .into_iter()
355        .fold(FxHashMap::default(), |mut acc, stop_connection| {
356            acc.entry(stop_connection.stop_id_1())
357                .or_default()
358                .insert(stop_connection.id());
359            acc
360        })
361}
362
363fn create_exchange_times_journey_map(
364    exchange_times_journey: &ResourceStorage<ExchangeTimeJourney>,
365) -> FxHashMap<(i32, i32, i32), FxHashSet<i32>> {
366    exchange_times_journey.entries().into_iter().fold(
367        FxHashMap::default(),
368        |mut acc, exchange_time| {
369            let key = (
370                exchange_time.stop_id(),
371                exchange_time.journey_id_1(),
372                exchange_time.journey_id_2(),
373            );
374
375            acc.entry(key).or_default().insert(exchange_time.id());
376            acc
377        },
378    )
379}
380
381fn create_exchange_times_administration_map(
382    exchange_times_administration: &ResourceStorage<ExchangeTimeAdministration>,
383) -> FxHashMap<(Option<i32>, String, String), i32> {
384    exchange_times_administration.entries().into_iter().fold(
385        FxHashMap::default(),
386        |mut acc, exchange_time| {
387            let key = (
388                exchange_time.stop_id(),
389                exchange_time.administration_1().into(),
390                exchange_time.administration_2().into(),
391            );
392
393            acc.insert(key, exchange_time.id());
394            acc
395        },
396    )
397}