save and restore savestates for entire a/b tests. launching from scratch

takes 33s in debug, this takes 19s

.gitignore

@@ -4,6 +4,7 @@
 
 editor_state.json
 
+data/ab_test_saves/*
 data/ab_tests/*
 data/edits/*
 data/input/*

abstutil/src/io.rs

@@ -303,7 +303,7 @@ pub fn find_next_file(orig: String) -> Option<String> {
     files.into_iter().find(|f| *f > orig)
 }
 
-fn list_dir(dir: &std::path::Path) -> Vec<String> {
+pub fn list_dir(dir: &std::path::Path) -> Vec<String> {
     let mut files: Vec<String> = Vec::new();
     match std::fs::read_dir(dir) {
         Ok(iter) => {

abstutil/src/lib.rs

@@ -12,7 +12,7 @@ pub use crate::collections::{contains_duplicates, retain_btreemap, wraparound_ge
 pub use crate::error::Error;
 pub use crate::io::{
     basename, deserialize_btreemap, deserialize_multimap, find_next_file, find_prev_file,
-    list_all_objects, load_all_objects, read_binary, read_json, save_binary_object,
+    list_all_objects, list_dir, load_all_objects, read_binary, read_json, save_binary_object,
     save_json_object, serialize_btreemap, serialize_multimap, to_json, write_binary, write_json,
     FileWithProgress,
 };

editor/src/abtest/mod.rs

@@ -9,8 +9,9 @@ use ezgui::{
     hotkey, Color, EventCtx, EventLoopMode, GeomBatch, GfxCtx, Key, ModalMenu, Text, Wizard,
 };
 use geom::{Circle, Distance, Duration, Line, PolyLine};
-use map_model::LANE_THICKNESS;
-use sim::TripID;
+use map_model::{Map, LANE_THICKNESS};
+use serde_derive::{Deserialize, Serialize};
+use sim::{Sim, TripID};
 
 pub struct ABTestMode {
     menu: ModalMenu,
@@ -22,6 +23,7 @@ pub struct ABTestMode {
     diff_all: Option<DiffAllTrips>,
     // TODO Not present in Setup state.
     common: CommonState,
+    test_name: String,
 }
 
 pub enum State {
@@ -31,7 +33,7 @@ pub enum State {
 }
 
 impl ABTestMode {
-    pub fn new(ctx: &mut EventCtx, ui: &mut UI) -> ABTestMode {
+    pub fn new(ctx: &mut EventCtx, ui: &mut UI, test_name: &str) -> ABTestMode {
         ui.primary.current_selection = None;
 
         ABTestMode {
@@ -48,6 +50,7 @@ impl ABTestMode {
                         (hotkey(Key::D), "diff all trips"),
                         (hotkey(Key::B), "stop diffing trips"),
                         (hotkey(Key::Q), "scoreboard"),
+                        (hotkey(Key::O), "save state"),
                     ],
                     CommonState::modal_menu_entries(),
                 ]
@@ -60,6 +63,7 @@ impl ABTestMode {
             diff_trip: None,
             diff_all: None,
             common: CommonState::new(),
+            test_name: test_name.to_string(),
         }
     }
 
@@ -133,6 +137,10 @@ impl ABTestMode {
                             return EventLoopMode::InputOnly;
                         }
 
+                        if mode.menu.action("save state") {
+                            mode.savestate(&mut state.ui.primary);
+                        }
+
                         if mode.diff_trip.is_some() {
                             if mode.menu.action("stop diffing trips") {
                                 mode.diff_trip = None;
@@ -263,6 +271,44 @@ impl ABTestMode {
             _ => unreachable!(),
         }
     }
+
+    fn savestate(&mut self, primary: &mut PerMapUI) {
+        // Temporarily move everything into this structure.
+        let blank_map = Map::blank();
+        let mut secondary = self.secondary.take().unwrap();
+        let ss = ABTestSavestate {
+            primary_map: std::mem::replace(&mut primary.map, Map::blank()),
+            primary_sim: std::mem::replace(
+                &mut primary.sim,
+                Sim::new(&blank_map, "run".to_string(), None),
+            ),
+            secondary_map: std::mem::replace(&mut secondary.map, Map::blank()),
+            secondary_sim: std::mem::replace(
+                &mut secondary.sim,
+                Sim::new(&blank_map, "run".to_string(), None),
+            ),
+        };
+
+        let path = format!(
+            "../data/ab_test_saves/{}/{}/{}.bin",
+            ss.primary_map.get_name(),
+            self.test_name,
+            ss.primary_sim.time()
+        );
+        abstutil::write_binary(&path, &ss).unwrap();
+        println!("Saved {}", path);
+
+        // Restore everything.
+        primary.sim = ss.primary_sim;
+        primary.map = ss.primary_map;
+        self.secondary = Some(PerMapUI {
+            map: ss.secondary_map,
+            draw_map: secondary.draw_map,
+            sim: ss.secondary_sim,
+            current_selection: secondary.current_selection,
+            current_flags: secondary.current_flags,
+        });
+    }
 }
 
 pub struct DiffOneTrip {
@@ -365,3 +411,11 @@ impl DiffAllTrips {
         batch.draw(g);
     }
 }
+
+#[derive(Serialize, Deserialize)]
+pub struct ABTestSavestate {
+    primary_map: Map,
+    primary_sim: Sim,
+    secondary_map: Map,
+    secondary_sim: Sim,
+}

editor/src/abtest/setup.rs

@@ -1,6 +1,7 @@
-use crate::abtest::{ABTestMode, State};
+use crate::abtest::{ABTestMode, ABTestSavestate, State};
 use crate::edit::apply_map_edits;
 use crate::game::{GameState, Mode};
+use crate::render::DrawMap;
 use crate::ui::{Flags, PerMapUI, UI};
 use ezgui::{hotkey, EventCtx, GfxCtx, Key, LogScroller, ModalMenu, Wizard, WrappedWizard};
 use geom::Duration;
@@ -11,6 +12,7 @@ use std::path::PathBuf;
 pub enum ABTestSetup {
     Pick(Wizard),
     Manage(ModalMenu, ABTest, LogScroller),
+    LoadSavestate(ABTest, Wizard),
 }
 
 impl ABTestSetup {
@@ -29,6 +31,7 @@ impl ABTestSetup {
                                     vec![
                                         (hotkey(Key::Escape), "quit"),
                                         (hotkey(Key::R), "run A/B test"),
+                                        (hotkey(Key::L), "load savestate"),
                                     ],
                                     ctx,
                                 ),
@@ -39,13 +42,37 @@ impl ABTestSetup {
                             state.mode = Mode::SplashScreen(Wizard::new(), None);
                         }
                     }
+                    ABTestSetup::LoadSavestate(ref test, ref mut wizard) => {
+                        if let Some(ss) = pick_savestate(test, &mut wizard.wrap(ctx)) {
+                            state.mode = launch_savestate(test, ss, &mut state.ui, ctx);
+                        } else if wizard.aborted() {
+                            // TODO Here's where we need to push and pop states.
+                            let scroller =
+                                LogScroller::new(test.test_name.clone(), test.describe());
+                            *setup = ABTestSetup::Manage(
+                                ModalMenu::new(
+                                    &format!("A/B Test Editor for {}", test.test_name),
+                                    vec![
+                                        (hotkey(Key::Escape), "quit"),
+                                        (hotkey(Key::R), "run A/B test"),
+                                        (hotkey(Key::L), "load savestate"),
+                                    ],
+                                    ctx,
+                                ),
+                                test.clone(),
+                                scroller,
+                            );
+                        }
+                    }
                     ABTestSetup::Manage(ref mut menu, test, ref mut scroller) => {
                         ctx.canvas.handle_event(ctx.input);
                         menu.handle_event(ctx, None);
                         if scroller.event(ctx.input) {
                             state.mode = Mode::SplashScreen(Wizard::new(), None);
                         } else if menu.action("run A/B test") {
                             state.mode = launch_test(test, &mut state.ui, ctx);
+                        } else if menu.action("load savestate") {
+                            *setup = ABTestSetup::LoadSavestate(test.clone(), Wizard::new());
                         }
                     }
                 },
@@ -60,6 +87,9 @@ impl ABTestSetup {
             ABTestSetup::Pick(wizard) => {
                 wizard.draw(g);
             }
+            ABTestSetup::LoadSavestate(_, wizard) => {
+                wizard.draw(g);
+            }
             ABTestSetup::Manage(ref menu, _, scroller) => {
                 scroller.draw(g);
                 menu.draw(g);
@@ -158,12 +188,53 @@ fn launch_test(test: &ABTest, ui: &mut UI, ctx: &mut EventCtx) -> Mode {
         },
     );
 
-    let mut mode = ABTestMode::new(ctx, ui);
+    let mut mode = ABTestMode::new(ctx, ui, &test.test_name);
     mode.state = State::Playing;
     mode.secondary = Some(secondary);
     Mode::ABTest(mode)
 }
 
+fn launch_savestate(test: &ABTest, ss_path: String, ui: &mut UI, ctx: &mut EventCtx) -> Mode {
+    ctx.loading_screen(
+        &format!("Launch A/B test from savestate {}", ss_path),
+        |ctx, mut timer| {
+            let ss: ABTestSavestate = abstutil::read_binary(&ss_path, &mut timer).unwrap();
+
+            timer.start("setup primary");
+            ui.primary.map = ss.primary_map;
+            ui.primary.sim = ss.primary_sim;
+            ui.primary.draw_map = DrawMap::new(
+                &ui.primary.map,
+                &ui.primary.current_flags,
+                &ui.cs,
+                ctx.prerender,
+                &mut timer,
+            );
+            timer.stop("setup primary");
+
+            timer.start("setup secondary");
+            let mut mode = ABTestMode::new(ctx, ui, &test.test_name);
+            mode.state = State::Playing;
+            mode.secondary = Some(PerMapUI {
+                draw_map: DrawMap::new(
+                    &ss.secondary_map,
+                    &ui.primary.current_flags,
+                    &ui.cs,
+                    ctx.prerender,
+                    &mut timer,
+                ),
+                map: ss.secondary_map,
+                sim: ss.secondary_sim,
+                current_selection: None,
+                // TODO Hack... can we just remove these?
+                current_flags: ui.primary.current_flags.clone(),
+            });
+            timer.stop("setup secondary");
+            Mode::ABTest(mode)
+        },
+    )
+}
+
 fn choose_scenario(map: &Map, wizard: &mut WrappedWizard, query: &str) -> Option<String> {
     let map_name = map.get_name().to_string();
     wizard
@@ -197,3 +268,21 @@ fn load_ab_test(map: &Map, wizard: &mut WrappedWizard, query: &str) -> Option<AB
         )
         .map(|(_, t)| t)
 }
+
+fn pick_savestate(test: &ABTest, wizard: &mut WrappedWizard) -> Option<String> {
+    let path = format!(
+        "../data/ab_test_saves/{}/{}/",
+        test.map_name, test.test_name
+    );
+    wizard
+        .choose_something_no_keys::<()>(
+            "Load which savestate?",
+            Box::new(move || {
+                abstutil::list_dir(std::path::Path::new(&path))
+                    .into_iter()
+                    .map(|f| (f, ()))
+                    .collect()
+            }),
+        )
+        .map(|(f, _)| f)
+}

editor/src/game.rs

@@ -262,7 +262,9 @@ fn splash_screen(
             x if x == tutorial => break Some(Mode::Tutorial(TutorialMode::new(ctx, ui))),
             x if x == debug => break Some(Mode::Debug(DebugMode::new(ctx, ui))),
             x if x == mission => break Some(Mode::Mission(MissionEditMode::new(ctx, ui))),
-            x if x == abtest => break Some(Mode::ABTest(ABTestMode::new(ctx, ui))),
+            x if x == abtest => {
+                break Some(Mode::ABTest(ABTestMode::new(ctx, ui, "unnamed a/b test")))
+            }
             x if x == about => {
                 if wizard.acknowledge(
                     "About A/B Street",

map_model/src/map.rs

@@ -8,7 +8,7 @@ use crate::{
 };
 use abstutil;
 use abstutil::{deserialize_btreemap, serialize_btreemap, Error, Timer};
-use geom::{Bounds, GPSBounds, Polygon};
+use geom::{Bounds, GPSBounds, Polygon, Pt2D};
 use serde_derive::{Deserialize, Serialize};
 use std::collections::{BTreeMap, BTreeSet, HashSet, VecDeque};
 use std::io;
@@ -55,6 +55,34 @@ impl Map {
         Ok(Map::create_from_raw(abstutil::basename(path), data, timer))
     }
 
+    // Just for temporary std::mem::replace tricks.
+    pub fn blank() -> Map {
+        Map {
+            roads: Vec::new(),
+            lanes: Vec::new(),
+            intersections: Vec::new(),
+            turns: BTreeMap::new(),
+            buildings: Vec::new(),
+            bus_stops: BTreeMap::new(),
+            bus_routes: Vec::new(),
+            areas: Vec::new(),
+            boundary_polygon: Polygon::new(&vec![
+                Pt2D::new(0.0, 0.0),
+                Pt2D::new(1.0, 0.0),
+                Pt2D::new(1.0, 1.0),
+            ]),
+            stop_signs: BTreeMap::new(),
+            traffic_signals: BTreeMap::new(),
+            gps_bounds: GPSBounds::new(),
+            bounds: Bounds::new(),
+            turn_lookup: Vec::new(),
+            pathfinder: None,
+            pathfinder_dirty: false,
+            name: "blank".to_string(),
+            edits: MapEdits::new("blank".to_string()),
+        }
+    }
+
     pub fn create_from_raw(name: String, data: raw_data::Map, timer: &mut Timer) -> Map {
         timer.start("raw_map to InitialMap");
         let gps_bounds = data.gps_bounds.clone();