Add KeptByAttribute to validate an item was kept due to a specific dependency (#3096)

The tests we have in the linker try to ensure that something can only be kept for a certain reason, but it can be hard to ensure it's not being kept for another reason. For example, sometimes we use typeof(TestType) to mark a type, but that also makes it relevant to variant casting, which can cause parts of TestType to be kept for unintended reasons.

This PR creates the KeptByAttribute which accepts a fully qualified string in Cecil format to indicate depenency provider (the thing that is creating the dependency that marks the item with the attribute), as well as a string representing the DependencyKind to specify the "reason" there was a dependency between the two. It also can accept a System.Type, or a System.Type + name of the member to indicate the dependency provider.

7 files changed, 128 insertions, 12 deletions

diff --git a/test/Mono.Linker.Tests.Cases.Expectations/Assertions/KeptByAttribute.cs b/test/Mono.Linker.Tests.Cases.Expectations/Assertions/KeptByAttribute.cs
new file mode 100644
index 000000000..d7871d9ca
--- /dev/null
+++ b/test/Mono.Linker.Tests.Cases.Expectations/Assertions/KeptByAttribute.cs
@@ -0,0 +1,35 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+
+namespace Mono.Linker.Tests.Cases.Expectations.Assertions
+{
+	[AttributeUsage (AttributeTargets.All, Inherited = false)]
+	public class KeptByAttribute : KeptAttribute
+	{
+		private KeptByAttribute () { }
+
+		/// <summary>
+		/// Place on an type member to indicate that the linker should log that the member is kept as a depenendency of <paramref name="dependencyProvider"/> with reason <paramref name="reason"/>.
+		/// </summary>
+		/// <param name="dependencyProvider">Cecil's FullName property of the item that provides the dependency that keeps the item</param>
+		/// <param name="reason">The string representation of the DependencyKind that is recorded as the reason for the dependency</param>
+		public KeptByAttribute (string dependencyProvider, string reason) { }
+
+		/// <summary>
+		/// Place on an type member to indicate that the linker should log that the member is kept as a depenendency of <paramref name="dependencyProviderType"/> with reason <paramref name="reason"/>.
+		/// </summary>
+		/// <param name="dependencyProviderType">The type that is providing the dependency that keeps the item</param>
+		/// <param name="reason">The string representation of the DependencyKind that is recorded as the reason for the dependency</param>
+		public KeptByAttribute (Type dependencyProviderType, string reason) { }
+
+		/// <summary>
+		/// Place on an type member to indicate that the linker should log that the member is kept as a depenendency of <paramref name="dependencyProviderType"/>.<paramref name="memberName"/> with reason <paramref name="reason"/>.
+		/// </summary>
+		/// <param name="dependencyProviderType">The declaring type of the member that is providing the dependency that keeps the item</param>
+		/// <param name="memberName">Cecil's 'Name' property of the member that provides the dependency that keeps the item</param>
+		/// <param name="reason">The string representation of the DependencyKind that is recorded as the reason for the dependency</param>
+		public KeptByAttribute (Type dependencyProviderType, string memberName, string reason) { }
+	}
+}
diff --git a/test/Mono.Linker.Tests.Cases/Inheritance.VirtualMethods/VirtualMethodGetsPreservedIfBaseMethodGetsInvoked.cs b/test/Mono.Linker.Tests.Cases/Inheritance.VirtualMethods/VirtualMethodGetsPreservedIfBaseMethodGetsInvoked.cs
index 70323315c..8b876f777 100644
--- a/test/Mono.Linker.Tests.Cases/Inheritance.VirtualMethods/VirtualMethodGetsPreservedIfBaseMethodGetsInvoked.cs
+++ b/test/Mono.Linker.Tests.Cases/Inheritance.VirtualMethods/VirtualMethodGetsPreservedIfBaseMethodGetsInvoked.cs
@@ -23,7 +23,7 @@ namespace Mono.Linker.Tests.Cases.Inheritance.VirtualMethods
 		[KeptBaseType (typeof (B))]
 		class A : B
 		{
-			[Kept]
+			[KeptBy (typeof (A), "OverrideOnInstantiatedType")]
 			public override void Foo ()
 			{
 			}
diff --git a/test/Mono.Linker.Tests.Cases/Inheritance.VirtualMethods/VirtualMethodGetsStrippedIfImplementingMethodGetsInvokedDirectly.cs b/test/Mono.Linker.Tests.Cases/Inheritance.VirtualMethods/VirtualMethodGetsStrippedIfImplementingMethodGetsInvokedDirectly.cs
index 76457ce53..eaae2a4aa 100644
--- a/test/Mono.Linker.Tests.Cases/Inheritance.VirtualMethods/VirtualMethodGetsStrippedIfImplementingMethodGetsInvokedDirectly.cs
+++ b/test/Mono.Linker.Tests.Cases/Inheritance.VirtualMethods/VirtualMethodGetsStrippedIfImplementingMethodGetsInvokedDirectly.cs
@@ -12,7 +12,9 @@ namespace Mono.Linker.Tests.Cases.Inheritance.VirtualMethods
 		[KeptMember (".ctor()")]
 		class B
 		{
-			[Kept] // TODO: Would be nice to be removed
+			// TODO: Would be nice to be removed
+			// https://github.com/dotnet/linker/issues/3097
+			[KeptBy (typeof (A), nameof (A.Foo), "BaseMethod")]
 			public virtual void Foo ()
 			{
 			}
@@ -22,7 +24,10 @@ namespace Mono.Linker.Tests.Cases.Inheritance.VirtualMethods
 		[KeptBaseType (typeof (B))]
 		class A : B
 		{
-			[Kept]
+			// Bug: https://github.com/dotnet/linker/issues/3078
+			// Linker should mark for DirectCall as well as OverrideOnInstantiatedType, not just OverrideOnInstantiatedType
+			//[KeptBy (typeof(A), nameof(Foo), DependencyKind.DirectCall)]
+			[KeptBy (typeof (A), "OverrideOnInstantiatedType")]
 			public override void Foo ()
 			{
 			}
diff --git a/test/Mono.Linker.Tests/TestCasesRunner/AssemblyChecker.cs b/test/Mono.Linker.Tests/TestCasesRunner/AssemblyChecker.cs
index 41210172f..26ad1309b 100644
--- a/test/Mono.Linker.Tests/TestCasesRunner/AssemblyChecker.cs
+++ b/test/Mono.Linker.Tests/TestCasesRunner/AssemblyChecker.cs
@@ -16,6 +16,7 @@ namespace Mono.Linker.Tests.TestCasesRunner
 	public class AssemblyChecker
 	{
 		readonly AssemblyDefinition originalAssembly, linkedAssembly;
+		readonly LinkedTestCaseResult linkedTestCase;
 
 		HashSet<string> linkedMembers;
 		readonly HashSet<string> verifiedGeneratedFields = new HashSet<string> ();
@@ -23,10 +24,11 @@ namespace Mono.Linker.Tests.TestCasesRunner
 		readonly HashSet<string> verifiedGeneratedTypes = new HashSet<string> ();
 		bool checkNames;
 
-		public AssemblyChecker (AssemblyDefinition original, AssemblyDefinition linked)
+		public AssemblyChecker (AssemblyDefinition original, AssemblyDefinition linked, LinkedTestCaseResult linkedTestCase)
 		{
 			this.originalAssembly = original;
 			this.linkedAssembly = linked;
+			this.linkedTestCase = linkedTestCase;
 
 			checkNames = original.MainModule.GetTypeReferences ().Any (attr =>
 				attr.Name == nameof (RemovedNameValueAttribute));
@@ -44,6 +46,7 @@ namespace Mono.Linker.Tests.TestCasesRunner
 
 			VerifyResources (originalAssembly, linkedAssembly);
 			VerifyReferences (originalAssembly, linkedAssembly);
+			VerifyKeptByAttributes (originalAssembly, originalAssembly.FullName);
 
 			linkedMembers = new HashSet<string> (linkedAssembly.MainModule.AllMembers ().Select (s => {
 				return s.FullName;
@@ -139,11 +142,76 @@ namespace Mono.Linker.Tests.TestCasesRunner
 			}
 		}
 
+		/// <summary>
+		/// Validates that all <see cref="KeptByAttribute"/> instances on a member are valid (i.e. the linker recorded a marked dependency described in the attribute)
+		/// </summary>
+		void VerifyKeptByAttributes (IMemberDefinition src, IMemberDefinition linked)
+		{
+			foreach (var keptByAttribute in src.CustomAttributes.Where (ca => ca.AttributeType.IsTypeOf<KeptByAttribute> ()))
+				VerifyKeptByAttribute (linked.FullName, keptByAttribute);
+		}
+
+		/// <summary>
+		/// Validates that all <see cref="KeptByAttribute"/> instances on an attribute provider are valid (i.e. the linker recorded a marked dependency described in the attribute)
+		/// <paramref name="src"/> is the attribute provider that may have a <see cref="KeptByAttribute"/>, and <paramref name="attributeProviderFullName"/> is the 'FullName' of <paramref name="src"/>.
+		/// </summary>
+		void VerifyKeptByAttributes (ICustomAttributeProvider src, string attributeProviderFullName)
+		{
+			foreach (var keptByAttribute in src.CustomAttributes.Where (ca => ca.AttributeType.IsTypeOf<KeptByAttribute> ()))
+				VerifyKeptByAttribute (attributeProviderFullName, keptByAttribute);
+		}
+
+		void VerifyKeptByAttribute (string keptAttributeProviderName, CustomAttribute attribute)
+		{
+			// public KeptByAttribute (string dependencyProvider, string reason) { }
+			// public KeptByAttribute (Type dependencyProvider, string reason) { }
+			// public KeptByAttribute (Type dependencyProvider, string memberName, string reason) { }
+
+			Assert.AreEqual (nameof (KeptByAttribute), attribute.AttributeType.Name);
+
+			// Create the expected TestDependencyRecorder.Dependency that should be in the recorded dependencies
+			TestDependencyRecorder.Dependency expectedDependency = new ();
+			expectedDependency.Target = keptAttributeProviderName;
+			expectedDependency.Marked = true;
+			if (attribute.ConstructorArguments.Count == 2) {
+				// public KeptByAttribute (string dependencyProvider, string reason) { }
+				// public KeptByAttribute (Type dependencyProvider, string reason) { }
+				if (attribute.ConstructorArguments[0].Type.IsTypeOf<string> ())
+					expectedDependency.Source = (string) attribute.ConstructorArguments[0].Value;
+				else if (attribute.ConstructorArguments[0].Type.IsTypeOf<Type> ())
+					expectedDependency.Source = ((TypeDefinition) attribute.ConstructorArguments[0].Value).FullName;
+				else
+					throw new NotImplementedException ("Unexpected KeptByAttribute ctor variant");
+
+				expectedDependency.DependencyKind = (string) attribute.ConstructorArguments[1].Value;
+			} else if (attribute.ConstructorArguments.Count == 3) {
+				// public KeptByAttribute (Type dependencyProvider, string memberName, string reason) { }
+				if (!attribute.ConstructorArguments[0].Type.IsTypeOf<Type> ())
+					throw new NotImplementedException ("Unexpected KeptByAttribute ctor variant");
+				var type = (TypeDefinition) attribute.ConstructorArguments[0].Value;
+				string memberName = (string) attribute.ConstructorArguments[1].Value;
+				var memberDefinition = type.AllMembers ().Where (m => m.Name == memberName).Single ();
+				expectedDependency.Source = memberDefinition.FullName;
+				expectedDependency.DependencyKind = (string) attribute.ConstructorArguments[2].Value;
+			} else {
+				throw new NotImplementedException ("Unexpected KeptByAttribute ctor variant");
+			}
+
+			foreach (var dep in this.linkedTestCase.Customizations.DependencyRecorder.Dependencies) {
+				if (dep == expectedDependency) {
+					return;
+				}
+			}
+			string errorMessage = $"{keptAttributeProviderName} was expected to be kept by {expectedDependency.Source} with reason {expectedDependency.DependencyKind.ToString ()}.";
+			Assert.Fail (errorMessage);
+		}
+
 		protected virtual void VerifyTypeDefinitionKept (TypeDefinition original, TypeDefinition linked)
 		{
 			if (linked == null)
 				Assert.Fail ($"Type `{original}' should have been kept");
 
+			VerifyKeptByAttributes (original, linked);
 			if (!original.IsInterface)
 				VerifyBaseType (original, linked);
 
@@ -313,6 +381,7 @@ namespace Mono.Linker.Tests.TestCasesRunner
 
 			Assert.AreEqual (src?.Constant, linked?.Constant, $"Field `{src}' value");
 
+			VerifyKeptByAttributes (src, linked);
 			VerifyPseudoAttributes (src, linked);
 			VerifyCustomAttributes (src, linked);
 		}
@@ -335,6 +404,7 @@ namespace Mono.Linker.Tests.TestCasesRunner
 
 			Assert.AreEqual (src?.Constant, linked?.Constant, $"Property `{src}' value");
 
+			VerifyKeptByAttributes (src, linked);
 			VerifyPseudoAttributes (src, linked);
 			VerifyCustomAttributes (src, linked);
 		}
@@ -367,6 +437,7 @@ namespace Mono.Linker.Tests.TestCasesRunner
 				linkedMembers.Remove (src.RemoveMethod.FullName);
 			}
 
+			VerifyKeptByAttributes (src, linked);
 			VerifyPseudoAttributes (src, linked);
 			VerifyCustomAttributes (src, linked);
 		}
@@ -430,6 +501,7 @@ namespace Mono.Linker.Tests.TestCasesRunner
 			VerifySecurityAttributes (src, linked);
 			VerifyArrayInitializers (src, linked);
 			VerifyMethodBody (src, linked);
+			VerifyKeptByAttributes (src, linked);
 		}
 
 		protected virtual void VerifyMethodBody (MethodDefinition src, MethodDefinition linked)
@@ -989,7 +1061,7 @@ namespace Mono.Linker.Tests.TestCasesRunner
 
 		protected virtual bool ShouldBeKept<T> (T member, string signature = null) where T : MemberReference, ICustomAttributeProvider
 		{
-			if (member.HasAttribute (nameof (KeptAttribute)))
+			if (member.HasAttribute (nameof (KeptAttribute)) || member.HasAttribute (nameof (KeptByAttribute)))
 				return true;
 
 			ICustomAttributeProvider cap = (ICustomAttributeProvider) member.DeclaringType;
diff --git a/test/Mono.Linker.Tests/TestCasesRunner/ResultChecker.cs b/test/Mono.Linker.Tests/TestCasesRunner/ResultChecker.cs
index 4a7de6de6..af4c6e22c 100644
--- a/test/Mono.Linker.Tests/TestCasesRunner/ResultChecker.cs
+++ b/test/Mono.Linker.Tests/TestCasesRunner/ResultChecker.cs
@@ -94,7 +94,7 @@ namespace Mono.Linker.Tests.TestCasesRunner
 					PerformOutputSymbolChecks (original, linkResult.OutputAssemblyPath.Parent);
 
 					if (!HasAttribute (original.MainModule.GetType (linkResult.TestCase.ReconstructedFullTypeName), nameof (SkipKeptItemsValidationAttribute))) {
-						CreateAssemblyChecker (original, linked).Verify ();
+						CreateAssemblyChecker (original, linked, linkResult).Verify ();
 					}
 				}
 
@@ -106,9 +106,9 @@ namespace Mono.Linker.Tests.TestCasesRunner
 			}
 		}
 
-		protected virtual AssemblyChecker CreateAssemblyChecker (AssemblyDefinition original, AssemblyDefinition linked)
+		protected virtual AssemblyChecker CreateAssemblyChecker (AssemblyDefinition original, AssemblyDefinition linked, LinkedTestCaseResult linkedTestCase)
 		{
-			return new AssemblyChecker (original, linked);
+			return new AssemblyChecker (original, linked, linkedTestCase);
 		}
 
 		void InitializeResolvers (LinkedTestCaseResult linkedResult)
diff --git a/test/Mono.Linker.Tests/TestCasesRunner/TestCaseMetadataProvider.cs b/test/Mono.Linker.Tests/TestCasesRunner/TestCaseMetadataProvider.cs
index a3290c4f6..6f0ba6d0c 100644
--- a/test/Mono.Linker.Tests/TestCasesRunner/TestCaseMetadataProvider.cs
+++ b/test/Mono.Linker.Tests/TestCasesRunner/TestCaseMetadataProvider.cs
@@ -89,8 +89,10 @@ namespace Mono.Linker.Tests.TestCasesRunner
 
 		public virtual void CustomizeLinker (LinkerDriver linker, LinkerCustomizations customizations)
 		{
-			if (_testCaseTypeDefinition.CustomAttributes.Any (attr =>
-				attr.AttributeType.Name == nameof (DependencyRecordedAttribute))) {
+			if (!_testCaseTypeDefinition.CustomAttributes.Any (a => a.AttributeType.IsTypeOf<SkipKeptItemsValidationAttribute> ())
+				|| _testCaseTypeDefinition.CustomAttributes.Any (attr =>
+				attr.AttributeType.Name == nameof (DependencyRecordedAttribute)
+				|| attr.AttributeType.Name == nameof (KeptByAttribute))) {
 				customizations.DependencyRecorder = new TestDependencyRecorder ();
 				customizations.CustomizeContext += context => {
 					context.Tracer.AddRecorder (customizations.DependencyRecorder);
diff --git a/test/Mono.Linker.Tests/TestCasesRunner/TestDependencyRecorder.cs b/test/Mono.Linker.Tests/TestCasesRunner/TestDependencyRecorder.cs
index cbf41bbe9..643832414 100644
--- a/test/Mono.Linker.Tests/TestCasesRunner/TestDependencyRecorder.cs
+++ b/test/Mono.Linker.Tests/TestCasesRunner/TestDependencyRecorder.cs
@@ -7,11 +7,12 @@ namespace Mono.Linker.Tests.TestCasesRunner
 {
 	public class TestDependencyRecorder : IDependencyRecorder
 	{
-		public struct Dependency
+		public record struct Dependency
 		{
 			public string Source;
 			public string Target;
 			public bool Marked;
+			public string DependencyKind;
 		}
 
 		public List<Dependency> Dependencies = new List<Dependency> ();
@@ -30,7 +31,8 @@ namespace Mono.Linker.Tests.TestCasesRunner
 			Dependencies.Add (new Dependency () {
 				Source = reason.Source?.ToString (),
 				Target = target.ToString (),
-				Marked = marked
+				Marked = marked,
+				DependencyKind = reason.Kind.ToString ()
 			});
 		}