fix potential desync on buff dispel

rhazarian · Nov 8, 2024 · 54f516b · 54f516b
1 parent 6389eac
commit 54f516b
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/src/engine/behaviour/unit.ts b/src/engine/behaviour/unit.ts
@@ -44,11 +44,11 @@ export abstract class UnitBehavior<PeriodicActionParameters extends any[] = any[
         // no-op
     }
 
-    protected onAbilityGained(ability: Ability): void {
+    public onAbilityGained(ability: Ability): void {
         // no-op
     }
 
-    protected onAbilityLost(ability: Ability): void {
+    public onAbilityLost(ability: Ability): void {
         // no-op
     }
 
@@ -83,6 +83,14 @@ export abstract class UnitBehavior<PeriodicActionParameters extends any[] = any[
             UnitBehavior.forAll(target, "onDamageReceived", source, event)
         })
 
+        Unit.abilityGainedEvent.addListener((source, target) => {
+            UnitBehavior.forAll(source, "onAbilityGained", target)
+        })
+
+        Unit.abilityLostEvent.addListener((source, target) => {
+            UnitBehavior.forAll(source, "onAbilityLost", target)
+        })
+
         Unit.deathEvent.addListener((target, source) => {
             if (source != undefined) {
                 UnitBehavior.forAll(source, "onKill", target)

diff --git a/src/engine/buff.ts b/src/engine/buff.ts
@@ -1300,25 +1300,24 @@ export class Buff<
     }
 
     static {
+        const destroyBuffIfNeeded = (buff: Buff) => {
+            if (getUnitAbility(buff[BuffPropertyKey.UNIT].handle, buff.typeId) != buff.handle) {
+                buff.destroy()
+            }
+        }
+
         checkBuff = (unit: Unit, buffTypeId: number): void => {
             const buffByTypeId = buffByTypeIdByUnit.get(unit)
             if (buffByTypeId != undefined) {
                 const buff = buffByTypeId.get(buffTypeId)
-                if (buff != undefined && getUnitAbility(unit.handle, buffTypeId) != buff.handle) {
-                    buff.destroy()
+                if (buff != undefined) {
+                    destroyBuffIfNeeded(buff)
                 }
             }
         }
 
         checkBuffs = (unit: Unit): void => {
-            const buffByTypeId = buffByTypeIdByUnit.get(unit)
-            if (buffByTypeId != undefined) {
-                for (const [, buff] of buffByTypeId) {
-                    if (getUnitAbility(unit.handle, buff.typeId) != buff.handle) {
-                        buff.destroy()
-                    }
-                }
-            }
+            Buff.forAll(unit, destroyBuffIfNeeded)
         }
 
         Unit.abilityChannelingStartEvent.addListener(