Add TESObjectREFR::GetEnchantment

powerof3 · Nov 25, 2024 · d160a87 · d160a87
1 parent eb78881
commit d160a87
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 2 deletions.
diff --git a/include/RE/A/ArrowProjectile.h b/include/RE/A/ArrowProjectile.h
@@ -1,12 +1,10 @@
 #pragma once
 
-#include "RE/F/FormTypes.h"
 #include "RE/M/MissileProjectile.h"
 
 namespace RE
 {
 	class AlchemyItem;
-	class EnchantmentItem;
 
 	class ArrowProjectile : public MissileProjectile
 	{

diff --git a/include/RE/Skyrim.h b/include/RE/Skyrim.h
@@ -1653,6 +1653,7 @@
 #include "RE/V/VirtualMachine.h"
 #include "RE/V/VoiceSpellFireHandler.h"
 #include "RE/W/WeaponAnimationGraphManagerHolder.h"
+#include "RE/W/WeaponEnchantmentController.h"
 #include "RE/W/WeatherType.h"
 #include "RE/W/WerewolfEffect.h"
 #include "RE/W/WerewolfFeedEffect.h"

diff --git a/include/RE/T/TESObjectREFR.h b/include/RE/T/TESObjectREFR.h
@@ -35,6 +35,7 @@ namespace RE
 	class BSFaceGenNiNode;
 	class BSFlattenedBoneTree;
 	class DialogueResponse;
+	class EnchantmentItem;
 	class Explosion;
 	class InventoryChanges;
 	class InventoryEntryData;
@@ -380,6 +381,7 @@ namespace RE
 		BGSEncounterZone*                       GetEncounterZone() const;
 		BGSLocation*                            GetEditorLocation() const;
 		bool                                    GetEditorLocation(NiPoint3& a_outPos, NiPoint3& a_outRot, TESForm*& a_outWorldOrCell, TESObjectCELL* a_fallback);
+		EnchantmentItem*                        GetEnchantment() const;
 		std::optional<double>                   GetEnchantmentCharge() const;
 		TESFaction*                             GetFactionOwner();
 		ObjectRefHandle                         GetHandle();

diff --git a/src/RE/T/TESObjectREFR.cpp b/src/RE/T/TESObjectREFR.cpp
@@ -275,6 +275,24 @@ namespace RE
 		return GetEditorLocation2(a_outPos, a_outRot, a_outWorldOrCell, a_fallback);
 	}
 
+	EnchantmentItem* TESObjectREFR::GetEnchantment() const
+	{
+		auto xEnch = extraList.GetByType<ExtraEnchantment>();
+		if (xEnch && xEnch->enchantment) {
+			return xEnch->enchantment;
+		}
+
+		auto obj = GetObjectReference();
+		if (obj) {
+			auto ench = obj->As<TESEnchantableForm>();
+			if (ench) {
+				return ench->formEnchanting;
+			}
+		}
+
+		return nullptr;
+	}
+
 	std::optional<double> TESObjectREFR::GetEnchantmentCharge() const
 	{
 		std::optional<double> result;