calculate less

Speed up validation by precalculating some stuff and toing a fast return
wherever possible

battlesnake/src/bin/seed-cracker.rs

@@ -1,12 +1,12 @@
-use std::{fs::File, path::PathBuf, sync::LazyLock};
+use std::{fs::File, path::PathBuf, sync::LazyLock, time::Instant};
 
 use az::{Az, WrappingAs};
 use battlesnake::types::{Coord, Direction, wire::Request};
 use bytemuck::cast_slice;
 use clap::Parser;
-use flamer::flame;
+//use flamer::flame;
 use hashbrown::HashMap;
-use log::{debug, error, info, trace, warn};
+use log::{debug, error, info, trace};
 use memmap2::{Mmap, MmapOptions};
 
 type DynError = Box<dyn std::error::Error>;
@@ -23,7 +23,7 @@ struct Args {
     file: PathBuf,
 }
 
-#[flame]
+//#[flame]
 fn main() -> Result<(), DynError> {
     env_logger::init();
 
@@ -39,38 +39,58 @@ fn main() -> Result<(), DynError> {
     let _ = *CHAIN;
 
     let mut rand = SeedRand::new(0);
+    let prepared = PreparedGame::prepare(&requests);
 
     if let Some(seed) = args.seed {
-        if validate(&mut rand, &requests, seed, true) {
+        info!("Validating using old method");
+        let start = Instant::now();
+        let valid = validate(&mut rand, &requests, seed, true);
+        let elapsed = start.elapsed();
+        info!("Validation took {elapsed:?}");
+        if valid {
+            println!("{seed} is the correct seed for this game");
+        } else {
+            println!("{seed} is not the correct seed for this game");
+        }
+        info!("Validating using prepared method");
+        let start = Instant::now();
+        let valid = prepared.validate(&mut rand, seed, true);
+        let elapsed = start.elapsed();
+        info!("Validation took {elapsed:?}");
+        if valid {
             println!("{seed} is the correct seed for this game");
         } else {
             println!("{seed} is not the correct seed for this game");
         }
     } else {
+        info!("Trying every seed from 1 to {}", i32::MAX);
+        let start = Instant::now();
         for seed in 1..i32::MAX {
-            if seed == 100 {
-                break;
-            }
             if seed % 1_000_000 == 0 {
                 debug!("Checking {seed:>10}");
             }
-            if validate(&mut rand, &requests, i64::from(seed), false) {
+            if prepared.validate(&mut rand, i64::from(seed), false) {
+                let elapsed = start.elapsed();
+                info!(
+                    "Trying {seed} seeds took {elapsed:?} ({:?} / seed)",
+                    elapsed / seed.az()
+                );
                 println!("The correct seed is {seed}");
                 break;
             }
         }
     }
 
-    flame::dump_html(File::create("flamegraph.html")?)?;
+    //flame::dump_html(File::create("flamegraph.html")?)?;
 
     Ok(())
 }
 
 #[allow(clippy::too_many_lines)]
-#[flame]
+//#[flame]
 fn validate(input_rand: &mut SeedRand, requests: &[Request], seed: i64, log: bool) -> bool {
     if log {
-        info!("Checking starting positions");
+        debug!("Checking starting positions");
     }
     let request = &requests[0];
     let rand = get_rand(input_rand, seed, 0);
@@ -90,12 +110,8 @@ fn validate(input_rand: &mut SeedRand, requests: &[Request], seed: i64, log: boo
         Coord { x: mx, y: md },
     ];
 
-    rand.shuffle(corner_points.len(), |i, j| {
-        corner_points.swap(i, j);
-    });
-    rand.shuffle(cardinal_points.len(), |i, j| {
-        cardinal_points.swap(i, j);
-    });
+    rand.shuffle(&mut corner_points);
+    rand.shuffle(&mut cardinal_points);
 
     let mut start_points = Vec::with_capacity(corner_points.len() + cardinal_points.len());
     if rand.int_n(2) == 0 {
@@ -106,12 +122,12 @@ fn validate(input_rand: &mut SeedRand, requests: &[Request], seed: i64, log: boo
         start_points.extend_from_slice(&corner_points);
     }
 
-    for i in 0..request.board.snakes.len() {
-        if start_points[i] != request.board.snakes[i].head {
+    for (snake, start_point) in request.board.snakes.iter().zip(start_points.iter()) {
+        if *start_point != snake.head {
             if log {
                 error!(
                     "Starting position is not the same: {} != {}",
-                    start_points[i], request.board.snakes[i].head
+                    start_point, snake.head
                 );
             }
             return false;
@@ -119,22 +135,16 @@ fn validate(input_rand: &mut SeedRand, requests: &[Request], seed: i64, log: boo
     }
 
     if log {
-        info!("Checking food placement");
+        debug!("Checking food placement");
     }
     for i in 1..requests.len() {
         let previous_request = &requests[i - 1];
         let request = &requests[i];
         if previous_request.turn + 1 != request.turn {
-            if log {
-                warn!(
-                    "nonconsecutive requests: {} + 1 != {}",
-                    previous_request.turn, request.turn
-                );
-            }
             continue;
         }
         if log {
-            info!("Checking turn {}", request.turn);
+            debug!("Checking turn {}", request.turn);
         }
         let rand = get_rand(input_rand, seed, i64::from(request.turn - 1));
         let mut forged_request = request.clone();
@@ -160,9 +170,6 @@ fn validate(input_rand: &mut SeedRand, requests: &[Request], seed: i64, log: boo
                 &forged_request,
             ));
         }
-        if log {
-            debug!("snake head: {:?}", request.board.snakes[0].head);
-        }
         if forged_request.board.food != request.board.food {
             let forged: Vec<_> = forged_request
                 .board
@@ -190,14 +197,216 @@ fn validate(input_rand: &mut SeedRand, requests: &[Request], seed: i64, log: boo
     true
 }
 
+#[derive(Debug, Clone)]
+struct PreparedGame {
+    corners: [Option<u8>; 4],
+    cardinals: [Option<u8>; 4],
+    corner_first: bool,
+    food_spawns: Vec<Option<FoodSpawn>>,
+    food_chance: u8,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+enum FoodSpawn {
+    Forced { selected: u16, free_spots: u16 },
+    Random { selected: u16, free_spots: u16 },
+}
+
+impl PreparedGame {
+    fn prepare(requests: &[Request]) -> Self {
+        debug!("Preparing initial board state");
+        let request = &requests[0];
+        assert_eq!(request.turn, 0);
+        let mn = 1;
+        let md = (request.board.width - 1) / 2;
+        let mx = request.board.width - 2;
+        let corners = [
+            Coord { x: mn, y: mn },
+            Coord { x: mn, y: mx },
+            Coord { x: mx, y: mn },
+            Coord { x: mx, y: mx },
+        ];
+        let cardinals = [
+            Coord { x: mn, y: md },
+            Coord { x: md, y: mn },
+            Coord { x: md, y: mx },
+            Coord { x: mx, y: md },
+        ];
+        let corner_first = corners.contains(&request.board.snakes[0].head);
+        let corners = corners.map(|start| {
+            request
+                .board
+                .snakes
+                .iter()
+                .position(|snake| snake.head == start)
+                .map(|pos| (pos % 4).az())
+        });
+        let cardinals = cardinals.map(|start| {
+            request
+                .board
+                .snakes
+                .iter()
+                .position(|snake| snake.head == start)
+                .map(|pos| (pos % 4).az())
+        });
+
+        debug!("Preparing food placement");
+        let food_chance = request.game.ruleset.settings.food_spawn_chance;
+        let food_min = request.game.ruleset.settings.minimum_food;
+
+        let food_spawns = requests
+            .windows(2)
+            .filter_map(|window| {
+                if let [previous, next] = window {
+                    if previous.turn + 1 == next.turn {
+                        Some((previous, next))
+                    } else {
+                        None
+                    }
+                } else {
+                    error!("Windows didn't return 2 requests: {requests:#?}");
+                    None
+                }
+            })
+            .map(|(previous, next)| {
+                // no new food, no food eaten
+                if previous.board.food == next.board.food {
+                    return None;
+                }
+
+                let diff: Vec<_> = next
+                    .board
+                    .food
+                    .iter()
+                    .filter(|next| !previous.board.food.contains(next))
+                    .collect();
+                if diff.is_empty() {
+                    // We didn't spawn any food. It was only eaten
+                    return None;
+                }
+
+                let mut tmp = next.clone();
+                tmp.board.food.clone_from(&previous.board.food);
+                let heads: Vec<_> = tmp.board.snakes.iter().map(|snake| snake.head).collect();
+                tmp.board.food.retain(|coord| !heads.contains(coord));
+                let free_spots = get_unoccupied_points(&tmp);
+
+                let selected = free_spots
+                    .iter()
+                    .position(|spot| spot == diff[0])
+                    .unwrap()
+                    .az();
+                let free_spots = free_spots.len().az();
+
+                if next.board.food.len() == usize::from(food_min) {
+                    Some(FoodSpawn::Forced {
+                        selected,
+                        free_spots,
+                    })
+                } else {
+                    Some(FoodSpawn::Random {
+                        selected,
+                        free_spots,
+                    })
+                }
+            })
+            .collect();
+
+        Self {
+            corners,
+            cardinals,
+            corner_first,
+            food_spawns,
+            food_chance,
+        }
+    }
+
+    fn validate(&self, input_rand: &mut SeedRand, seed: i64, log: bool) -> bool {
+        if log {
+            debug!("Checking Start positions");
+        }
+        let mut corners = [None; 4];
+        for (a, b) in self.corners.iter().zip(corners.iter_mut()) {
+            *b = a.map(usize::from);
+        }
+        let mut cardinals = [None; 4];
+        for (a, b) in self.cardinals.iter().zip(cardinals.iter_mut()) {
+            *b = a.map(usize::from);
+        }
+        let rand = get_rand(input_rand, seed, 0);
+        if !rand.shuffle_check(&mut corners) {
+            if log {
+                error!("Corners are shuffled wrong");
+            }
+            return false;
+        }
+        if !rand.shuffle_check(&mut cardinals) {
+            if log {
+                error!("Cardinals are shuffled wrong");
+            }
+            return false;
+        }
+        if rand.int_n(2) == usize::from(self.corner_first) {
+            if log {
+                error!("Corners and cardinals are flipped: {}", self.corner_first);
+            }
+            return false;
+        }
+
+        self.food_spawns.iter().enumerate().all(|(turn, spawn)| {
+            if log {
+                debug!("Checking turn {turn}");
+            }
+            let rand = get_rand(input_rand, seed, (turn).az());
+
+            match spawn {
+                Some(
+                    spawn @ (FoodSpawn::Forced {
+                        selected,
+                        free_spots,
+                    }
+                    | FoodSpawn::Random {
+                        selected,
+                        free_spots,
+                    }),
+                ) => {
+                    if matches!(spawn, FoodSpawn::Random { .. })
+                        && 100 - rand.int_n(100) >= usize::from(self.food_chance)
+                    {
+                        if log {
+                            error!("Food not spawned, when some should have been spawned");
+                        }
+                        return false;
+                    }
+
+                    let mut foods = vec![None; usize::from(*free_spots)];
+                    foods[usize::from(*selected)] = Some(0);
+                    let correct = rand.shuffle_check(&mut foods);
+                    if log && !correct {
+                        error!("Wrong food selected: {foods:?}");
+                    }
+                    correct
+                }
+                None => {
+                    let correct = 100 - rand.int_n(100) >= usize::from(self.food_chance);
+                    if log && !correct {
+                        error!("Food spawned, when none should have been spawned");
+                    }
+                    correct
+                }
+            }
+        })
+    }
+}
+
 #[must_use]
-#[flame]
+//#[flame]
 fn get_rand(rand: &mut SeedRand, seed: i64, turn: i64) -> &mut SeedRand {
     rand.rand.seed(seed + turn);
     rand
 }
 
-#[flame]
+//#[flame]
 fn check_food_needing_placement(rand: &mut impl Rand, request: &Request) -> usize {
     let min_food = request.game.ruleset.settings.minimum_food as usize;
     let food_spawn_chance = request.game.ruleset.settings.food_spawn_chance as usize;
@@ -213,13 +422,13 @@ fn check_food_needing_placement(rand: &mut impl Rand, request: &Request) -> usiz
     0
 }
 
-#[flame]
+//#[flame]
 fn place_food_randomly(rand: &mut impl Rand, n: usize, request: &Request) -> Vec<Coord> {
     let unoccupied_points = get_unoccupied_points(request);
     place_food_randomly_at_positions(rand, n, unoccupied_points)
 }
 
-#[flame]
+//#[flame]
 fn place_food_randomly_at_positions(
     rand: &mut impl Rand,
     n: usize,
@@ -227,9 +436,7 @@ fn place_food_randomly_at_positions(
 ) -> Vec<Coord> {
     let n = n.min(positions.len());
 
-    rand.shuffle(positions.len(), |i, j| {
-        positions.swap(i, j);
-    });
+    rand.shuffle(&mut positions);
 
     let mut out = Vec::new();
     let mut i = 0;
@@ -240,7 +447,7 @@ fn place_food_randomly_at_positions(
     out
 }
 
-#[flame]
+//#[flame]
 fn get_unoccupied_points(request: &Request) -> Vec<Coord> {
     let possible_moves: Vec<_> = request
         .board
@@ -275,7 +482,10 @@ trait Rand {
     fn range(&mut self, min: usize, max: usize) -> usize;
     /// shuffle an array of n values using the swap function
     #[allow(unused)]
-    fn shuffle(&mut self, n: usize, swap: impl FnMut(usize, usize));
+    fn shuffle<T>(&mut self, data: &mut [T]);
+    /// checks shuffle an array of n values using the swap function
+    #[allow(unused)]
+    fn shuffle_check(&mut self, data: &mut [Option<usize>]) -> bool;
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -284,7 +494,7 @@ struct SeedRand {
 }
 
 impl SeedRand {
-    #[flame]
+    //#[flame]
     fn new(seed: i64) -> Self {
         Self {
             rand: GoRand::new(GoSource::new(seed)),
@@ -293,19 +503,23 @@ impl SeedRand {
 }
 
 impl Rand for SeedRand {
-    #[flame]
+    //#[flame]
     fn int_n(&mut self, n: usize) -> usize {
         self.rand.int_n(n)
     }
 
-    #[flame]
+    //#[flame]
     fn range(&mut self, min: usize, max: usize) -> usize {
         self.rand.int_n(max - min + 1) + min
     }
 
-    #[flame]
-    fn shuffle(&mut self, n: usize, swap: impl FnMut(usize, usize)) {
-        self.rand.shuffle(n, swap);
+    //#[flame]
+    fn shuffle<T>(&mut self, data: &mut [T]) {
+        self.rand.shuffle(data);
+    }
+
+    fn shuffle_check(&mut self, data: &mut [Option<usize>]) -> bool {
+        self.rand.shuffle_check(data)
     }
 }
 
@@ -315,26 +529,26 @@ struct GoRand {
 }
 
 impl GoRand {
-    #[flame]
+    //#[flame]
     const fn new(src: GoSource) -> Self {
         Self { src }
     }
 
     // Seed uses the provided seed value to initialize the generator to a deterministic state.
     // Seed should not be called concurrently with any other [Rand] method.
-    #[flame]
+    //#[flame]
     fn seed(&mut self, seed: i64) {
         self.src.seed(seed);
     }
 
     /// returns a non-negative pseudo-random 63-bit integer
-    #[flame]
+    //#[flame]
     fn int63(&mut self) -> u64 {
         self.src.int63()
     }
 
     /// returns a non-negative pseudo-random 63-bit integer
-    #[flame]
+    //#[flame]
     fn int31(&mut self) -> u32 {
         (self.src.int63() >> 32) as u32
     }
@@ -343,7 +557,7 @@ impl GoRand {
     ///
     /// # Panics
     /// if n <= 0
-    #[flame]
+    //#[flame]
     fn int63_n(&mut self, n: u64) -> u64 {
         assert!((n > 0), "invalid argument to int63_n");
         if n.is_power_of_two() {
@@ -362,7 +576,7 @@ impl GoRand {
     ///
     /// # Panics
     /// if n <= 0
-    #[flame]
+    //#[flame]
     fn int31_n(&mut self, n: u32) -> u32 {
         assert!((n > 0), "invalid argument to in32_n");
         if n.is_power_of_two() {
@@ -377,12 +591,12 @@ impl GoRand {
         v % n
     }
 
-    #[flame]
+    //#[flame]
     fn uint32(&mut self) -> u32 {
         (self.int63() >> 31).az()
     }
 
-    #[flame]
+    //#[flame]
     fn int31_n2(&mut self, n: u32) -> u32 {
         let mut v = self.uint32();
         let mut prod = u64::from(v) * u64::from(n);
@@ -402,7 +616,7 @@ impl GoRand {
     ///
     /// # Panics
     /// if n <= 0
-    #[flame]
+    //#[flame]
     fn int_n(&mut self, n: usize) -> usize {
         assert!((n > 0), "invalid argument to int_n");
         if n < (1 << 31) {
@@ -413,34 +627,65 @@ impl GoRand {
     }
 
     #[allow(unused)]
-    #[flame]
+    //#[flame]
     fn int(&mut self) -> usize {
         self.int63().az()
     }
 
     // Shuffle pseudo-randomizes the order of elements. n is the number of elements. Shouffle
     // panics if n < 0. swap swaps the elements with indices i and j.
-    #[flame]
-    fn shuffle(&mut self, n: usize, mut swap: impl FnMut(usize, usize)) {
+    //#[flame]
+    fn shuffle<T>(&mut self, data: &mut [T]) {
         // Fisher-Yates shuffle: https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle
         // Shuffle really ought not be called with n that doesn't fit in 32 bits.
         // Not only will it take a very long time, but with 2³¹! possible permutations,
         // there's no way that any PRNG can have a big enough internal state to
         // generate even a minuscule percentage of the possible permutations.
         // Nevertheless, the right API signature accepts an int n, so handle it as best we can.
-        let mut i = n.saturating_sub(1);
+        let mut i = data.len().saturating_sub(1);
         while i > (1 << 31) - 1 - 1 {
             let j = self.int63_n((i + 1) as u64).az();
-            swap(i, j);
+            data.swap(i, j);
             i -= 1;
         }
         let mut i: u32 = i.az();
         while i > 0 {
             let j = self.int31_n2(i + 1);
-            swap(i.az(), j.az());
+            data.swap(i.az(), j.az());
             i -= 1;
         }
     }
+
+    // Shuffle the array and check if it is deshuffled afterwards
+    fn shuffle_check(&mut self, data: &mut [Option<usize>]) -> bool {
+        let mut i = data.len().saturating_sub(1);
+        while i > (1 << 31) - 1 - 1 {
+            let j = self.int63_n((i + 1) as u64).az();
+            let datum = data[j];
+            if let Some(new_i) = datum {
+                if new_i != i {
+                    return false;
+                }
+            }
+            data.swap(i, j);
+            i -= 1;
+        }
+        let mut i: u32 = i.az();
+        while i > 0 {
+            let j = self.int31_n2(i + 1);
+            let i_usize = i.az();
+            let j_usize = j.az();
+            let datum = data[j_usize];
+            if let Some(i) = datum {
+                if i != i_usize {
+                    return false;
+                }
+            }
+            data.swap(i_usize, j_usize);
+            i -= 1;
+        }
+        true
+    }
 }
 
 static CHAIN_RAW: LazyLock<Mmap> = LazyLock::new(|| unsafe {
@@ -468,7 +713,7 @@ impl GoSource {
     const MAX: u64 = 1 << 63;
     const MASK: u64 = Self::MAX - 1;
 
-    #[flame]
+    //#[flame]
     fn new(seed: i64) -> Self {
         let seed: i32 = (seed % i64::from(i32::MAX)).az();
         let seed = match seed {
@@ -490,13 +735,13 @@ impl GoSource {
 }
 
 impl GoSource {
-    #[flame]
+    //#[flame]
     fn int63(&mut self) -> u64 {
         self.uint64() & Self::MASK
     }
 
     #[allow(unused)]
-    #[flame]
+    //#[flame]
     fn seed(&mut self, seed: i64) {
         let seed: i32 = (seed % i64::from(i32::MAX)).az();
         let seed = match seed {
@@ -514,7 +759,7 @@ impl GoSource {
         self.vec.clear();
     }
 
-    #[flame]
+    //#[flame]
     fn uint64(&mut self) -> u64 {
         self.tap = self.tap.checked_sub(1).unwrap_or(Self::LEN - 1);
         self.feed = self.feed.checked_sub(1).unwrap_or(Self::LEN - 1);
@@ -536,7 +781,7 @@ impl GoSource {
     }
 }
 
-#[flame]
+//#[flame]
 fn seed_one(base: usize, i: u16) -> i64 {
     let i = usize::from(i);
     let base = base + 20 + i * 3;
@@ -554,7 +799,7 @@ fn seed_one(base: usize, i: u16) -> i64 {
     u ^ GoSource::COOKED[i]
 }
 
-#[flame]
+//#[flame]
 const fn seedrand(x: i32) -> i32 {
     const A: i32 = 48_271;
     const Q: i32 = 44_488;