Split shape storage (timetable + routed) (#321)

* split shape storage (timetable + routed)

* formatting

* fix assert

Author: pablohoch

include/nigiri/loader/gtfs/shape.h

@@ -14,6 +14,9 @@ using relative_shape_idx_t =
     cista::strong<std::uint32_t, struct relative_shape_idx_>;
 
 struct shape_loader_state {
+  scoped_shape_idx_t get_scoped_shape_idx(relative_shape_idx_t const i) const {
+    return to_scoped_shape_idx(get_shape_idx(i), shape_source::kTimetable);
+  }
   shape_idx_t get_shape_idx(relative_shape_idx_t const i) const {
     return shape_idx_t{to_idx(i) + index_offset_};
   }
@@ -28,4 +31,4 @@ struct shape_loader_state {
 
 shape_loader_state parse_shapes(std::string_view const, shapes_storage&);
 
-}  // namespace nigiri::loader::gtfs
+}  // namespace nigiri::loader::gtfs

include/nigiri/shapes_storage.h

@@ -17,26 +17,25 @@ namespace nigiri {
 struct shapes_storage {
   shapes_storage(std::filesystem::path,
                  cista::mmap::protection,
-                 bool keep_shape_data = false);
+                 bool keep_routed_shape_data = false);
 
   cista::mmap mm(char const* file, bool keep = false);
 
-  std::span<geo::latlng const> get_shape(shape_idx_t) const;
+  std::span<geo::latlng const> get_shape(scoped_shape_idx_t) const;
   std::span<geo::latlng const> get_shape(trip_idx_t) const;
   std::span<geo::latlng const> get_shape(trip_idx_t,
                                          interval<stop_idx_t> const&) const;
 
-  std::pair<std::span<geo::latlng const>, shape_idx_t> get_shape_with_idx(
-      trip_idx_t) const;
-  std::pair<std::span<geo::latlng const>, shape_idx_t> get_shape_with_idx(
-      trip_idx_t, interval<stop_idx_t> const&) const;
+  std::pair<std::span<geo::latlng const>, scoped_shape_idx_t>
+      get_shape_with_idx(trip_idx_t) const;
+  std::pair<std::span<geo::latlng const>, scoped_shape_idx_t>
+  get_shape_with_idx(trip_idx_t, interval<stop_idx_t> const&) const;
 
-  shape_idx_t get_shape_idx(trip_idx_t) const;
-  shape_source get_shape_source(shape_idx_t) const;
+  scoped_shape_idx_t get_shape_idx(trip_idx_t) const;
 
   shape_offset_idx_t add_offsets(std::vector<shape_offset_t> const&);
   void add_trip_shape_offsets(
-      trip_idx_t, cista::pair<shape_idx_t, shape_offset_idx_t> const&);
+      trip_idx_t, cista::pair<scoped_shape_idx_t, shape_offset_idx_t> const&);
 
   geo::box get_bounding_box(route_idx_t) const;
   std::optional<geo::box> get_bounding_box(route_idx_t,
@@ -46,12 +45,12 @@ struct shapes_storage {
   std::filesystem::path p_;
 
   mm_paged_vecvec<shape_idx_t, geo::latlng> data_;
+  mm_paged_vecvec<shape_idx_t, geo::latlng> routed_data_;
   mm_vecvec<shape_offset_idx_t, shape_offset_t, std::uint64_t> offsets_;
-  mm_vec_map<trip_idx_t, cista::pair<shape_idx_t, shape_offset_idx_t>>
+  mm_vec_map<trip_idx_t, cista::pair<scoped_shape_idx_t, shape_offset_idx_t>>
       trip_offset_indices_;
   mm_vec_map<route_idx_t, geo::box> route_bboxes_;
   mm_vecvec<route_idx_t, geo::box, std::uint64_t> route_segment_bboxes_;
-  mm_vec_map<shape_idx_t, shape_source> shape_sources_;
 };
 
 }  // namespace nigiri

include/nigiri/types.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <cassert>
 #include <chrono>
 #include <cinttypes>
 #include <variant>
@@ -169,6 +170,8 @@ using route_idx_t = cista::strong<std::uint32_t, struct _route_idx>;
 using section_idx_t = cista::strong<std::uint32_t, struct _section_idx>;
 using section_db_idx_t = cista::strong<std::uint32_t, struct _section_db_idx>;
 using shape_idx_t = cista::strong<std::uint32_t, struct _shape_idx>;
+using scoped_shape_idx_t =
+    cista::strong<std::uint32_t, struct _scoped_shape_idx>;
 using shape_offset_t = cista::strong<std::uint32_t, struct _shape_offset>;
 using shape_offset_idx_t =
     cista::strong<std::uint32_t, struct _shape_offset_idx>;
@@ -459,6 +462,35 @@ enum class shape_source : std::uint8_t {
   kRouted = 2
 };
 
+constexpr auto const kShapeSourceBit = std::uint32_t{1U} << 31U;
+constexpr auto const kShapeIndexMask = kShapeSourceBit - 1U;
+
+inline shape_source get_shape_source(scoped_shape_idx_t const idx) {
+  if (idx == scoped_shape_idx_t::invalid()) {
+    return shape_source::kNone;
+  }
+  return (to_idx(idx) & kShapeSourceBit) != 0U ? shape_source::kRouted
+                                               : shape_source::kTimetable;
+}
+
+inline shape_idx_t get_local_shape_idx(scoped_shape_idx_t const idx) {
+  return idx == scoped_shape_idx_t::invalid()
+             ? shape_idx_t::invalid()
+             : shape_idx_t{to_idx(idx) & kShapeIndexMask};
+}
+
+inline scoped_shape_idx_t to_scoped_shape_idx(shape_idx_t const local_idx,
+                                              shape_source const source) {
+  if (local_idx == shape_idx_t::invalid() || source == shape_source::kNone) {
+    return scoped_shape_idx_t::invalid();
+  }
+  auto const local = to_idx(local_idx);
+  assert((local & kShapeSourceBit) == 0U);
+  return scoped_shape_idx_t{
+      local |
+      (source == shape_source::kRouted ? kShapeSourceBit : std::uint32_t{0U})};
+}
+
 }  // namespace nigiri
 
 #include <iomanip>

