Fix sprite owner array not being rebuilt on save loads

While other sprite owner updating code caught most of the cases,
hidden units would lose their sprite owner on load. (Especially
when loading fresh from startup)

src/bw/mod.rs

@@ -68,6 +68,10 @@ pub fn give_ai(unit: bw_dat::Unit) {
     unsafe { samase::give_ai(*unit) }
 }
 
+pub fn unit_array() -> (*mut Unit, usize) {
+    unsafe { samase::unit_array() }
+}
+
 whack_hooks!(stdcall, 0x00400000,
     0x00488AF0 => increment_death_scores(@edi *mut Unit, @edx u8);
     0x004465C0 => choose_placement_position(u32, u32, *mut Point, u32, @ecx *mut Unit) -> u32;

src/globals.rs

@@ -52,6 +52,7 @@ pub unsafe extern fn init_game() {
     globals.iscript_state = iscript::IscriptState::from_script(&iscript);
     iscript::set_as_bw_script(iscript);
     *Globals::get("init") = globals;
+    iscript::rebuild_sprite_owners();
 }
 
 pub unsafe extern fn save(set_data: unsafe extern fn(*const u8, usize)) {
@@ -81,6 +82,7 @@ pub unsafe extern fn load(ptr: *const u8, len: usize) -> u32 {
     let iscript = iscript::load_iscript(true);
     iscript::set_as_bw_script(iscript);
     *Globals::get("load") = data;
+    iscript::rebuild_sprite_owners();
     1
 }
 

src/iscript.rs

@@ -820,6 +820,30 @@ impl SpriteOwnerMap {
     }
 }
 
+pub fn rebuild_sprite_owners() {
+    let mut map = SPRITE_OWNER_MAP.lock("rebuild_sprite_owners");
+    map.unit_mapping.clear();
+    map.bullet_mapping.clear();
+    unsafe {
+        let (units, len) = bw::unit_array();
+        for i in 0..len {
+            let unit = units.add(i);
+            let sprite = (*unit).sprite;
+            if !sprite.is_null() {
+                map.add_unit(Unit::from_ptr(unit).unwrap(), sprite);
+            }
+        }
+        let mut bullet = bw::first_active_bullet();
+        while !bullet.is_null() {
+            let sprite = (*bullet).sprite;
+            if !sprite.is_null() {
+                map.add_bullet(bullet, sprite);
+            }
+            bullet = (*bullet).next;
+        }
+    }
+}
+
 unsafe fn invalid_aice_command(iscript: *mut bw::Iscript, image: *mut bw::Image, offset: CodePos) {
     let sprite_id = match Sprite::from_ptr((*image).parent) {
         Some(sprite) => sprite.id().0,

src/samase.rs

@@ -164,6 +164,14 @@ pub unsafe fn give_ai(unit: *mut bw::Unit) {
     (GIVE_AI.0.unwrap())(unit)
 }
 
+static mut UNIT_ARRAY_LEN: GlobalFunc<extern fn(*mut *mut bw::Unit, *mut usize)> = GlobalFunc(None);
+pub unsafe fn unit_array() -> (*mut bw::Unit, usize) {
+    let mut arr = null_mut();
+    let mut len = 0;
+    (UNIT_ARRAY_LEN.0.unwrap())(&mut arr, &mut len);
+    (arr, len)
+}
+
 pub struct SamaseBox {
     data: NonNull<u8>,
     len: usize,
@@ -267,6 +275,7 @@ pub unsafe extern fn samase_plugin_init(api: *const PluginApi) {
     FINISH_UNIT_PRE.0 = Some(mem::transmute(((*api).finish_unit_pre)()));
     FINISH_UNIT_POST.0 = Some(mem::transmute(((*api).finish_unit_post)()));
     GIVE_AI.0 = Some(mem::transmute(((*api).give_ai)()));
+    UNIT_ARRAY_LEN.0 = Some(mem::transmute(((*api).unit_array_len)()));
     let result = ((*api).extend_save)(
         "aice\0".as_ptr(),
         Some(crate::globals::save),