From a7c9098af2df4f64a23701af022ec2ddc087897d Mon Sep 17 00:00:00 2001
From: MetricExpansion <>
Date: Sun, 15 Nov 2020 11:55:09 -0800
Subject: [PATCH] Fix memory bug.

---
 include/ArmorAddonOverrideService.h |  2 +-
 src/ArmorAddonOverrideService.cpp   | 22 ++++++++--------------
 src/OutfitSystem.cpp                | 10 +++++++---
 3 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/include/ArmorAddonOverrideService.h b/include/ArmorAddonOverrideService.h
index 6b59d2d..c4dc466 100644
--- a/include/ArmorAddonOverrideService.h
+++ b/include/ArmorAddonOverrideService.h
@@ -108,7 +108,7 @@ class ArmorAddonOverrideService {
       void setLocationOutfit(LocationType location, const char* name);
       void unsetLocationOutfit(LocationType location);
       std::optional<cobb::istring> getLocationOutfit(LocationType location);
-      LocationType checkLocationType(const std::unordered_set<std::string>& keywords, const WeatherFlags& weather_flags);
+      std::optional<LocationType> checkLocationType(const std::unordered_set<std::string>& keywords, const WeatherFlags& weather_flags);
       //
       bool shouldOverride() const noexcept;
       void getOutfitNames(std::vector<std::string>& out, bool favoritesOnly = false) const;
diff --git a/src/ArmorAddonOverrideService.cpp b/src/ArmorAddonOverrideService.cpp
index 13b4473..9ddee8c 100644
--- a/src/ArmorAddonOverrideService.cpp
+++ b/src/ArmorAddonOverrideService.cpp
@@ -235,20 +235,14 @@ std::optional<cobb::istring> ArmorAddonOverrideService::getLocationOutfit(Locati
     }
 }
 
-LocationType ArmorAddonOverrideService::checkLocationType(const std::unordered_set<std::string>& keywords, const WeatherFlags& weather_flags) {
-    if (keywords.count("LocTypeHabitation")) {
-        return LocationType::Town;
-    } else if (keywords.count("LocTypeDungeon")) {
-        return LocationType::Dungeon;
-    }
-    // By now, we know we're in no location more specific than "World"
-    //
-    // We will now perform weather checks and return the value *only if* that
-    // weather has an outfit assigned.
-    if (weather_flags.snowy && locationOutfits.count(LocationType::WorldSnowy)) {
-        return LocationType::WorldSnowy;
-    }
-    return LocationType::World;
+#define CHECK_LOCATION(TYPE, CHECK_CODE) if (locationOutfits.count(LocationType::TYPE) && (CHECK_CODE)) return std::optional<LocationType>(LocationType::TYPE);
+
+std::optional<LocationType> ArmorAddonOverrideService::checkLocationType(const std::unordered_set<std::string>& keywords, const WeatherFlags& weather_flags) {
+    CHECK_LOCATION(Town, keywords.count("LocTypeHabitation"));
+    CHECK_LOCATION(Dungeon, keywords.count("LocTypeDungeon"));
+    CHECK_LOCATION(WorldSnowy, weather_flags.snowy);
+    CHECK_LOCATION(World, true);
+    return std::optional<LocationType>();
 }
 
 bool ArmorAddonOverrideService::shouldOverride() const noexcept {
diff --git a/src/OutfitSystem.cpp b/src/OutfitSystem.cpp
index 6bd0201..f1c53f0 100644
--- a/src/OutfitSystem.cpp
+++ b/src/OutfitSystem.cpp
@@ -596,7 +596,7 @@ extern SKSESerializationInterface* g_Serialization;
             }
             return result;
         }
-        LocationType identifyLocation(RE::BGSLocation* location, RE::TESWeather* weather) {
+        std::optional<LocationType> identifyLocation(RE::BGSLocation* location, RE::TESWeather* weather) {
             // Just a helper function to classify a location.
             // TODO: Think of a better place than this since we're not exposing it to Papyrus.
             auto& service = ArmorAddonOverrideService::GetInstance();
@@ -627,7 +627,9 @@ extern SKSESerializationInterface* g_Serialization;
             return service.checkLocationType(keywords, weather_flags);
         }
         UInt32 IdentifyLocationType(VMClassRegistry* registry, UInt32 stackId, StaticFunctionTag*, BGSLocation* location_skse, TESWeather* weather_skse) {
-            return static_cast<UInt32>(identifyLocation((RE::BGSLocation*) location_skse, (RE::TESWeather*) weather_skse));
+            // NOTE: Identify the location for Papyrus. In the event no location is identified, we lie to Papyrus and say "World".
+            //       Therefore, Papyrus cannot assume that locations returned have an outfit assigned, at least not for "World".
+            return static_cast<UInt32>(identifyLocation((RE::BGSLocation*) location_skse, (RE::TESWeather*) weather_skse).value_or(LocationType::World));
         }
         void SetOutfitUsingLocation(VMClassRegistry* registry, UInt32 stackId, StaticFunctionTag*, BGSLocation* location_skse, TESWeather* weather_skse) {
             // NOTE: Location can be NULL.
@@ -642,7 +644,9 @@ extern SKSESerializationInterface* g_Serialization;
                 sprintf_s(message, "SOS: This location is a %s.", locationName);
                 RE::DebugNotification(message, nullptr, false);
                 */
-                service.setOutfitUsingLocation(location);
+                if (location.has_value()) {
+                    service.setOutfitUsingLocation(location.value());
+                }
             }
 
         }