src/loader/gtfs/shape.cc

@@ -44,7 +44,6 @@ shape_loader_state parse_shapes(std::string_view const data,
             shapes.emplace_back_empty();
             states.distances_.emplace_back();
             seq.emplace_back();
-            shapes_data.shape_sources_.emplace_back(shape_source::kTimetable);
             return idx;
           });
           auto polyline = shapes[shape_idx];

src/loader/gtfs/shape_prepare.cc

@@ -164,7 +164,8 @@ void process_task(timetable const& tt,
                   shape_loader_state const& shape_states,
                   relative_shape_idx_t const i,
                   task& t) {
-  auto const shape = shapes_data.get_shape(shape_states.get_shape_idx(i));
+  auto const shape =
+      shapes_data.get_shape(shape_states.get_scoped_shape_idx(i));
   auto const& shape_distances = shape_states.distances_[i];
   for (auto& x : t) {
     auto& r = x.result_;
@@ -228,7 +229,8 @@ void assign_shape_offsets(shapes_storage& shapes_data,
     auto const shape_idx = trip.shape_idx_;
     if (shape_idx == shape_idx_t::invalid()) {
       shapes_data.add_trip_shape_offsets(
-          trip_idx, cista::pair{shape_idx, shape_offset_idx_t::invalid()});
+          trip_idx, cista::pair{scoped_shape_idx_t::invalid(),
+                                shape_offset_idx_t::invalid()});
     } else {
       auto const task = tasks[states.get_relative_idx(shape_idx)];
       auto const x = std::ranges::lower_bound(
@@ -238,11 +240,13 @@ void assign_shape_offsets(shapes_storage& shapes_data,
       if (x != end(task) &&
           x->result_.shape_offset_idx_ != shape_offset_idx_t::invalid()) {
         shapes_data.add_trip_shape_offsets(
-            trip_idx, cista::pair{shape_idx, x->result_.shape_offset_idx_});
+            trip_idx, cista::pair{to_scoped_shape_idx(shape_idx,
+                                                      shape_source::kTimetable),
+                                  x->result_.shape_offset_idx_});
       } else {
         shapes_data.add_trip_shape_offsets(
-            trip_idx,
-            cista::pair{shape_idx_t::invalid(), shape_offset_idx_t::invalid()});
+            trip_idx, cista::pair{scoped_shape_idx_t::invalid(),
+                                  shape_offset_idx_t::invalid()});
       }
     }
   }
@@ -283,12 +287,14 @@ void assign_bounding_boxes(timetable const& tt,
                                          .rt_ = rt_transport_idx_t::invalid()}};
       frun.for_each_trip([&](trip_idx_t const trip_idx,
                              interval<stop_idx_t> const absolute_range) {
-        auto const [shape_idx, offset_idx] =
+        auto const [scoped_shape_idx, offset_idx] =
             shapes_data.trip_offset_indices_[trip_idx];
-        if (shape_idx == shape_idx_t::invalid() ||
+        if (scoped_shape_idx == scoped_shape_idx_t::invalid() ||
+            get_shape_source(scoped_shape_idx) != shape_source::kTimetable ||
             offset_idx == shape_offset_idx_t::invalid()) {
           return;
         }
+        auto const shape_idx = get_local_shape_idx(scoped_shape_idx);
         auto const& task = tasks[shape_states.get_relative_idx(shape_idx)];
         auto const it = utl::find_if(task, [&](stop_seq_dist const& s) {
           return s.result_.shape_offset_idx_ == offset_idx;

src/loader/netex/load_timetable.cc

@@ -1450,8 +1450,8 @@ void load_timetable(loader_config const& config,
       tt.trip_stop_seq_numbers_.add_back_sized(0U);
       if (shapes_data != nullptr) {
         shapes_data->add_trip_shape_offsets(
-            trip_idx,
-            cista::pair{shape_idx_t::invalid(), shape_offset_idx_t::invalid()});
+            trip_idx, cista::pair{scoped_shape_idx_t::invalid(),
+                                  shape_offset_idx_t::invalid()});
       }
 
       auto const all_destinations_equal = utl::all_of(

src/rt/frun.cc

@@ -661,7 +661,7 @@ void frun::for_each_shape_point(
       auto const [shape, shape_idx] = shapes_data->get_shape_with_idx(
           trip_idx, absolute_range << absolute_trip_offset);
       if (!shape.empty()) {
-        current_shape_source = shapes_data->shape_sources_[shape_idx];
+        current_shape_source = get_shape_source(shape_idx);
         return shape;
       }
     }

src/shapes_storage.cc

@@ -17,24 +17,28 @@ namespace nigiri {
 
 shapes_storage::shapes_storage(std::filesystem::path path,
                                cista::mmap::protection const mode,
-                               bool keep_shape_data)
+                               bool const keep_routed_shape_data)
     : mode_{mode},
       p_{[&]() {
         fs::create_directories(path);
         return std::move(path);
       }()},
       data_{mm_paged_vecvec_helper<shape_idx_t, geo::latlng>::data_t{
-                mm_vec<geo::latlng>{mm("shapes_data.bin", keep_shape_data)}},
+                mm_vec<geo::latlng>{mm("shapes_data.bin")}},
             mm_vec<cista::page<std::uint64_t, std::uint32_t>>{
-                mm("shapes_idx.bin", keep_shape_data)}},
+                mm("shapes_idx.bin")}},
+      routed_data_{mm_paged_vecvec_helper<shape_idx_t, geo::latlng>::data_t{
+                       mm_vec<geo::latlng>{mm("routed_shapes_data.bin",
+                                              keep_routed_shape_data)}},
+                   mm_vec<cista::page<std::uint64_t, std::uint32_t>>{
+                       mm("routed_shapes_idx.bin", keep_routed_shape_data)}},
       offsets_{mm_vec<shape_offset_t>{mm("shape_offsets_data.bin")},
                mm_vec<std::uint64_t>{mm("shape_offsets_idx.bin")}},
       trip_offset_indices_{mm("shape_trip_offsets.bin")},
       route_bboxes_{mm("shape_route_bboxes.bin")},
       route_segment_bboxes_{
           mm_vec<geo::box>{mm("shape_route_segment_bboxes_data.bin")},
-          mm_vec<std::uint64_t>{mm("shape_route_segment_bboxes_idx.bin")}},
-      shape_sources_{mm("shape_sources.bin", keep_shape_data)} {}
+          mm_vec<std::uint64_t>{mm("shape_route_segment_bboxes_idx.bin")}} {}
 
 cista::mmap shapes_storage::mm(char const* file, bool const keep) {
   auto const p = (p_ / file);
@@ -45,31 +49,50 @@ cista::mmap shapes_storage::mm(char const* file, bool const keep) {
           : mode_};
 }
 
-std::tuple<std::span<geo::latlng const>, shape_idx_t, shape_offset_idx_t>
+std::tuple<std::span<geo::latlng const>, scoped_shape_idx_t, shape_offset_idx_t>
 get_shape(shapes_storage const& storage, trip_idx_t const trip_idx) {
   if (trip_idx == trip_idx_t::invalid() ||
       trip_idx >= storage.trip_offset_indices_.size()) {
     return {};
   }
   auto const [shape_idx, offset_idx] = storage.trip_offset_indices_[trip_idx];
-  assert((shape_idx == shape_idx_t::invalid()) ==
+  assert((shape_idx == scoped_shape_idx_t::invalid()) ==
          (offset_idx == shape_offset_idx_t::invalid()));
   if (offset_idx == shape_offset_idx_t::invalid()) {
     return {};
   }
   return std::tuple{storage.get_shape(shape_idx), shape_idx, offset_idx};
 }
 
-std::span<geo::latlng const> shapes_storage::get_shape(
-    shape_idx_t const shape_idx) const {
+std::span<geo::latlng const> get_shape_from_data(
+    mm_paged_vecvec<shape_idx_t, geo::latlng> const& data,
+    shape_idx_t const shape_idx) {
   if (shape_idx == shape_idx_t::invalid() ||
-      static_cast<std::size_t>(to_idx(shape_idx)) > data_.size()) {
+      static_cast<std::size_t>(to_idx(shape_idx)) >= data.size()) {
     return {};
   }
-  auto const shape = data_[shape_idx];
+  auto const shape = data[shape_idx];
   return {begin(shape), end(shape)};
 }
 
+std::span<geo::latlng const> shapes_storage::get_shape(
+    scoped_shape_idx_t const shape_idx) const {
+  if (shape_idx == scoped_shape_idx_t::invalid()) {
+    return {};
+  }
+
+  auto const local_shape_idx = get_local_shape_idx(shape_idx);
+  switch (get_shape_source(shape_idx)) {
+    case shape_source::kTimetable:
+      return get_shape_from_data(data_, local_shape_idx);
+    case shape_source::kRouted:
+      return get_shape_from_data(routed_data_, local_shape_idx);
+    case shape_source::kNone: return {};
+  }
+
+  return {};
+}
+
 std::span<geo::latlng const> shapes_storage::get_shape(
     trip_idx_t const trip_idx) const {
   auto const [shape, _shape_idx, _offset_idx] =
@@ -98,24 +121,25 @@ std::span<geo::latlng const> shapes_storage::get_shape(
   return get_subshape(*this, shape, offset_idx, range);
 }
 
-std::pair<std::span<geo::latlng const>, shape_idx_t>
+std::pair<std::span<geo::latlng const>, scoped_shape_idx_t>
 shapes_storage::get_shape_with_idx(trip_idx_t const trip_idx) const {
   auto const [shape, shape_idx, _] = nigiri::get_shape(*this, trip_idx);
   return std::pair{shape, shape_idx};
 }
 
-std::pair<std::span<geo::latlng const>, shape_idx_t>
+std::pair<std::span<geo::latlng const>, scoped_shape_idx_t>
 shapes_storage::get_shape_with_idx(trip_idx_t const trip_idx,
                                    interval<stop_idx_t> const& range) const {
   auto const [shape, shape_idx, offset_idx] =
       nigiri::get_shape(*this, trip_idx);
   if (shape.empty()) {
-    return std::pair{shape, shape_idx_t::invalid()};
+    return std::pair{shape, scoped_shape_idx_t::invalid()};
   }
   return std::pair{get_subshape(*this, shape, offset_idx, range), shape_idx};
 }
 
-shape_idx_t shapes_storage::get_shape_idx(trip_idx_t const trip_idx) const {
+scoped_shape_idx_t shapes_storage::get_shape_idx(
+    trip_idx_t const trip_idx) const {
   if (trip_idx == trip_idx_t::invalid() ||
       trip_idx >= trip_offset_indices_.size()) {
     return {};
@@ -133,7 +157,7 @@ shape_offset_idx_t shapes_storage::add_offsets(
 
 void shapes_storage::add_trip_shape_offsets(
     [[maybe_unused]] trip_idx_t const trip_idx,
-    cista::pair<shape_idx_t, shape_offset_idx_t> const& offset_idx) {
+    cista::pair<scoped_shape_idx_t, shape_offset_idx_t> const& offset_idx) {
   assert(trip_idx == trip_offset_indices_.size());
   trip_offset_indices_.emplace_back(offset_idx);
 }
@@ -155,13 +179,4 @@ std::optional<geo::box> shapes_storage::get_bounding_box(
                                  : std::optional<geo::box>{std::nullopt};
 }
 
-shape_source shapes_storage::get_shape_source(
-    shape_idx_t const shape_idx) const {
-  if (shape_idx == shape_idx_t::invalid() ||
-      static_cast<std::size_t>(to_idx(shape_idx)) >= shape_sources_.size()) {
-    return shape_source::kNone;
-  }
-  return shape_sources_[shape_idx];
-}
-
 }  // namespace nigiri

test/loader/gtfs/shape_test.cc

@@ -43,7 +43,8 @@ TEST(gtfs, shape_get_existing_shape_points) {
                 {51.543652, 7.217830},
                 {51.478609, 7.223275},
             }),
-            shapes_data.get_shape(shapes.at("243")));
+            shapes_data.get_shape(to_scoped_shape_idx(
+                shapes.at("243"), shape_source::kTimetable)));
 
   EXPECT_EQ((geo::polyline{
                 {50.553822, 6.356876},
@@ -54,7 +55,8 @@ TEST(gtfs, shape_get_existing_shape_points) {
                 {50.578249, 6.383394},
                 {50.581956, 6.379866},
             }),
-            shapes_data.get_shape(shapes.at("3105")));
+            shapes_data.get_shape((to_scoped_shape_idx(
+                shapes.at("3105"), shape_source::kTimetable))));
 }
 
 TEST(gtfs, shape_not_ascending_sequence) {
@@ -71,7 +73,8 @@ TEST(gtfs, shape_not_ascending_sequence) {
   auto const& shapes = shape_states.id_map_;
 
   EXPECT_EQ((geo::polyline{{50.636259, 6.473668}, {50.636512, 6.473487}}),
-            shapes_data.get_shape(shapes.at("1")));
+            shapes_data.get_shape((to_scoped_shape_idx(
+                shapes.at("1"), shape_source::kTimetable))));
 }
 
 TEST(gtfs, shape_shuffled_rows) {
@@ -131,7 +134,8 @@ TEST(gtfs, shape_shuffled_rows) {
            }},
       };
   for (auto [id, polyline] : shape_points) {
-    EXPECT_EQ(polyline, shapes_data.get_shape(shapes.at(id)));
+    EXPECT_EQ(polyline, shapes_data.get_shape((to_scoped_shape_idx(
+                            shapes.at(id), shape_source::kTimetable))));
   }
 }
 
@@ -152,5 +156,6 @@ TEST(gtfs, shape_delay_insert_no_ascending_sequence) {
   EXPECT_NE(shape_idx, end(shapes));
   EXPECT_NE(shapes.find("2"), end(shapes));
   EXPECT_EQ((geo::polyline{{50.636259, 6.473668}, {50.636512, 6.473487}}),
-            shapes_data.get_shape(shape_idx->second));
-}
+            shapes_data.get_shape((to_scoped_shape_idx(
+                shape_idx->second, shape_source::kTimetable))));
+}