From f8b2826d389ec2b94281cab0ae8bb3ab1661c07b Mon Sep 17 00:00:00 2001
From: James Tanner <jimmyt857@gmail.com>
Date: Sat, 14 Oct 2023 10:35:25 -0700
Subject: [PATCH] Other small cleanups

---
 sim/core/energy.go |  5 ++++-
 sim/core/unit.go   | 16 +++++++---------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/sim/core/energy.go b/sim/core/energy.go
index 94117fc5d2..787fba89f6 100644
--- a/sim/core/energy.go
+++ b/sim/core/energy.go
@@ -28,7 +28,7 @@ type energyBar struct {
 	// [10, 15, 20, 30, 60, 85]
 	energyDecisionThresholds []int
 
-	// Slice with len == maxEnergy with each index corresponding to an amount of energy. Looks like this:
+	// Slice with len == maxEnergy+1 with each index corresponding to an amount of energy. Looks like this:
 	// [0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, ...]
 	// Increments by 1 at each value of energyDecisionThresholds.
 	cumulativeEnergyDecisionThresholds []int
@@ -72,6 +72,9 @@ func (unit *Unit) EnableEnergyBar(maxEnergy float64, onEnergyGain OnEnergyGain)
 
 // Computes the energy thresholds.
 func (eb *energyBar) setupEnergyThresholds() {
+	if eb.unit == nil {
+		return
+	}
 	if !eb.unit.IsUsingAPL {
 		return
 	}
diff --git a/sim/core/unit.go b/sim/core/unit.go
index 33ba6789e6..e8cd51bfc9 100644
--- a/sim/core/unit.go
+++ b/sim/core/unit.go
@@ -439,15 +439,13 @@ func (unit *Unit) finalize() {
 		spell.finalize()
 	}
 
-	if unit.HasEnergyBar() {
-		// For now, restrict this optimization to rogues only. Ferals will require
-		// some extra logic to handle their ExcessEnergy() calc.
-		agent := unit.Env.Raid.GetPlayerFromUnit(unit)
-		if agent != nil && agent.GetCharacter().Class == proto.Class_ClassRogue {
-			unit.Env.RegisterPostFinalizeEffect(func() {
-				unit.energyBar.setupEnergyThresholds()
-			})
-		}
+	// For now, restrict this optimization to rogues only. Ferals will require
+	// some extra logic to handle their ExcessEnergy() calc.
+	agent := unit.Env.Raid.GetPlayerFromUnit(unit)
+	if agent != nil && agent.GetCharacter().Class == proto.Class_ClassRogue {
+		unit.Env.RegisterPostFinalizeEffect(func() {
+			unit.energyBar.setupEnergyThresholds()
+		})
 	}
 }