PR feedback

- Rename IPerAssemblyStep to IMarkAssemblyStep
- Keep SubStepsDispatcher

src/linker/Linker.Steps/IMarkAssemblyStep.cs

@@ -0,0 +1,33 @@
+// 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 Mono.Cecil;
+
+namespace Mono.Linker.Steps
+{
+
+	/// <summary>
+	/// Extensibility point for custom steps that run during MarkStep,
+	/// for every marked assembly.
+	/// </summary>
+	public interface IMarkAssemblyStep
+	{
+		/// <summary>
+		/// Initialize the per-assembly step. The intended use is just to do
+		/// simple setup of global state. The main processing logic should be
+		/// done by ProcessAssembly. Initialize is called at the beginning of
+		/// pipeline processing, before any of the normal pipeline steps (IStep) run.
+		/// </summary>
+		void Initialize (LinkContext context);
+
+		/// <summary>
+		/// Process an assembly. This should perform the main logic of the step,
+		/// including any modifications to the assembly or to the global Annotations
+		/// state. ProcessAssembly is called exactly once, when an assembly is marked.
+		/// It is called before the main MarkStep processing happens for the assembly,
+		/// but other steps may already have marked parts of the assembly through Annotations.
+		/// </summary>
+		void ProcessAssembly (AssemblyDefinition assembly);
+	}
+}

src/linker/Linker.Steps/MarkAssemblySubStepsDispatcher.cs

@@ -0,0 +1,195 @@
+// 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 IMarkAssemblyStep
+	// and is invoked during MarkStep when an assembly gets marked.
+	//
+	public abstract class MarkAssemblySubStepsDispatcher : IMarkAssemblyStep
+	{
+		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 MarkAssemblySubStepsDispatcher ()
+		{
+			substeps = new List<ISubStep> ();
+		}
+
+		protected MarkAssemblySubStepsDispatcher (IEnumerable<ISubStep> subSteps)
+		{
+			substeps = new List<ISubStep> (subSteps);
+		}
+
+		public void Add (ISubStep substep)
+		{
+			substeps.Add (substep);
+		}
+
+		void IMarkAssemblyStep.Initialize (LinkContext context)
+		{
+			InitializeSubSteps (context);
+		}
+
+		void IMarkAssemblyStep.ProcessAssembly (AssemblyDefinition 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 = null;
+			on_types = null;
+			on_fields = null;
+			on_methods = null;
+			on_properties = null;
+			on_events = null;
+
+			foreach (var substep in substeps)
+				CategorizeSubStep (substep, assembly);
+		}
+
+		void CategorizeSubStep (ISubStep substep, AssemblyDefinition assembly)
+		{
+			if (!substep.IsActiveFor (assembly))
+				return;
+
+			CategorizeTarget (substep, SubStepTargets.Assembly, ref on_assemblies);
+			CategorizeTarget (substep, SubStepTargets.Type, ref on_types);
+			CategorizeTarget (substep, SubStepTargets.Field, ref on_fields);
+			CategorizeTarget (substep, SubStepTargets.Method, ref on_methods);
+			CategorizeTarget (substep, SubStepTargets.Property, ref on_properties);
+			CategorizeTarget (substep, SubStepTargets.Event, ref on_events);
+		}
+
+		static void CategorizeTarget (ISubStep substep, SubStepTargets target, ref List<ISubStep> list)
+		{
+			if (!Targets (substep, target))
+				return;
+
+			if (list == null)
+				list = new List<ISubStep> ();
+
+			list.Add (substep);
+		}
+
+		static bool Targets (ISubStep substep, SubStepTargets target) => (substep.Targets & target) == target;
+	}
+}

src/linker/Linker.Steps/SubStepsDispatcher.cs

@@ -14,7 +14,7 @@ namespace Mono.Linker.Steps
 	// consist of multiple steps. It simplifies their implementation as well as the
 	// way how to hook them into the pipeline of existing steps.
 	//
-	public abstract class SubStepsDispatcher : IPerAssemblyStep
+	public abstract class SubStepsDispatcher : IStep
 	{
 		readonly List<ISubStep> substeps;
 
@@ -40,29 +40,28 @@ public void Add (ISubStep substep)
 			substeps.Add (substep);
 		}
 
-		void IPerAssemblyStep.Initialize (LinkContext context)
+		void IStep.Process (LinkContext context)
 		{
 			InitializeSubSteps (context);
-		}
 
-		void IPerAssemblyStep.ProcessAssembly (AssemblyDefinition assembly)
-		{
-			BrowseAssembly (assembly);
+			BrowseAssemblies (context.GetAssemblies ());
 		}
 
 		static bool HasSubSteps (List<ISubStep> substeps) => substeps?.Count > 0;
 
