From a1cd01cd3d1a06ec0e013ec93b7bc88d09fc490d Mon Sep 17 00:00:00 2001
From: Brian Fiete <bfiete@gmail.com>
Date: Fri, 17 Jan 2025 17:32:25 -0800
Subject: [PATCH] Fixed 'ref' and 'in' this parameter for extension methods

---
 IDEHelper/Compiler/BfAutoComplete.cpp   |  3 +++
 IDEHelper/Compiler/BfDefBuilder.cpp     |  7 +++++++
 IDEHelper/Compiler/BfExprEvaluator.cpp  |  9 ++++++++-
 IDEHelper/Tests/src/ExtensionMethods.bf | 25 +++++++++++++++++++++++++
 4 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/IDEHelper/Compiler/BfAutoComplete.cpp b/IDEHelper/Compiler/BfAutoComplete.cpp
index 2ba202b3..3594f962 100644
--- a/IDEHelper/Compiler/BfAutoComplete.cpp
+++ b/IDEHelper/Compiler/BfAutoComplete.cpp
@@ -1336,6 +1336,9 @@ void BfAutoComplete::AddExtensionMethods(BfTypeInstance* targetType, BfTypeInsta
 			continue;
 
 		auto thisType = methodInstance->GetParamType(0);
+		if (thisType->IsRef())
+			thisType = thisType->GetUnderlyingType();
+
 		bool paramValidated = false;
 		if (methodInstance->GetNumGenericParams() > 0)
 		{
diff --git a/IDEHelper/Compiler/BfDefBuilder.cpp b/IDEHelper/Compiler/BfDefBuilder.cpp
index 2aad2665..b8220d40 100644
--- a/IDEHelper/Compiler/BfDefBuilder.cpp
+++ b/IDEHelper/Compiler/BfDefBuilder.cpp
@@ -426,7 +426,14 @@ BfMethodDef* BfDefBuilder::CreateMethodDef(BfMethodDeclaration* methodDeclaratio
 	methodDef->mBody = methodDeclaration->mBody;
 
 	if ((methodDeclaration->mThisToken != NULL) && (!methodDeclaration->mParams.IsEmpty()))
+	{
 		methodDef->mMethodType = BfMethodType_Extension;
+		if (!methodDef->mIsStatic)
+		{
+			methodDef->mIsStatic = true;
+			Fail("Extension methods must be declared 'static'", methodDef->GetRefNode());
+		}
+	}
 
 	HashContext signatureHashCtx;
 	HashNode(signatureHashCtx, methodDeclaration, methodDef->mBody);
diff --git a/IDEHelper/Compiler/BfExprEvaluator.cpp b/IDEHelper/Compiler/BfExprEvaluator.cpp
index 37a3015a..c60e866e 100644
--- a/IDEHelper/Compiler/BfExprEvaluator.cpp
+++ b/IDEHelper/Compiler/BfExprEvaluator.cpp
@@ -2274,6 +2274,12 @@ bool BfMethodMatcher::CheckMethod(BfTypeInstance* targetTypeInstance, BfTypeInst
 			}
 			else
 			{
+				if ((checkMethod->mMethodType == BfMethodType_Extension) && (argIdx == -1))
+				{
+					if ((wantType->IsRef()) && (!argTypedValue.mType->IsRef()))
+						wantType = wantType->GetUnderlyingType();
+				}
+
 				if ((wantType->IsRef()) && (!argTypedValue.mType->IsRef()) &&
 					((mAllowImplicitRef) || (wantType->IsIn())))
 					wantType = wantType->GetUnderlyingType();
@@ -10582,7 +10588,8 @@ BfTypedValue BfExprEvaluator::MatchMethod(BfAstNode* targetSrc, BfMethodBoundExp
 		callTargetType = moduleMethodInstance.mMethodInstance->GetParamType(0);
 		if ((callTargetType->IsRef()) && (target.IsAddr()) && (!target.IsReadOnly()) && (target.mType->IsValueType()))
 		{
-			target = BfTypedValue(target.mValue, mModule->CreateRefType(target.mType));
+			auto refType = (BfRefType*)callTargetType;
+			target = BfTypedValue(target.mValue, mModule->CreateRefType(target.mType, refType->mRefKind));
 		}
 	}
 
diff --git a/IDEHelper/Tests/src/ExtensionMethods.bf b/IDEHelper/Tests/src/ExtensionMethods.bf
index 90238e28..0bf9591a 100644
--- a/IDEHelper/Tests/src/ExtensionMethods.bf
+++ b/IDEHelper/Tests/src/ExtensionMethods.bf
@@ -12,8 +12,21 @@ namespace Tests
 			return 123;
 		}
 
+		public static void AddTo(this ref StructA @this)
+		{
+			@this.mA += 1000;
+		}
+
+		public static void AddTo2(this in StructA @this)
+		{
+#unwarn
+			(&@this).mA += 1000;
+		}
+
 		class ClassA
 		{
+			public int mA = 10000;
+
 			public static void Test()
 			{
 				Test.Assert("Abc".Remove(1.2f, 2.3f) == 123);
@@ -29,6 +42,11 @@ namespace Tests
 
 		}
 
+		struct StructA
+		{
+			public int mA;
+		}
+
 		[Test]
 		public static void TestBasics()
 		{
@@ -46,6 +64,13 @@ namespace Tests
 			float a = 1.2f;
 			float b = 2.3f;
 			Test.Assert(a.CompareIt(b) < 0);
+
+			StructA sa = .() { mA = 123 };
+
+			sa.AddTo();
+			Test.Assert(sa.mA == 1123);
+			sa.AddTo2();
+			Test.Assert(sa.mA == 2123);
 		}
 	}