Change to a registration-based model

Author: sbomer

src/linker/Linker.Steps/IMarkAssemblyStep.cs

@@ -0,0 +1,20 @@
+// 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>
+	/// 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/IMarkHandler.cs

@@ -12,10 +12,11 @@ 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 IMarkAssemblyStep
-	// and is invoked during MarkStep when an assembly gets marked.
+	// 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 MarkAssemblySubStepsDispatcher : IMarkAssemblyStep
+	public abstract class MarkHandlerDispatcher : IMarkHandler
 	{
 		readonly List<ISubStep> substeps;
 
@@ -26,12 +27,12 @@ public abstract class MarkAssemblySubStepsDispatcher : IMarkAssemblyStep
 		List<ISubStep> on_properties;
 		List<ISubStep> on_events;
 
-		protected MarkAssemblySubStepsDispatcher ()
+		protected MarkHandlerDispatcher ()
 		{
 			substeps = new List<ISubStep> ();
 		}
 
-		protected MarkAssemblySubStepsDispatcher (IEnumerable<ISubStep> subSteps)
+		protected MarkHandlerDispatcher (IEnumerable<ISubStep> subSteps)
 		{
 			substeps = new List<ISubStep> (subSteps);
 		}
@@ -41,14 +42,10 @@ public void Add (ISubStep substep)
 			substeps.Add (substep);
 		}
 
-		void IMarkAssemblyStep.Initialize (LinkContext context)
+		void IMarkHandler.Initialize (LinkContext context, MarkContext markContext)
 		{
 			InitializeSubSteps (context);
-		}
-
-		void IMarkAssemblyStep.ProcessAssembly (AssemblyDefinition assembly)
-		{
-			BrowseAssembly (assembly);
+			markContext.RegisterMarkAssemblyAction (assembly => BrowseAssembly (assembly));
 		}
 
 		static bool HasSubSteps (List<ISubStep> substeps) => substeps?.Count > 0;

src/linker/Linker.Steps/MarkAssemblySubStepsDispatcher.cs

@@ -0,0 +1,22 @@
+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);
+	}
+}

src/linker/Linker.Steps/MarkContext.cs

