Add a custom steps API to run logic on mark (#1774)

This adds a new API which lets custom steps register an action to be called when a member gets marked.

This includes a helper, MarkSubStepsDispatcher, which behaves similarly to SubStepDispatcher and is used to scan in marked assemblies.

* PR feedback

- Rename IPerAssemblyStep to IMarkAssemblyStep
- Keep SubStepsDispatcher

* Use --custom-step for both step interfaces

* Change to a registration-based model

* Rename test file

* Cleanup

* Add another handler for marked methods

* PR feedback

- Add test for IMarkHandler
- Add missing license header

* Add a warning to the API docs

* PR feedback

- MarkHandlerDispatcher -> MarkSubStepsDispatcher
- Make Initialize virtual so that it can be called by derived types

* Fix analyzer warning

* PR feedback

- Always initialize substep lists
- Actually make Initialize virtual (and test this)
- Remove default ctor and Add method

src/linker/Linker.Steps/IMarkHandler.cs

@@ -0,0 +1,21 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+namespace Mono.Linker.Steps
+{
+
+	/// <summary>
+	/// This API supports the product infrastructure and is not intended to be used directly from your code.
+	/// Extensibility point for custom logic that run during MarkStep, for marked members.
+	/// </summary>
+	public interface IMarkHandler
+	{
+		/// <summary>
+		/// Initialize is called at the beginning of MarkStep. This should be
+		/// used to perform global setup, and register callbacks through the
+		/// MarkContext.Register* methods) to be called when pieces of IL are marked.
+		/// </summary>
+		void Initialize (LinkContext context, MarkContext markContext);
+	}
+}

src/linker/Linker.Steps/MarkContext.cs

@@ -0,0 +1,31 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using Mono.Cecil;
+
+namespace Mono.Linker.Steps
+{
+	/// <summary>
+	/// Context which can be used to register actions to call during MarkStep
+	/// when various members are marked.
+	/// </summary>
+	public abstract class MarkContext
+	{
+		/// <summary>
+		/// Register a callback that will be invoked once for each marked assembly.
+		/// </summary>
+		public abstract void RegisterMarkAssemblyAction (Action<AssemblyDefinition> action);
+
+		/// <summary>
+		/// Register a callback that will be invoked once for each marked type.
+		/// </summary>
+		public abstract void RegisterMarkTypeAction (Action<TypeDefinition> action);
+
+		/// <summary>
+		/// Register a callback that will be invoked once for each marked method.
+		/// </summary>
+		public abstract void RegisterMarkMethodAction (Action<MethodDefinition> action);
+	}
+}

src/linker/Linker.Steps/MarkStep.cs

@@ -56,6 +56,7 @@ public partial class MarkStep : IStep
 
 		readonly List<(TypeDefinition Type, MethodBody Body, Instruction Instr)> _pending_isinst_instr;
 		UnreachableBlocksOptimizer _unreachableBlocksOptimizer;
+		MarkStepContext _markContext;
 
 #if DEBUG
 		static readonly DependencyKind[] _entireTypeReasons = new DependencyKind[] {
@@ -190,6 +191,7 @@ public virtual void Process (LinkContext context)
 		{
 			_context = context;
 			_unreachableBlocksOptimizer = new UnreachableBlocksOptimizer (_context);
+			_markContext = new MarkStepContext ();
 
 			Initialize ();
 			Process ();
@@ -199,6 +201,7 @@ public virtual void Process (LinkContext context)
 		void Initialize ()
 		{
 			InitializeCorelibAttributeXml ();
+			_context.Pipeline.InitializeMarkHandlers (_context, _markContext);
 
 			ProcessMarkedPending ();
 		}
@@ -1296,6 +1299,9 @@ protected void MarkAssembly (AssemblyDefinition assembly, DependencyInfo reason)
 
 			EmbeddedXmlInfo.ProcessDescriptors (assembly, _context);
 
+			foreach (Action<AssemblyDefinition> handleMarkAssembly in _markContext.MarkAssemblyActions)
+				handleMarkAssembly (assembly);
+
 			// Security attributes do not respect the attributes XML
 			if (_context.StripSecurity)
 				RemoveSecurity.ProcessAssembly (assembly, _context);
@@ -1616,6 +1622,10 @@ protected internal virtual TypeDefinition MarkType (TypeReference reference, Dep
 				return null;
 
 			MarkModule (type.Scope as ModuleDefinition, new DependencyInfo (DependencyKind.ScopeOfType, type));
+
+			foreach (Action<TypeDefinition> handleMarkType in _markContext.MarkTypeActions)
+				handleMarkType (type);
+
 			MarkType (type.BaseType, new DependencyInfo (DependencyKind.BaseType, type), type);
 			if (type.DeclaringType != null)
 				MarkType (type.DeclaringType, new DependencyInfo (DependencyKind.DeclaringType, type), type);
@@ -2655,6 +2665,9 @@ protected virtual void ProcessMethod (MethodDefinition method, in DependencyInfo
 
 			_unreachableBlocksOptimizer.ProcessMethod (method);
 
+			foreach (Action<MethodDefinition> handleMarkMethod in _markContext.MarkMethodActions)
+				handleMarkMethod (method);
+
 			if (!markedForCall)
 				MarkType (method.DeclaringType, new DependencyInfo (DependencyKind.DeclaringType, method), method);
 			MarkCustomAttributes (method, new DependencyInfo (DependencyKind.CustomAttribute, method), method);

src/linker/Linker.Steps/MarkStepContext.cs

@@ -0,0 +1,40 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Generic;
+using Mono.Cecil;
+
+namespace Mono.Linker.Steps
+{
+	public class MarkStepContext : MarkContext
+	{
+
+		public List<Action<AssemblyDefinition>> MarkAssemblyActions { get; }
+		public List<Action<TypeDefinition>> MarkTypeActions { get; }
+		public List<Action<MethodDefinition>> MarkMethodActions { get; }
+
+		public MarkStepContext ()
+		{
+			MarkAssemblyActions = new List<Action<AssemblyDefinition>> ();
+			MarkTypeActions = new List<Action<TypeDefinition>> ();
+			MarkMethodActions = new List<Action<MethodDefinition>> ();
+		}
+
+		public override void RegisterMarkAssemblyAction (Action<AssemblyDefinition> action)
+		{
+			MarkAssemblyActions.Add (action);
+		}
+
+		public override void RegisterMarkTypeAction (Action<TypeDefinition> action)
+		{
+			MarkTypeActions.Add (action);
+		}
+
+		public override void RegisterMarkMethodAction (Action<MethodDefinition> action)
+		{
+			MarkMethodActions.Add (action);
+		}
+	}
+}

src/linker/Linker.Steps/MarkSubStepsDispatcher.cs

@@ -0,0 +1,179 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Collections.Generic;
+
+using Mono.Cecil;
+using Mono.Collections.Generic;
+
+namespace Mono.Linker.Steps
+{
+	//
+	// Dispatcher for SubSteps which only need to run on marked assemblies.
+	// This simplifies the implementation of linker custom steps, in the same
+	// way that SubStepsDispatcher does, but it implements IMarkHandler
+	// and registers a callback that gets invoked during MarkStep when an
+	// assembly gets marked.
+	//
+	public abstract class MarkSubStepsDispatcher : IMarkHandler
+	{
+		readonly List<ISubStep> substeps;
+
+		List<ISubStep> on_assemblies;
+		List<ISubStep> on_types;
+		List<ISubStep> on_fields;
+		List<ISubStep> on_methods;
+		List<ISubStep> on_properties;
+		List<ISubStep> on_events;
+
+		protected MarkSubStepsDispatcher (IEnumerable<ISubStep> subSteps)
+		{
+			substeps = new List<ISubStep> (subSteps);
+		}
+
+		public virtual void Initialize (LinkContext context, MarkContext markContext)
+		{
+			InitializeSubSteps (context);
+			markContext.RegisterMarkAssemblyAction (assembly => BrowseAssembly (assembly));
+		}
+
+		static bool HasSubSteps (List<ISubStep> substeps) => substeps?.Count > 0;
+
+		void BrowseAssembly (AssemblyDefinition assembly)
+		{
+			CategorizeSubSteps (assembly);
+
+			if (HasSubSteps (on_assemblies))
+				DispatchAssembly (assembly);
+
+			if (!ShouldDispatchTypes ())
+				return;
+
+			BrowseTypes (assembly.MainModule.Types);
+		}
+
+		bool ShouldDispatchTypes ()
+		{
+			return HasSubSteps (on_types)
+				|| HasSubSteps (on_fields)
+				|| HasSubSteps (on_methods)
+				|| HasSubSteps (on_properties)
+				|| HasSubSteps (on_events);
+		}
+
+		void BrowseTypes (Collection<TypeDefinition> types)
+		{
+			foreach (TypeDefinition type in types) {
+				DispatchType (type);
+
+				if (type.HasFields && HasSubSteps (on_fields)) {
+					foreach (FieldDefinition field in type.Fields)
+						DispatchField (field);
+				}
+
+				if (type.HasMethods && HasSubSteps (on_methods)) {
+					foreach (MethodDefinition method in type.Methods)
+						DispatchMethod (method);
+				}
+
+				if (type.HasProperties && HasSubSteps (on_properties)) {
+					foreach (PropertyDefinition property in type.Properties)
+						DispatchProperty (property);
+				}
+
+				if (type.HasEvents && HasSubSteps (on_events)) {
+					foreach (EventDefinition @event in type.Events)
+						DispatchEvent (@event);
+				}
+
+				if (type.HasNestedTypes)
+					BrowseTypes (type.NestedTypes);
+			}
+		}
+
+		void DispatchAssembly (AssemblyDefinition assembly)
+		{
+			foreach (var substep in on_assemblies) {
+				substep.ProcessAssembly (assembly);
+			}
+		}
+
+		void DispatchType (TypeDefinition type)
+		{
+			foreach (var substep in on_types) {
+				substep.ProcessType (type);
+			}
+		}
+
+		void DispatchField (FieldDefinition field)
+		{
+			foreach (var substep in on_fields) {
+				substep.ProcessField (field);
+			}
+		}
+
+		void DispatchMethod (MethodDefinition method)
+		{
+			foreach (var substep in on_methods) {
+				substep.ProcessMethod (method);
+			}
+		}
+
+		void DispatchProperty (PropertyDefinition property)
+		{
+			foreach (var substep in on_properties) {
+				substep.ProcessProperty (property);
+			}
+		}
+
+		void DispatchEvent (EventDefinition @event)
+		{
+			foreach (var substep in on_events) {
+				substep.ProcessEvent (@event);
+			}
+		}
+
+		void InitializeSubSteps (LinkContext context)
+		{
+			foreach (var substep in substeps)
+				substep.Initialize (context);
+		}
+
+		void CategorizeSubSteps (AssemblyDefinition assembly)
+		{
+			on_assemblies = new List<ISubStep> ();
+			on_types = new List<ISubStep> ();
+			on_fields = new List<ISubStep> ();
+			on_methods = new List<ISubStep> ();
+			on_properties = new List<ISubStep> ();
+			on_events = new List<ISubStep> ();
+
+			foreach (var substep in substeps)
+				CategorizeSubStep (substep, assembly);
+		}
+
+		void CategorizeSubStep (ISubStep substep, AssemblyDefinition assembly)
+		{
+			if (!substep.IsActiveFor (assembly))
+				return;
+
+			CategorizeTarget (substep, SubStepTargets.Assembly, on_assemblies);
+			CategorizeTarget (substep, SubStepTargets.Type, on_types);
+			CategorizeTarget (substep, SubStepTargets.Field, on_fields);
+			CategorizeTarget (substep, SubStepTargets.Method, on_methods);
+			CategorizeTarget (substep, SubStepTargets.Property, on_properties);
+			CategorizeTarget (substep, SubStepTargets.Event, on_events);
+		}
+
+		static void CategorizeTarget (ISubStep substep, SubStepTargets target, List<ISubStep> list)
+		{
+			if (!Targets (substep, target))
+				return;
+
+			list.Add (substep);
+		}
+
+		static bool Targets (ISubStep substep, SubStepTargets target) => (substep.Targets & target) == target;
+	}
+}