Refactor stages and context

maplibre/src/context.rs

@@ -63,8 +63,4 @@ pub struct MapContext {
 
     pub tile_cache: TileCache,
     pub renderer: Renderer,
-    pub scheduler: Box<dyn ScheduleMethod>,
-
-    pub message_receiver: mpsc::Receiver<TessellateMessage>,
-    pub shared_thread_state: SharedThreadState,
 }

maplibre/src/map_schedule.rs

@@ -26,10 +26,6 @@ pub struct PrematureMapContext {
     style: Style,
 
     tile_cache: TileCache,
-    scheduler: Box<dyn ScheduleMethod>,
-
-    message_receiver: mpsc::Receiver<TessellateMessage>,
-    shared_thread_state: SharedThreadState,
 
     wgpu_settings: WgpuSettings,
     renderer_settings: RendererSettings,
@@ -38,12 +34,12 @@ pub struct PrematureMapContext {
 pub enum EventuallyMapContext {
     Full(MapContext),
     Premature(PrematureMapContext),
-    Empty,
+    _Uninitialized,
 }
 
 impl EventuallyMapContext {
     pub fn make_full(&mut self, renderer: Renderer) {
-        let context = mem::replace(self, EventuallyMapContext::Empty);
+        let context = mem::replace(self, EventuallyMapContext::_Uninitialized);
 
         match context {
             EventuallyMapContext::Full(_) => {}
@@ -51,11 +47,7 @@ impl EventuallyMapContext {
                 view_state,
                 style,
                 tile_cache,
-                scheduler,
+                ..
-                message_receiver,
-                shared_thread_state,
-                wgpu_settings,
-                renderer_settings,
             }) => {
                 mem::replace(
                     self,
@@ -64,13 +56,10 @@ impl EventuallyMapContext {
                         style,
                         tile_cache,
                         renderer,
-                        scheduler,
-                        message_receiver,
-                        shared_thread_state,
                     }),
                 );
             }
-            EventuallyMapContext::Empty => {}
+            EventuallyMapContext::_Uninitialized => {}
         }
     }
 }
@@ -84,7 +73,7 @@ where
 {
     map_window_config: MWC,
 
-    map_context: EventuallyMapContext,
+    pub map_context: EventuallyMapContext,
 
     schedule: Schedule,
 
@@ -112,20 +101,13 @@ where
     ) -> Self {
         let view_state = ViewState::new(&window_size);
         let tile_cache = TileCache::new();
-
         let mut schedule = Schedule::default();
-        let client: SourceClient<HC> = SourceClient::Http(HttpSourceClient::new(http_client));
+
-        register_stages(&mut schedule, client);
+        let http_source_client: HttpSourceClient<HC> = HttpSourceClient::new(http_client);
+        //register_stages(&mut schedule, http_source_client, Box::new(scheduler));
+
         register_render_stages(&mut schedule);
 
-        let (message_sender, message_receiver) = mpsc::channel();
-
-        let scheduler = Box::new(scheduler.take());
-        let shared_thread_state = SharedThreadState {
-            tile_request_state: Arc::new(Mutex::new(TileRequestState::new())),
-            message_sender,
-            geometry_index: Arc::new(Mutex::new(GeometryIndex::new())),
-        };
         Self {
             map_window_config,
             map_context: match renderer {
@@ -133,10 +115,7 @@ where
                     view_state,
                     style,
                     tile_cache,
-                    scheduler,
-                    shared_thread_state,
                     wgpu_settings,
-                    message_receiver,
                     renderer_settings,
                 }),
                 Some(renderer) => EventuallyMapContext::Full(MapContext {
@@ -144,9 +123,6 @@ where
                     style,
                     tile_cache,
                     renderer,
-                    scheduler,
-                    shared_thread_state,
-                    message_receiver,
                 }),
             },
             schedule,
@@ -226,7 +202,7 @@ where
                 &self.map_context.make_full(renderer);
                 true
             }
-            EventuallyMapContext::Empty => false,
+            EventuallyMapContext::_Uninitialized => false,
         }
     }
 

maplibre/src/stages/mod.rs

@@ -1,15 +1,42 @@
 //! [Stages](Stage) for requesting and preparing data
 
