From d160a8780f705792f824b5a9566c57edc71e1f13 Mon Sep 17 00:00:00 2001
From: powerof3 <powerof3@users.noreply.github.com>
Date: Mon, 25 Nov 2024 06:38:41 +0530
Subject: [PATCH] Add `TESObjectREFR::GetEnchantment`

---
 include/RE/A/ArrowProjectile.h |  2 --
 include/RE/Skyrim.h            |  1 +
 include/RE/T/TESObjectREFR.h   |  2 ++
 src/RE/T/TESObjectREFR.cpp     | 18 ++++++++++++++++++
 4 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/include/RE/A/ArrowProjectile.h b/include/RE/A/ArrowProjectile.h
index 3cfc50648..b10ae3a69 100644
--- 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
index 38197b515..128b3872a 100644
--- 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
index c37a3dad9..cd81fb59f 100644
--- 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
index 526ecbead..ea270d9ed 100644
--- 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;