-		void BrowseAssembly (AssemblyDefinition assembly)
+		void BrowseAssemblies (IEnumerable<AssemblyDefinition> assemblies)
 		{
-			CategorizeSubSteps (assembly);
+			foreach (var assembly in assemblies) {
+				CategorizeSubSteps (assembly);
 
-			if (HasSubSteps (on_assemblies))
-				DispatchAssembly (assembly);
+				if (HasSubSteps (on_assemblies))
+					DispatchAssembly (assembly);
 
-			if (!ShouldDispatchTypes ())
-				return;
+				if (!ShouldDispatchTypes ())
+					continue;
 
-			BrowseTypes (assembly.MainModule.Types);
+				BrowseTypes (assembly.MainModule.Types);
+			}
 		}
 
 		bool ShouldDispatchTypes ()

src/linker/Linker/Driver.cs

@@ -267,7 +267,7 @@ protected int SetupContext (ILogger customLogger = null)
 
 						continue;
 
-					case "--per-assembly-step":
+					case "--mark-assembly-step":
 						if (!GetStringParam (token, l => per_assembly_steps.Push (l)))
 							return -1;
 
@@ -699,7 +699,7 @@ protected int SetupContext (ILogger customLogger = null)
 			}
 
 			foreach (string per_assembly_step in per_assembly_steps) {
-				if (!AddPerAssemblyStep (p, per_assembly_step))
+				if (!AddMarkAssemblyStep (p, per_assembly_step))
 					return -1;
 			}
 
@@ -801,16 +801,16 @@ protected virtual void AddXmlDependencyRecorder (LinkContext context, string fil
 			context.Tracer.AddRecorder (new XmlDependencyRecorder (context, fileName));
 		}
 
-		protected bool AddPerAssemblyStep (Pipeline pipeline, string arg)
+		protected bool AddMarkAssemblyStep (Pipeline pipeline, string arg)
 		{
 			if (!TryGetCustomAssembly (ref arg, out Assembly custom_assembly))
 				return false;
 
-			var step = ResolveStep<IPerAssemblyStep> (arg, custom_assembly);
+			var step = ResolveStep<IMarkAssemblyStep> (arg, custom_assembly);
 			if (step == null)
 				return false;
 
-			pipeline.AppendPerAssemblyStep (step);
+			pipeline.AppendMarkAssemblyStep (step);
 			return true;
 		}
 
@@ -1140,7 +1140,7 @@ static void Usage ()
 			Console.WriteLine ("                            TYPE,PATH_TO_ASSEMBLY: Add user defined type as last step to the pipeline");
 			Console.WriteLine ("                            -NAME:TYPE,PATH_TO_ASSEMBLY: Inserts step type before existing step with name");
 			Console.WriteLine ("                            +NAME:TYPE,PATH_TO_ASSEMBLY: Add step type after existing step");
-			Console.WriteLine ("  --per-assembly-step CFG   Add a custom per-assembly step <config> to the existing pipeline");
+			Console.WriteLine ("  --mark-assembly-step CFG  Add a custom per-assembly step <config> to the existing pipeline");
 			Console.WriteLine ("                            Step can use one of following configurations");
 			Console.WriteLine ("                            TYPE,PATH_TO_ASSEMBLY: Add user defined type as last step to the per-assembly pipeline");
 			Console.WriteLine ("  --custom-data KEY=VALUE   Populates context data set with user specified key-value pair");

src/linker/Linker/Pipeline.cs

@@ -38,12 +38,12 @@ public class Pipeline
 	{
 
 		readonly List<IStep> _steps;
-		readonly List<IPerAssemblyStep> _perAssemblySteps;
+		readonly List<IMarkAssemblyStep> _markAssemblySteps;
 
 		public Pipeline ()
 		{
 			_steps = new List<IStep> ();
-			_perAssemblySteps = new List<IPerAssemblyStep> ();
+			_markAssemblySteps = new List<IMarkAssemblyStep> ();
 		}
 
 		public void PrependStep (IStep step)
@@ -56,9 +56,9 @@ public void AppendStep (IStep step)
 			_steps.Add (step);
 		}
 
-		public void AppendPerAssemblyStep (IPerAssemblyStep step)
+		public void AppendMarkAssemblyStep (IMarkAssemblyStep step)
 		{
-			_perAssemblySteps.Add (step);
+			_markAssemblySteps.Add (step);
 		}
 
 		public void AddStepBefore (Type target, IStep step)
@@ -130,7 +130,7 @@ public void RemoveStep (Type target)
 
 		public void Process (LinkContext context)
 		{
-			foreach (var step in _perAssemblySteps) {
+			foreach (var step in _markAssemblySteps) {
 				step.Initialize (context);
 			}
 
@@ -143,7 +143,7 @@ public void Process (LinkContext context)
 
 		public void ProcessAssembly (AssemblyDefinition assembly)
 		{
-			foreach (var step in _perAssemblySteps) {
+			foreach (var step in _markAssemblySteps) {
 				step.ProcessAssembly (assembly);
 			}
 		}

src/linker/ref/Linker.Steps/IPerAssemblyStep.cs → src/linker/ref/Linker.Steps/IMarkAssemblyStep.cs

@@ -7,7 +7,7 @@
 namespace Mono.Linker.Steps
 {
 
-	public interface IPerAssemblyStep
+	public interface IMarkAssemblyStep
 	{
 		void Initialize (LinkContext context);
 		void ProcessAssembly (AssemblyDefinition assembly);