refactoring a WizardState to replace LOTS of boilerplate. converting

some easy first cases

editor/src/abtest/score.rs

@@ -1,9 +1,8 @@
-use crate::game::{State, Transition};
+use crate::game::{State, Transition, WizardState};
 use crate::ui::PerMapUI;
 use crate::ui::UI;
 use ezgui::{
     hotkey, EventCtx, GfxCtx, HorizontalAlignment, Key, ModalMenu, Text, VerticalAlignment, Wizard,
-    WrappedWizard,
 };
 use geom::Duration;
 use itertools::Itertools;
@@ -83,19 +82,13 @@ impl Scoreboard {
 }
 
 impl State for Scoreboard {
-    fn event(&mut self, ctx: &mut EventCtx, ui: &mut UI) -> Transition {
+    fn event(&mut self, ctx: &mut EventCtx, _: &mut UI) -> Transition {
         self.menu.handle_event(ctx, None);
         if self.menu.action("quit") {
             return Transition::Pop;
         }
         if self.menu.action("browse trips") {
-            return Transition::Push(Box::new(BrowseTrips {
+            return Transition::Push(WizardState::new(Box::new(browse_trips)));
-                trips: CompareTrips::new(
-                    ui.primary.sim.get_finished_trips(),
-                    ui.secondary.as_ref().unwrap().sim.get_finished_trips(),
-                ),
-                wizard: Wizard::new(),
-            }));
         }
         Transition::Keep
     }
@@ -109,28 +102,8 @@ impl State for Scoreboard {
     }
 }
 
-struct BrowseTrips {
+fn browse_trips(wiz: &mut Wizard, ctx: &mut EventCtx, ui: &mut UI) -> Option<Transition> {
-    trips: CompareTrips,
+    let mut wizard = wiz.wrap(ctx);
-    wizard: Wizard,
-}
-
-impl State for BrowseTrips {
-    fn event(&mut self, ctx: &mut EventCtx, _: &mut UI) -> Transition {
-        if pick_trip(&self.trips, &mut self.wizard.wrap(ctx)).is_some() {
-            // TODO show more details...
-            return Transition::Pop;
-        } else if self.wizard.aborted() {
-            return Transition::Pop;
-        }
-        Transition::Keep
-    }
-
-    fn draw(&self, g: &mut GfxCtx, _: &UI) {
-        self.wizard.draw(g);
-    }
-}
-
-fn pick_trip(trips: &CompareTrips, wizard: &mut WrappedWizard) -> Option<TripID> {
     let mode = wizard
         .choose_something_no_keys::<TripMode>(
             "Browse which trips?",
@@ -145,6 +118,10 @@ fn pick_trip(trips: &CompareTrips, wizard: &mut WrappedWizard) -> Option<TripID>
         )?
         .1;
     // TODO Ewwww. Can't do this inside choices_generator because trips isn't &'a static.
+    let trips = CompareTrips::new(
+        ui.primary.sim.get_finished_trips(),
+        ui.secondary.as_ref().unwrap().sim.get_finished_trips(),
+    );
     let mut filtered: Vec<&(TripID, TripMode, Duration, Duration)> = trips
         .finished_trips
         .iter()
@@ -156,12 +133,12 @@ fn pick_trip(trips: &CompareTrips, wizard: &mut WrappedWizard) -> Option<TripID>
         .into_iter()
         .map(|(id, _, t1, t2)| (format!("{} taking {} vs {}", id, t1, t2), *id))
         .collect();
-    wizard
+    wizard.choose_something_no_keys::<TripID>(
-        .choose_something_no_keys::<TripID>(
+        "Examine which trip?",
-            "Examine which trip?",
+        Box::new(move || choices.clone()),
-            Box::new(move || choices.clone()),
+    )?;
-        )
+    // TODO show more details...
-        .map(|(_, id)| id)
+    Some(Transition::Pop)
 }
 
 pub struct CompareTrips {

editor/src/edit/mod.rs

@@ -3,7 +3,7 @@ mod traffic_signals;
 
 use crate::common::CommonState;
 use crate::debug::DebugMode;
-use crate::game::{State, Transition};
+use crate::game::{State, Transition, WizardState};
 use crate::helpers::{ColorScheme, ID};
 use crate::render::{
     DrawCtx, DrawIntersection, DrawLane, DrawMap, DrawOptions, DrawTurn, Renderable,
@@ -99,13 +99,9 @@ impl State for EditMode {
 
         // TODO Only if current edits are unsaved
         if self.menu.action("save edits") {
-            return Transition::Push(Box::new(Saving {
+            return Transition::Push(WizardState::new(Box::new(save_edits)));
-                wizard: Wizard::new(),
-            }));
         } else if self.menu.action("load different edits") {
-            return Transition::Push(Box::new(Loading {
+            return Transition::Push(WizardState::new(Box::new(load_edits)));
-                wizard: Wizard::new(),
-            }));
         }
 
         if let Some(ID::Lane(id)) = ui.primary.current_selection {
@@ -292,79 +288,10 @@ impl State for EditMode {
     }
 }
 
-struct Saving {
+fn save_edits(wiz: &mut Wizard, ctx: &mut EventCtx, ui: &mut UI) -> Option<Transition> {
-    wizard: Wizard,
+    let map = &mut ui.primary.map;
-}
+    let mut wizard = wiz.wrap(ctx);
 
-impl State for Saving {
-    fn event(&mut self, ctx: &mut EventCtx, ui: &mut UI) -> Transition {
-        ctx.canvas.handle_event(ctx.input);
-        if save_edits(self.wizard.wrap(ctx), &mut ui.primary.map).is_some() || self.wizard.aborted()
-        {
-            Transition::Pop
-        } else {
-            Transition::Keep
-        }
-    }
-
-    fn draw(&self, g: &mut GfxCtx, _: &UI) {
-        // TODO Still draw the diffs, yo
-        self.wizard.draw(g);
-    }
-}
-
-struct Loading {
-    wizard: Wizard,
-}
-
-impl State for Loading {
-    fn event(&mut self, ctx: &mut EventCtx, ui: &mut UI) -> Transition {
-        ctx.canvas.handle_event(ctx.input);
-        if let Some(new_edits) = load_edits(
-            &ui.primary.map,
-            &mut self.wizard.wrap(ctx),
-            "Load which map edits?",
-        ) {
-            apply_map_edits(&mut ui.primary, &ui.cs, ctx, new_edits);
-            Transition::Pop
-        } else if self.wizard.aborted() {
-            Transition::Pop
-        } else {
-            Transition::Keep
-        }
-    }
-
-    fn draw(&self, g: &mut GfxCtx, _: &UI) {
-        // TODO Still draw the diffs, yo
-        self.wizard.draw(g);
-    }
-}
-
-struct BulkEditLanes {
-    road: RoadID,
-    wizard: Wizard,
-}
-
-impl State for BulkEditLanes {
-    fn event(&mut self, ctx: &mut EventCtx, ui: &mut UI) -> Transition {
-        ctx.canvas.handle_event(ctx.input);
-        if let Some(edits) = bulk_edit(self.road, &mut self.wizard.wrap(ctx), &ui.primary.map) {
-            apply_map_edits(&mut ui.primary, &ui.cs, ctx, edits);
-            Transition::Pop
-        } else if self.wizard.aborted() {
-            Transition::Pop
-        } else {
-            Transition::Keep
-        }
-    }
-
-    fn draw(&self, g: &mut GfxCtx, _: &UI) {
-        // TODO Still draw the diffs, yo
-        self.wizard.draw(g);
-    }
-}
-
-fn save_edits(mut wizard: WrappedWizard, map: &mut Map) -> Option<()> {
     let rename = if map.get_edits().edits_name == "no_edits" {
         Some(wizard.input_string("Name these map edits")?)
     } else {
@@ -386,7 +313,25 @@ fn save_edits(mut wizard: WrappedWizard, map: &mut Map) -> Option<()> {
         }
         map.get_edits().save();
     }
-    Some(())
+    Some(Transition::Pop)
+}
+
+fn load_edits(wiz: &mut Wizard, ctx: &mut EventCtx, ui: &mut UI) -> Option<Transition> {
+    let map = &mut ui.primary.map;
+    let mut wizard = wiz.wrap(ctx);
+
+    // TODO Exclude current
+    let map_name = map.get_name().to_string();
+    let (_, new_edits) = wizard.choose_something_no_keys::<MapEdits>(
+        "Load which map edits?",
+        Box::new(move || {
+            let mut list = abstutil::load_all_objects("edits", &map_name);
+            list.push(("no_edits".to_string(), MapEdits::new(map_name.clone())));
+            list
+        }),
+    )?;
+    apply_map_edits(&mut ui.primary, &ui.cs, ctx, new_edits);
+    Some(Transition::Pop)
 }
 
 // For lane editing
@@ -532,19 +477,28 @@ pub fn apply_map_edits(
     bundle.map.simplify_edits(&mut timer);
 }
 
-fn load_edits(map: &Map, wizard: &mut WrappedWizard, query: &str) -> Option<MapEdits> {
+struct BulkEditLanes {
-    // TODO Exclude current?
+    road: RoadID,
-    let map_name = map.get_name().to_string();
+    wizard: Wizard,
-    wizard
+}
-        .choose_something_no_keys::<MapEdits>(
+
-            query,
+impl State for BulkEditLanes {
-            Box::new(move || {
+    fn event(&mut self, ctx: &mut EventCtx, ui: &mut UI) -> Transition {
-                let mut list = abstutil::load_all_objects("edits", &map_name);
+        ctx.canvas.handle_event(ctx.input);
-                list.push(("no_edits".to_string(), MapEdits::new(map_name.clone())));
+        if let Some(edits) = bulk_edit(self.road, &mut self.wizard.wrap(ctx), &ui.primary.map) {
-                list
+            apply_map_edits(&mut ui.primary, &ui.cs, ctx, edits);
-            }),
+            Transition::Pop
-        )
+        } else if self.wizard.aborted() {
-        .map(|(_, e)| e)
+            Transition::Pop
+        } else {
+            Transition::Keep
+        }
+    }
+
+    fn draw(&self, g: &mut GfxCtx, _: &UI) {
+        // TODO Still draw the diffs, yo
+        self.wizard.draw(g);
+    }
 }
 
 fn bulk_edit(r: RoadID, wizard: &mut WrappedWizard, map: &Map) -> Option<MapEdits> {

editor/src/game.rs

@@ -2,7 +2,7 @@ use crate::render::DrawOptions;
 use crate::sandbox::SandboxMode;
 use crate::splash_screen::SplashScreen;
 use crate::ui::{Flags, ShowEverything, UI};
-use ezgui::{Canvas, EventCtx, EventLoopMode, GfxCtx, GUI};
+use ezgui::{Canvas, EventCtx, EventLoopMode, GfxCtx, Wizard, GUI};
 
 // This is the top-level of the GUI logic. This module should just manage interactions between the
 // top-level game states.
@@ -195,3 +195,38 @@ pub enum Transition {
     PushWithMode(Box<State>, EventLoopMode),
     ReplaceWithMode(Box<State>, EventLoopMode),
 }
+
+// TODO Maybe let callers stash expensive data computed once here, and let the cb borrow it.
+// Use cases: both BrowseTrips's
+pub struct WizardState {
+    wizard: Wizard,
+    // Returning None means stay in this WizardState
+    cb: Box<Fn(&mut Wizard, &mut EventCtx, &mut UI) -> Option<Transition>>,
+}
+
+impl WizardState {
+    pub fn new(
+        cb: Box<Fn(&mut Wizard, &mut EventCtx, &mut UI) -> Option<Transition>>,
+    ) -> Box<State> {
+        Box::new(WizardState {
+            wizard: Wizard::new(),
+            cb,
+        })
+    }
+}
+
+impl State for WizardState {
+    fn event(&mut self, ctx: &mut EventCtx, ui: &mut UI) -> Transition {
+        ctx.canvas.handle_event(ctx.input);
+        if let Some(t) = (self.cb)(&mut self.wizard, ctx, ui) {
+            return t;
+        } else if self.wizard.aborted() {
+            return Transition::Pop;
+        }
+        Transition::Keep
+    }
+
+    fn draw(&self, g: &mut GfxCtx, _: &UI) {
+        self.wizard.draw(g);
+    }
+}

editor/src/sandbox/mod.rs

@@ -7,11 +7,9 @@ use crate::common::{
 };
 use crate::debug::DebugMode;
 use crate::edit::EditMode;
-use crate::game::{State, Transition};
+use crate::game::{State, Transition, WizardState};
 use crate::ui::{ShowEverything, UI};
-use ezgui::{
+use ezgui::{hotkey, lctrl, EventCtx, EventLoopMode, GfxCtx, Key, ModalMenu, Text, Wizard};
-    hotkey, lctrl, EventCtx, EventLoopMode, GfxCtx, Key, ModalMenu, Text, Wizard, WrappedWizard,
-};
 use geom::Duration;
 use sim::Sim;
 
@@ -175,10 +173,7 @@ impl State for SandboxMode {
                 });
             }
             if self.menu.action("pick a savestate to load") {
-                return Transition::Push(Box::new(LoadSavestate {
+                return Transition::Push(WizardState::new(Box::new(load_savestate)));
-                    path: ui.primary.sim.save_dir(),
-                    wizard: Wizard::new(),
-                }));
             }
 
             if let Some(t) = time_controls(ctx, ui, &mut self.menu) {
@@ -213,41 +208,22 @@ impl State for SandboxMode {
     }
 }
 
-struct LoadSavestate {
+fn load_savestate(wiz: &mut Wizard, ctx: &mut EventCtx, ui: &mut UI) -> Option<Transition> {
-    path: String,
+    let path = ui.primary.sim.save_dir();
-    wizard: Wizard,
-}
 
-impl State for LoadSavestate {
+    let (ss, _) = wiz.wrap(ctx).choose_something_no_keys::<()>(
-    fn event(&mut self, ctx: &mut EventCtx, ui: &mut UI) -> Transition {
+        "Load which savestate?",
-        if let Some(ss) = pick_savestate(&self.path, &mut self.wizard.wrap(ctx)) {
+        Box::new(move || {
-            ctx.loading_screen("load savestate", |ctx, mut timer| {
+            abstutil::list_dir(std::path::Path::new(&path))
-                ui.primary.sim = Sim::load_savestate(ss, &mut timer).unwrap();
+                .into_iter()
-                ui.recalculate_current_selection(ctx);
+                .map(|f| (f, ()))
-            });
+                .collect()
-            return Transition::Pop;
+        }),
-        } else if self.wizard.aborted() {
+    )?;
-            return Transition::Pop;
-        }
-        Transition::Keep
-    }
 
-    fn draw(&self, g: &mut GfxCtx, _: &UI) {
+    ctx.loading_screen("load savestate", |ctx, mut timer| {
-        self.wizard.draw(g);
+        ui.primary.sim = Sim::load_savestate(ss, &mut timer).expect("Can't load savestate");
-    }
+        ui.recalculate_current_selection(ctx);
-}
+    });
-
+    Some(Transition::Pop)
-fn pick_savestate(path: &str, wizard: &mut WrappedWizard) -> Option<String> {
-    let path_copy = path.to_string();
-    wizard
-        .choose_something_no_keys::<()>(
-            "Load which savestate?",
-            Box::new(move || {
-                abstutil::list_dir(std::path::Path::new(&path_copy))
-                    .into_iter()
-                    .map(|f| (f, ()))
-                    .collect()
-            }),
-        )
-        .map(|(f, _)| f)
 }

editor/src/sandbox/score.rs

@@ -1,12 +1,11 @@
-use crate::game::{State, Transition};
+use crate::game::{State, Transition, WizardState};
 use crate::ui::UI;
 use ezgui::{
     hotkey, EventCtx, GfxCtx, HorizontalAlignment, Key, ModalMenu, Text, VerticalAlignment, Wizard,
-    WrappedWizard,
 };
 use geom::{Duration, DurationHistogram};
 use itertools::Itertools;
-use sim::{FinishedTrips, TripID, TripMode};
+use sim::{TripID, TripMode};
 
 pub struct Scoreboard {
     menu: ModalMenu,
@@ -47,16 +46,13 @@ impl Scoreboard {
 }
 
 impl State for Scoreboard {
-    fn event(&mut self, ctx: &mut EventCtx, ui: &mut UI) -> Transition {
+    fn event(&mut self, ctx: &mut EventCtx, _: &mut UI) -> Transition {
         self.menu.handle_event(ctx, None);
         if self.menu.action("quit") {
             return Transition::Pop;
         }
         if self.menu.action("browse trips") {
-            return Transition::Push(Box::new(BrowseTrips {
+            return Transition::Push(WizardState::new(Box::new(browse_trips)));
-                trips: ui.primary.sim.get_finished_trips(),
-                wizard: Wizard::new(),
-            }));
         }
         Transition::Keep
     }
@@ -70,42 +66,21 @@ impl State for Scoreboard {
     }
 }
 
-struct BrowseTrips {
+fn browse_trips(wiz: &mut Wizard, ctx: &mut EventCtx, ui: &mut UI) -> Option<Transition> {
-    trips: FinishedTrips,
+    let mut wizard = wiz.wrap(ctx);
-    wizard: Wizard,
+    let (_, mode) = wizard.choose_something_no_keys::<TripMode>(
-}
+        "Browse which trips?",
-
+        Box::new(|| {
-impl State for BrowseTrips {
+            vec![
-    fn event(&mut self, ctx: &mut EventCtx, _: &mut UI) -> Transition {
+                ("walk".to_string(), TripMode::Walk),
-        if pick_trip(&self.trips, &mut self.wizard.wrap(ctx)).is_some() {
+                ("bike".to_string(), TripMode::Bike),
-            // TODO show trip departure, where it started and ended
+                ("transit".to_string(), TripMode::Transit),
-            return Transition::Pop;
+                ("drive".to_string(), TripMode::Drive),
-        } else if self.wizard.aborted() {
+            ]
-            return Transition::Pop;
+        }),
-        }
+    )?;
-        Transition::Keep
-    }
-
-    fn draw(&self, g: &mut GfxCtx, _: &UI) {
-        self.wizard.draw(g);
-    }
-}
-
-fn pick_trip(trips: &FinishedTrips, wizard: &mut WrappedWizard) -> Option<TripID> {
-    let mode = wizard
-        .choose_something_no_keys::<TripMode>(
-            "Browse which trips?",
-            Box::new(|| {
-                vec![
-                    ("walk".to_string(), TripMode::Walk),
-                    ("bike".to_string(), TripMode::Bike),
-                    ("transit".to_string(), TripMode::Transit),
-                    ("drive".to_string(), TripMode::Drive),
-                ]
-            }),
-        )?
-        .1;
     // TODO Ewwww. Can't do this inside choices_generator because trips isn't &'a static.
+    let trips = ui.primary.sim.get_finished_trips();
     let mut filtered: Vec<&(TripID, TripMode, Duration)> = trips
         .finished_trips
         .iter()
@@ -118,10 +93,10 @@ fn pick_trip(trips: &FinishedTrips, wizard: &mut WrappedWizard) -> Option<TripID
         // TODO Show percentile for time
         .map(|(id, _, dt)| (format!("{} taking {}", id, dt), *id))
         .collect();
-    wizard
+    wizard.choose_something_no_keys::<TripID>(
-        .choose_something_no_keys::<TripID>(
+        "Examine which trip?",
-            "Examine which trip?",
+        Box::new(move || choices.clone()),
-            Box::new(move || choices.clone()),
+    )?;
-        )
+    // TODO show trip departure, where it started and ended
-        .map(|(_, id)| id)
+    Some(Transition::Pop)
 }

editor/src/sandbox/spawner.rs

@@ -1,14 +1,13 @@
 use crate::common::CommonState;
-use crate::game::{State, Transition};
+use crate::game::{State, Transition, WizardState};
 use crate::helpers::ID;
 use crate::render::DrawOptions;
 use crate::ui::{ShowEverything, UI};
 use abstutil::Timer;
-use ezgui::{hotkey, EventCtx, GfxCtx, Key, ModalMenu, Wizard, WrappedWizard};
+use ezgui::{hotkey, EventCtx, GfxCtx, Key, ModalMenu, Wizard};
 use geom::{Duration, PolyLine};
 use map_model::{
-    BuildingID, IntersectionID, IntersectionType, LaneType, Map, PathRequest, Position,
+    BuildingID, IntersectionID, IntersectionType, LaneType, PathRequest, Position, LANE_THICKNESS,
-    LANE_THICKNESS,
 };
 use rand::seq::SliceRandom;
 use rand::Rng;
@@ -94,9 +93,7 @@ impl AgentSpawner {
             }
             None => {
                 if ui.primary.sim.is_empty() && sandbox_menu.action("start a scenario") {
-                    return Some(Box::new(InstantiateScenario {
+                    return Some(WizardState::new(Box::new(instantiate_scenario)));
-                        wizard: Wizard::new(),
-                    }));
                 }
             }
             _ => {}
@@ -358,51 +355,17 @@ fn spawn_agents_around(i: IntersectionID, ui: &mut UI, ctx: &EventCtx) {
     ui.recalculate_current_selection(ctx);
 }
 
-// TODO Dedupe with code from mission/mod.rs.
+fn instantiate_scenario(wiz: &mut Wizard, ctx: &mut EventCtx, ui: &mut UI) -> Option<Transition> {
-struct InstantiateScenario {
+    let num_agents = ui.primary.current_flags.num_agents;
-    wizard: Wizard,
-}
-
-impl State for InstantiateScenario {
-    fn event(&mut self, ctx: &mut EventCtx, ui: &mut UI) -> Transition {
-        if let Some(scenario) = pick_scenario(
-            ui.primary.current_flags.num_agents,
-            &ui.primary.map,
-            &mut self.wizard.wrap(ctx),
-        ) {
-            ctx.loading_screen("instantiate scenario", |_, timer| {
-                scenario.instantiate(
-                    &mut ui.primary.sim,
-                    &ui.primary.map,
-                    &mut ui.primary.current_flags.sim_flags.make_rng(),
-                    timer,
-                );
-                ui.primary.sim.step(&ui.primary.map, SMALL_DT);
-            });
-            return Transition::Pop;
-        } else if self.wizard.aborted() {
-            return Transition::Pop;
-        }
-        Transition::Keep
-    }
-
-    fn draw(&self, g: &mut GfxCtx, _: &UI) {
-        self.wizard.draw(g);
-    }
-}
-
-fn pick_scenario(
-    num_agents: Option<usize>,
-    map: &Map,
-    wizard: &mut WrappedWizard,
-) -> Option<Scenario> {
     let builtin = if let Some(n) = num_agents {
         format!("random scenario with {} agents", n)
     } else {
         "random scenario with some agents".to_string()
     };
+    let map = &ui.primary.map;
     let map_name = map.get_name().to_string();
-    let (_, scenario_name) = wizard.choose_something_no_keys::<String>(
+
+    let (_, scenario_name) = wiz.wrap(ctx).choose_something_no_keys::<String>(
         "Instantiate which scenario?",
         Box::new(move || {
             let mut list = vec![
@@ -413,7 +376,8 @@ fn pick_scenario(
             list
         }),
     )?;
-    Some(if scenario_name == "builtin" {
+
+    let scenario = if scenario_name == "builtin" {
         if let Some(n) = num_agents {
             Scenario::scaled_run(map, n)
         } else {
@@ -427,5 +391,15 @@ fn pick_scenario(
             &mut Timer::throwaway(),
         )
         .unwrap()
-    })
+    };
+    ctx.loading_screen("instantiate scenario", |_, timer| {
+        scenario.instantiate(
+            &mut ui.primary.sim,
+            &ui.primary.map,
+            &mut ui.primary.current_flags.sim_flags.make_rng(),
+            timer,
+        );
+        ui.primary.sim.step(&ui.primary.map, SMALL_DT);
+    });
+    Some(Transition::Pop)
 }