-use crate::io::source_client::SourceClient;
+use crate::io::geometry_index::GeometryIndex;
+use crate::io::shared_thread_state::SharedThreadState;
+use crate::io::source_client::{HttpSourceClient, SourceClient};
+use crate::io::tile_request_state::TileRequestState;
+use crate::io::TessellateMessage;
 use crate::schedule::Schedule;
 use crate::stages::populate_tile_store_stage::PopulateTileStore;
-use crate::HTTPClient;
+use crate::{HttpClient, ScheduleMethod, Scheduler};
 use request_stage::RequestStage;
+use std::sync::{mpsc, Arc, Mutex};
 
 mod populate_tile_store_stage;
 mod request_stage;
 
-pub fn register_stages<HC: HTTPClient>(schedule: &mut Schedule, source_client: SourceClient<HC>) {
+pub type MessageSender = mpsc::Sender<TessellateMessage>;
-    schedule.add_stage("request", RequestStage::new(source_client));
+pub type MessageReceiver = mpsc::Receiver<TessellateMessage>;
-    schedule.add_stage("populate_tile_store", PopulateTileStore::default());
+
+pub fn register_stages<HC: HttpClient, SM: ScheduleMethod>(
+    schedule: &mut Schedule,
+    http_source_client: HttpSourceClient<HC>,
+    scheduler: Box<Scheduler<SM>>,
+) {
+    let (message_sender, message_receiver): (MessageSender, MessageReceiver) = mpsc::channel();
+    let shared_thread_state = SharedThreadState {
+        tile_request_state: Arc::new(Mutex::new(TileRequestState::new())),
+        message_sender,
+        geometry_index: Arc::new(Mutex::new(GeometryIndex::new())),
+    };
+
+    let scheduler = Box::new(scheduler.take());
+
+    schedule.add_stage(
+        "request",
+        RequestStage::new(shared_thread_state.clone(), http_source_client, scheduler),
+    );
+    schedule.add_stage(
+        "populate_tile_store",
+        PopulateTileStore::new(shared_thread_state, message_receiver),
+    );
 }

maplibre/src/stages/populate_tile_store_stage.rs

@@ -1,23 +1,29 @@
 //! Receives data from async threads and populates the [`crate::io::tile_cache::TileCache`].
 
 use crate::context::MapContext;
+use crate::io::shared_thread_state::SharedThreadState;
 use crate::io::{TessellateMessage, TileTessellateMessage};
 use crate::schedule::Stage;
+use crate::stages::MessageReceiver;
+use std::sync::mpsc;
 
-#[derive(Default)]
+pub struct PopulateTileStore {
-pub struct PopulateTileStore {}
+    shared_thread_state: SharedThreadState,
+    message_receiver: MessageReceiver,
+}
 
-impl Stage for PopulateTileStore {
+impl PopulateTileStore {
-    fn run(
+    pub fn new(shared_thread_state: SharedThreadState, message_receiver: MessageReceiver) -> Self {
-        &mut self,
+        Self {
-        MapContext {
-            tile_cache,
             shared_thread_state,
             message_receiver,
-            ..
+        }
-        }: &mut MapContext,
+    }
-    ) {
+}
-        if let Ok(result) = message_receiver.try_recv() {
+
+impl Stage for PopulateTileStore {
+    fn run(&mut self, MapContext { tile_cache, .. }: &mut MapContext) {
+        if let Ok(result) = self.message_receiver.try_recv() {
             match result {
                 TessellateMessage::Layer(layer_result) => {
                     tracing::trace!(
@@ -29,7 +35,7 @@ impl Stage for PopulateTileStore {
                 }
                 TessellateMessage::Tile(TileTessellateMessage { request_id, coords }) => loop {
                     if let Ok(mut tile_request_state) =
-                        shared_thread_state.tile_request_state.try_lock()
+                        self.shared_thread_state.tile_request_state.try_lock()
                     {
                         tile_request_state.finish_tile_request(request_id);
                         tracing::trace!("Tile at {} finished loading", coords);

maplibre/src/stages/request_stage.rs

@@ -4,28 +4,36 @@ use crate::context::MapContext;
 use crate::coords::{ViewRegion, WorldTileCoords};
 use crate::error::Error;
 use crate::io::shared_thread_state::SharedThreadState;
-use crate::io::source_client::SourceClient;
+use crate::io::source_client::{HttpSourceClient, SourceClient};
 use crate::io::tile_cache::TileCache;
 use crate::io::TileRequest;
 use crate::schedule::Stage;
-use crate::{HTTPClient, ScheduleMethod, Style};
+use crate::{HttpClient, ScheduleMethod, Style};
 use std::collections::HashSet;
 
 pub struct RequestStage<HC>
 where
-    HC: HTTPClient,
+    HC: HttpClient,
 {
-    pub source_client: SourceClient<HC>,
+    shared_thread_state: SharedThreadState,
-    pub try_failed: bool,
+    scheduler: Box<dyn ScheduleMethod>,
+    http_source_client: HttpSourceClient<HC>,
+    try_failed: bool,
 }
 
 impl<HC> RequestStage<HC>
 where
-    HC: HTTPClient,
+    HC: HttpClient,
 {
-    pub fn new(source_client: SourceClient<HC>) -> Self {
+    pub fn new(
+        shared_thread_state: SharedThreadState,
+        http_source_client: HttpSourceClient<HC>,
+        scheduler: Box<dyn ScheduleMethod>,
+    ) -> Self {
         Self {
-            source_client,
+            shared_thread_state,
+            scheduler,
+            http_source_client,
             try_failed: false,
         }
     }
@@ -33,7 +41,7 @@ where
 
 impl<HC> Stage for RequestStage<HC>
 where
-    HC: HTTPClient,
+    HC: HttpClient,
 {
     fn run(
         &mut self,
@@ -41,8 +49,6 @@ where
             view_state,
             style,
             tile_cache,
-            scheduler,
-            shared_thread_state,
             ..
         }: &mut MapContext,
     ) {
@@ -59,13 +65,7 @@ where
         {
             if let Some(view_region) = &view_region {
                 // FIXME: We also need to request tiles from layers above if we are over the maximum zoom level
-                self.try_failed = self.request_tiles_in_view(
+                self.try_failed = self.request_tiles_in_view(tile_cache, style, view_region);
-                    tile_cache,
-                    style,
-                    shared_thread_state,
-                    scheduler,
-                    view_region,
-                );
             }
         }
 
@@ -76,7 +76,7 @@ where
 
 impl<HC> RequestStage<HC>
 where
-    HC: HTTPClient,
+    HC: HttpClient,
 {
     /// Request tiles which are currently in view.
     #[tracing::instrument(skip_all)]
@@ -84,8 +84,6 @@ where
         &self,
         tile_cache: &TileCache,
         style: &Style,
-        shared_thread_state: &SharedThreadState,
-        scheduler: &Box<dyn ScheduleMethod>,
         view_region: &ViewRegion,
     ) -> bool {
         let mut try_failed = false;
@@ -99,13 +97,7 @@ where
             if coords.build_quad_key().is_some() {
                 // TODO: Make tesselation depend on style?
                 try_failed = self
-                    .try_request_tile(
+                    .try_request_tile(tile_cache, &coords, &source_layers)
-                        tile_cache,
-                        shared_thread_state,
-                        scheduler,
-                        &coords,
-                        &source_layers,
-                    )
                     .unwrap();
             }
         }
@@ -115,8 +107,6 @@ where
     fn try_request_tile(
         &self,
         tile_cache: &TileCache,
-        shared_thread_state: &SharedThreadState,
-        scheduler: &Box<dyn ScheduleMethod>,
         coords: &WorldTileCoords,
         layers: &HashSet<String>,
     ) -> Result<bool, Error> {
@@ -124,7 +114,7 @@ where
             return Ok(false);
         }
 
-        if let Ok(mut tile_request_state) = shared_thread_state.tile_request_state.try_lock() {
+        if let Ok(mut tile_request_state) = self.shared_thread_state.tile_request_state.try_lock() {
             if let Some(request_id) = tile_request_state.start_tile_request(TileRequest {
                 coords: *coords,
                 layers: layers.clone(),
@@ -141,12 +131,12 @@ where
                     );
                 }*/
 
-                let client = self.source_client.clone();
+                let client = SourceClient::Http(self.http_source_client.clone());
                 let coords = *coords;
 
-                scheduler
+                self.scheduler
                     .schedule(
-                        shared_thread_state.clone(),
+                        self.shared_thread_state.clone(),
                         Box::new(move |state: SharedThreadState| {
                             Box::pin(async move {
                                 match client.fetch(&coords).await {