@@ -57,6 +57,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[] {
@@ -183,6 +184,7 @@ public virtual void Process (LinkContext context)
 		{
 			_context = context;
 			_unreachableBlocksOptimizer = new UnreachableBlocksOptimizer (_context);
+			_markContext = new MarkStepContext ();
 
 			Initialize ();
 			Process ();
@@ -191,6 +193,8 @@ public virtual void Process (LinkContext context)
 
 		void Initialize ()
 		{
+			_context.Pipeline.InitializeMarkHandlers (_context, _markContext);
+
 			foreach (AssemblyDefinition assembly in _context.GetAssemblies ())
 				InitializeAssembly (assembly);
 
@@ -1236,7 +1240,8 @@ protected void MarkAssembly (AssemblyDefinition assembly, DependencyInfo reason)
 			if (CheckProcessed (assembly))
 				return;
 
-			_context.Pipeline.ProcessAssembly (assembly);
+			foreach (Action<AssemblyDefinition> handleMarkAssembly in _markContext.MarkAssemblyActions)
+				handleMarkAssembly (assembly);
 
 			var action = _context.Annotations.GetAction (assembly);
 			if (action == AssemblyAction.Copy || action == AssemblyAction.Save) {
@@ -1558,6 +1563,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);
 			MarkType (type.DeclaringType, new DependencyInfo (DependencyKind.DeclaringType, type), type);
 			MarkCustomAttributes (type, new DependencyInfo (DependencyKind.CustomAttribute, type), type);

src/linker/Linker.Steps/MarkStep.cs

@@ -0,0 +1,29 @@
+using System;
+using System.Collections.Generic;
+using Mono.Cecil;
+
+namespace Mono.Linker.Steps
+{
+	public class MarkStepContext : MarkContext
+	{
+
+		public readonly List<Action<AssemblyDefinition>> MarkAssemblyActions;
+		public readonly List<Action<TypeDefinition>> MarkTypeActions;
+
+		public MarkStepContext ()
+		{
+			MarkAssemblyActions = new List<Action<AssemblyDefinition>> ();
+			MarkTypeActions = new List<Action<TypeDefinition>> ();
+		}
+
+		public override void RegisterMarkAssemblyAction (Action<AssemblyDefinition> action)
+		{
+			MarkAssemblyActions.Add (action);
+		}
+
+		public override void RegisterMarkTypeAction (Action<TypeDefinition> action)
+		{
+			MarkTypeActions.Add (action);
+		}
+	}
+}

src/linker/Linker.Steps/MarkStepContext.cs

@@ -788,16 +788,16 @@ protected virtual void AddXmlDependencyRecorder (LinkContext context, string fil
 			context.Tracer.AddRecorder (new XmlDependencyRecorder (context, fileName));
 		}
 
-		protected bool AddMarkAssemblyStep (Pipeline pipeline, string arg)
+		protected bool AddMarkHandler (Pipeline pipeline, string arg)
 		{
 			if (!TryGetCustomAssembly (ref arg, out Assembly custom_assembly))
 				return false;
 
-			var step = ResolveStep<IMarkAssemblyStep> (arg, custom_assembly);
+			var step = ResolveStep<IMarkHandler> (arg, custom_assembly);
 			if (step == null)
 				return false;
 
-			pipeline.AppendMarkAssemblyStep (step);
+			pipeline.AppendMarkHandler (step);
 			return true;
 		}
 
@@ -855,7 +855,7 @@ protected bool AddCustomStep (Pipeline pipeline, string arg)
 					pipeline.AppendStep (customStep);
 					return true;
 				}
-			
+
 				IStep target = FindStep (pipeline, targetName);
 				if (target == null) {
 					context.LogError ($"Pipeline step '{targetName}' could not be found", 1026);
@@ -867,24 +867,24 @@ protected bool AddCustomStep (Pipeline pipeline, string arg)
 				else
 					pipeline.AddStepAfter (target, customStep);
 
-			} else if (typeof (IMarkAssemblyStep).IsAssignableFrom (stepType)) {
+			} else if (typeof (IMarkHandler).IsAssignableFrom (stepType)) {
 
-				var customStep = (IMarkAssemblyStep) Activator.CreateInstance (stepType);
+				var customStep = (IMarkHandler) Activator.CreateInstance (stepType);
 				if (targetName == null) {
-					pipeline.AppendMarkAssemblyStep (customStep);
+					pipeline.AppendMarkHandler (customStep);
 					return true;
 				}
 
-				IMarkAssemblyStep target = FindMarkAssemblyStep (pipeline, targetName);
+				IMarkHandler target = FindMarkHandler (pipeline, targetName);
 				if (target == null) {
 					context.LogError ($"Pipeline step '{targetName}' could not be found", 1026);
 					return false;
 				}
 
 				if (before)
-					pipeline.AddMarkAssemblyStepBefore (target, customStep);
+					pipeline.AddMarkHandlerBefore (target, customStep);
 				else
-					pipeline.AddMarkAssemblyStepAfter (target, customStep);
+					pipeline.AddMarkHandlerAfter (target, customStep);
 
 			} else {
 				context.LogError ($"Custom step '{stepType}' is incompatible with this linker version", 1028);
@@ -905,9 +905,9 @@ static IStep FindStep (Pipeline pipeline, string name)
 			return null;
 		}
 
-		static IMarkAssemblyStep FindMarkAssemblyStep (Pipeline pipeline, string name)
+		static IMarkHandler FindMarkHandler (Pipeline pipeline, string name)
 		{
-			foreach (IMarkAssemblyStep step in pipeline.GetMarkAssemblySteps ()) {
+			foreach (IMarkHandler step in pipeline.MarkHandlers) {
 				Type t = step.GetType ();
 				if (t.Name == name)
 					return step;

src/linker/Linker/Driver.cs

@@ -32,7 +32,7 @@
 using System.IO;
 using Mono.Cecil;
 using Mono.Cecil.Cil;
-
+using Mono.Linker.Steps;
 namespace Mono.Linker
 {
 
@@ -200,6 +200,8 @@ public ISymbolWriterProvider SymbolWriterProvider {
 
 		public string AssemblyListFile { get; set; }
 
+		public List<IMarkHandler> MarkHandlers { get; }
+
 		public LinkContext (Pipeline pipeline, ILogger logger)
 		{
 			_pipeline = pipeline;
@@ -234,6 +236,7 @@ public LinkContext (Pipeline pipeline, ILogger logger)
 			GeneralWarnAsError = false;
 			WarnAsError = new Dictionary<int, bool> ();
 			WarnVersion = WarnVersion.Latest;
+			MarkHandlers = new List<IMarkHandler> ();
 
 			// See https://github.com/mono/linker/issues/612
 			const CodeOptimizations defaultOptimizations =