Allow ai to peek past first request if it has money to satisfy it

without affecting earlier requests

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "aise"
-version = "2.34.5"
+version = "2.35.0"
 authors = ["Markus Heikkinen <ittevien@gmail.com>"]
 edition = "2018"
 

src/ai.rs

@@ -129,23 +129,14 @@ impl PlayerAi {
     }
 
     fn push_request(&self, req: bw::AiSpendingRequest) {
-        let requests = unsafe { &mut (*self.0).requests[..] };
-        let mut pos = unsafe {
-            if usize::from((*self.0).request_count) < requests.len() {
-                (*self.0).request_count += 1;
-            }
-            isize::from((*self.0).request_count) - 1
-        };
-        let mut prev = pos;
-        requests[pos as usize] = req;
-        pos = (pos - 1) >> 1;
-        while prev > 0 {
-            if requests[prev as usize].priority <= requests[pos as usize].priority {
-                break;
+        unsafe {
+            let requests = &mut (*self.0).requests[..];
+            let request_count = usize::from((*self.0).request_count);
+            if request_count >= requests.len() {
+                return;
             }
-            requests.swap(prev as usize, pos as usize);
-            prev = pos;
-            pos = (pos - 1) >> 1;
+            (*self.0).request_count += 1;
+            request_heap_push(requests, request_count, req);
         }
     }
 
@@ -188,6 +179,25 @@ impl PlayerAi {
         }
     }
 
+    pub fn copy_requests(&self) -> SpendingRequestsCopy {
+        unsafe {
+            SpendingRequestsCopy {
+                buffer: (*self.0).requests,
+                count: (*self.0).request_count,
+                deleted: 0,
+            }
+        }
+    }
+
+    pub fn set_to_copied_requests(&self, copy: &SpendingRequestsCopy) {
+        unsafe {
+            let len = copy.count as usize;
+            let slice = &copy.buffer[..len];
+            (&mut (*self.0).requests[..len]).copy_from_slice(slice);
+            (*self.0).request_count = copy.count;
+        }
+    }
+
     pub fn has_resources(&self, game: Game, players: *mut bw::Player, cost: &Cost) -> bool {
         unsafe {
             let player = self.1;
@@ -202,6 +212,16 @@ impl PlayerAi {
         }
     }
 
+    pub fn available_resources(&self) -> AiResources {
+        unsafe {
+            AiResources {
+                minerals: (*self.0).minerals_available,
+                gas: (*self.0).gas_available,
+                supply: (*self.0).supply_available,
+            }
+        }
+    }
+
     /// spent_money is true when the ai actually built something, false on failures
     pub fn remove_resource_need(&self, cost: &Cost, spent_money: bool) {
         unsafe {
@@ -228,23 +248,7 @@ impl PlayerAi {
                 return;
             }
         }
-        requests[0] = requests[new_count];
-        let mut prev = 0;
-        let mut pos = 1;
-        while pos <= new_count - 1 {
-            let index =
-                if pos + 1 >= new_count || requests[pos].priority > requests[pos + 1].priority {
-                    pos
-                } else {
-                    pos + 1
-                };
-            if requests[prev].priority >= requests[index].priority {
-                break;
-            }
-            requests.swap(prev, index);
-            prev = index;
-            pos = (index << 1) + 1;
-        }
+        request_heap_pop(requests);
     }
 
     fn is_guard_being_trained(&self, guard: *mut bw::GuardAi) -> bool {
@@ -293,6 +297,106 @@ impl PlayerAi {
     }
 }
 
+pub struct SpendingRequestsCopy {
+    buffer: [bw::AiSpendingRequest; 0x3f],
+    count: u8,
+    /// Deleted requests are placed at end of buffer in reverse, as long
+    /// as there is space for them
+    deleted: u8,
+}
+
+impl SpendingRequestsCopy {
+    pub fn first_request(&self) -> Option<bw::AiSpendingRequest> {
+        if self.count == 0 {
+            None
+        } else {
+            Some(self.buffer[0])
+        }
+    }
+
+    pub fn pop_request_save(&mut self) {
+        let request = match self.first_request() {
+            Some(s) => s,
+            None => return,
+        };
+        self.pop_request_dont_save();
+        if (self.count as usize + self.deleted as usize) < self.buffer.len() {
+            let free_index = self.buffer.len() - 1 - self.deleted as usize;
+            self.buffer[free_index] = request;
+            self.deleted += 1;
+        }
+    }
+
+    pub fn pop_request_dont_save(&mut self) {
+        if self.count == 0 {
+            return;
+        }
+        request_heap_pop(&mut self.buffer[..(self.count as usize)]);
+        self.count -= 1;
+    }
+
+    pub fn restore_popped_requests(&mut self) {
+        while self.deleted > 0 {
+            let index = self.buffer.len() - self.deleted as usize;
+            let request = self.buffer[index];
+            self.deleted -= 1;
+            self.push_request(request);
+        }
+    }
+
+    pub fn push_request(&mut self, request: bw::AiSpendingRequest) {
+        let count = usize::from(self.count);
+        if count >= self.buffer.len() {
+            return;
+        }
+        self.count += 1;
+        request_heap_push(&mut self.buffer, count, request);
+    }
+}
+
+fn request_heap_push(
+    requests: &mut [bw::AiSpendingRequest],
+    current_size: usize,
+    req: bw::AiSpendingRequest,
+) {
+    let mut pos = current_size as isize;
+    let mut prev = pos;
+    requests[pos as usize] = req;
+    pos = (pos - 1) >> 1;
+    while prev > 0 {
+        if requests[prev as usize].priority <= requests[pos as usize].priority {
+            break;
+        }
+        requests.swap(prev as usize, pos as usize);
+        prev = pos;
+        pos = (pos - 1) >> 1;
+    }
+}
+
+fn request_heap_pop(requests: &mut [bw::AiSpendingRequest]) {
+    if requests.len() < 2 {
+        return;
+    }
+    let new_count = requests.len() - 1;
+    requests[0] = requests[new_count];
+    let mut prev = 0;
+    let mut pos = 1;
+    while pos <= new_count - 1 {
+        let index =
+            if pos + 1 >= new_count || requests[pos].priority > requests[pos + 1].priority {
+                pos
+            } else {
+                pos + 1
+            };
+        if requests[prev].priority >= requests[index].priority {
+            break;
+        }
+        requests.swap(prev, index);
+        prev = index;
+        pos = (index << 1) + 1;
+    }
+}
+
 /// Returns unit ids which are counted for request satisfication for request of `unit_id`
 /// So practically higher tier buildings with morph.
 pub fn request_equivalent_unit_ids(unit_id: UnitId) -> [UnitId; 4] {
@@ -403,9 +507,30 @@ pub fn count_units(player: u8, unit_id: UnitId, game: Game) -> u32 {
     morphing + existing
 }
 
+pub struct AiResources {
+    pub minerals: u32,
+    pub gas: u32,
+    pub supply: u32,
+}
+
+impl AiResources {
+    pub fn has_enough_for_cost(&self, cost: &Cost) -> bool {
+        self.minerals >= cost.minerals &&
+            self.gas >= cost.gas &&
+            self.supply >= cost.supply
+    }
+
+    pub fn reduce_cost(&mut self, cost: &Cost) {
+        self.minerals = self.minerals.saturating_sub(cost.minerals);
+        self.gas = self.gas.saturating_sub(cost.gas);
+        self.supply = self.supply.saturating_sub(cost.supply);
+    }
+}
+
 pub struct Cost {
     pub minerals: u32,
     pub gas: u32,
+    /// Includes supply cost for both units if the unit is dual birth unit.
     pub supply: u32,
     pub races: RaceFlags,
 }

src/ai_spending.rs

@@ -20,38 +20,112 @@ pub unsafe fn frame_hook(
     for player in 0..8 {
         let ai_mode = &ai_mode[player as usize];
         let ai = player_ai.player(player);
+        let mut remaining_money = ai.available_resources();
         while let Some(request) = ai.first_request() {
-            let can = can_satisfy_request(game, players, &ai, player, &request, ai_mode).is_ok();
-            let mut handled = false;
-            if can {
-                // Handle building morphs since preplaced colonies don't check for requests
-                // (Since their order switches to idle because they don't have a town on frame 0)
-                // Could handle other requests as well, but no need at the moment.
-                if handle_building_morph(game, players, &ai, &request, player) == Ok(()) {
-                    handled = true;
-                }
-                // Handle military/guard training as well since train opcode is otherwise
-                // bad at spreading work across buildings
-                if handle_training(game, players, &ai, unit_search, &request, player) == Ok(()) {
-                    handled = true;
-                }
-                if !handled {
-                    break;
-                }
-            }
             let level = if request.ty == 5 {
                 game.upgrade_level(player as u8, UpgradeId(request.id))
                     .saturating_add(1)
             } else {
                 0
             };
             let cost = ai::request_cost(&request, level);
+            if !remaining_money.has_enough_for_cost(&cost) {
+                if ai_mode.wait_for_resources {
+                    try_handle_single_resource_request(game, players, &ai, unit_search, ai_mode);
+                    break;
+                }
+            }
+            let can = can_satisfy_request(game, players, &ai, player, &request, ai_mode).is_ok();
+            let mut handled = false;
+            if can {
+                handled = try_handle_request(game, players, &ai, unit_search, &request);
+                if !handled {
+                    break;
+                } else {
+                    remaining_money.reduce_cost(&cost);
+                }
+            }
             ai.remove_resource_need(&cost, handled);
             ai.pop_request()
         }
     }
 }
 
+fn try_handle_request(
+    game: Game,
+    players: *mut bw::Player,
+    ai: &PlayerAi,
+    unit_search: &UnitSearch,
+    request: &bw::AiSpendingRequest,
+) -> bool {
+    let player = ai.1;
+    // Handle building morphs since preplaced colonies don't check for requests
+    // (Since their order switches to idle because they don't have a town on frame 0)
+    // Could handle other requests as well, but no need at the moment.
+    if handle_building_morph(game, players, &ai, &request, player) == Ok(()) {
+        return true;
+    }
+    // Handle military/guard training as well since train opcode is otherwise
+    // bad at spreading work across buildings
+    if handle_training(game, players, &ai, unit_search, &request, player) == Ok(()) {
+        return true;
+    }
+    false
+}
+
+unsafe fn try_handle_single_resource_request(
+    game: Game,
+    players: *mut bw::Player,
+    ai: &PlayerAi,
+    unit_search: &UnitSearch,
+    ai_mode: &AiMode,
+) {
+    // If the first request cannot be satisfied, but there would be enough
+    // money for single-resource request, handle that request instead.
+    // E.g. protoss has 400 minerals and 100 gas, requests are
+    // [build robo fac, train zealot], it should train the zealot without
+    // waiting for robo gas.
+    //
+    // Since the request queue is awkwardly a heap, this pops requests from copied heap,
+    // and if any requests manage to be handled, re-pushes any popped requests back
+    // and writes that to global state.
+    let mut remaining_money = ai.available_resources();
+    let mut request_copy = ai.copy_requests();
+    let mut any_handled = false;
+    let player = ai.1;
+    while let Some(request) = request_copy.first_request() {
+        if remaining_money.minerals == 0 && remaining_money.gas == 0 {
+            break;
+        }
+        let level = if request.ty == 5 {
+            game.upgrade_level(player as u8, UpgradeId(request.id))
+                .saturating_add(1)
+        } else {
+            0
+        };
+        let cost = ai::request_cost(&request, level);
+        if remaining_money.has_enough_for_cost(&cost) {
+            let can = can_satisfy_request(game, players, &ai, player, &request, ai_mode).is_ok();
+            if can {
+                let handled = try_handle_request(game, players, ai, unit_search, &request);
+                if handled {
+                    any_handled = true;
+                    ai.remove_resource_need(&cost, true);
+                    remaining_money.reduce_cost(&cost);
+                    request_copy.pop_request_dont_save();
+                    continue;
+                }
+            }
+        }
+        remaining_money.reduce_cost(&cost);
+        request_copy.pop_request_save();
+    }
+    if any_handled {
+        request_copy.restore_popped_requests();
+        ai.set_to_copied_requests(&request_copy);
+    }
+}
+
 fn handle_building_morph(
     game: Game,
     players: *mut bw::Player,