Skip to content

Commit 40abde7

Browse files
sbomersbomer
authored
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
1 parent 59eef1f commit 40abde7

File tree

15 files changed

+681
-44
lines changed

15 files changed

+681
-44
lines changed

‎src/linker/Linker.Steps/IMarkHandler.cs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
namespace Mono.Linker.Steps
6+
{
7+
8+
/// <summary>
9+
/// This API supports the product infrastructure and is not intended to be used directly from your code.
10+
/// Extensibility point for custom logic that run during MarkStep, for marked members.
11+
/// </summary>
12+
public interface IMarkHandler
13+
{
14+
/// <summary>
15+
/// Initialize is called at the beginning of MarkStep. This should be
16+
/// used to perform global setup, and register callbacks through the
17+
/// MarkContext.Register* methods) to be called when pieces of IL are marked.
18+
/// </summary>
19+
void Initialize (LinkContext context, MarkContext markContext);
20+
}
21+
}

‎src/linker/Linker.Steps/MarkContext.cs‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
using System;
6+
using Mono.Cecil;
7+
8+
namespace Mono.Linker.Steps
9+
{
10+
/// <summary>
11+
/// Context which can be used to register actions to call during MarkStep
12+
/// when various members are marked.
13+
/// </summary>
14+
public abstract class MarkContext
15+
{
16+
/// <summary>
17+
/// Register a callback that will be invoked once for each marked assembly.
18+
/// </summary>
19+
public abstract void RegisterMarkAssemblyAction (Action<AssemblyDefinition> action);
20+
21+
/// <summary>
22+
/// Register a callback that will be invoked once for each marked type.
23+
/// </summary>
24+
public abstract void RegisterMarkTypeAction (Action<TypeDefinition> action);
25+
26+
/// <summary>
27+
/// Register a callback that will be invoked once for each marked method.
28+
/// </summary>
29+
public abstract void RegisterMarkMethodAction (Action<MethodDefinition> action);
30+
}
31+
}

‎src/linker/Linker.Steps/MarkStep.cs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ public partial class MarkStep : IStep
5656

5757
readonly List<(TypeDefinition Type, MethodBody Body, Instruction Instr)> _pending_isinst_instr;
5858
UnreachableBlocksOptimizer _unreachableBlocksOptimizer;
59+
MarkStepContext _markContext;
5960

6061
#if DEBUG
6162
static readonly DependencyKind[] _entireTypeReasons = new DependencyKind[] {
@@ -190,6 +191,7 @@ public virtual void Process (LinkContext context)
190191
{
191192
_context = context;
192193
_unreachableBlocksOptimizer = new UnreachableBlocksOptimizer (_context);
194+
_markContext = new MarkStepContext ();
193195

194196
Initialize ();
195197
Process ();
@@ -199,6 +201,7 @@ public virtual void Process (LinkContext context)
199201
void Initialize ()
200202
{
201203
InitializeCorelibAttributeXml ();
204+
_context.Pipeline.InitializeMarkHandlers (_context, _markContext);
202205

203206
ProcessMarkedPending ();
204207
}
@@ -1296,6 +1299,9 @@ protected void MarkAssembly (AssemblyDefinition assembly, DependencyInfo reason)
12961299

12971300
EmbeddedXmlInfo.ProcessDescriptors (assembly, _context);
12981301

1302+
foreach (Action<AssemblyDefinition> handleMarkAssembly in _markContext.MarkAssemblyActions)
1303+
handleMarkAssembly (assembly);
1304+
12991305
// Security attributes do not respect the attributes XML
13001306
if (_context.StripSecurity)
13011307
RemoveSecurity.ProcessAssembly (assembly, _context);
@@ -1616,6 +1622,10 @@ protected internal virtual TypeDefinition MarkType (TypeReference reference, Dep
16161622
return null;
16171623

16181624
MarkModule (type.Scope as ModuleDefinition, new DependencyInfo (DependencyKind.ScopeOfType, type));
1625+
1626+
foreach (Action<TypeDefinition> handleMarkType in _markContext.MarkTypeActions)
1627+
handleMarkType (type);
1628+
16191629
MarkType (type.BaseType, new DependencyInfo (DependencyKind.BaseType, type), type);
16201630
if (type.DeclaringType != null)
16211631
MarkType (type.DeclaringType, new DependencyInfo (DependencyKind.DeclaringType, type), type);
@@ -2655,6 +2665,9 @@ protected virtual void ProcessMethod (MethodDefinition method, in DependencyInfo
26552665

26562666
_unreachableBlocksOptimizer.ProcessMethod (method);
26572667

2668+
foreach (Action<MethodDefinition> handleMarkMethod in _markContext.MarkMethodActions)
2669+
handleMarkMethod (method);
2670+
26582671
if (!markedForCall)
26592672
MarkType (method.DeclaringType, new DependencyInfo (DependencyKind.DeclaringType, method), method);
26602673
MarkCustomAttributes (method, new DependencyInfo (DependencyKind.CustomAttribute, method), method);

‎src/linker/Linker.Steps/MarkStepContext.cs‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
using System;
6+
using System.Collections.Generic;
7+
using Mono.Cecil;
8+
9+
namespace Mono.Linker.Steps
10+
{
11+
public class MarkStepContext : MarkContext
12+
{
13+
14+
public List<Action<AssemblyDefinition>> MarkAssemblyActions { get; }
15+
public List<Action<TypeDefinition>> MarkTypeActions { get; }
16+
public List<Action<MethodDefinition>> MarkMethodActions { get; }
17+
18+
public MarkStepContext ()
19+
{
20+
MarkAssemblyActions = new List<Action<AssemblyDefinition>> ();
21+
MarkTypeActions = new List<Action<TypeDefinition>> ();
22+
MarkMethodActions = new List<Action<MethodDefinition>> ();
23+
}
24+
25+
public override void RegisterMarkAssemblyAction (Action<AssemblyDefinition> action)
26+
{
27+
MarkAssemblyActions.Add (action);
28+
}
29+
30+
public override void RegisterMarkTypeAction (Action<TypeDefinition> action)
31+
{
32+
MarkTypeActions.Add (action);
33+
}
34+
35+
public override void RegisterMarkMethodAction (Action<MethodDefinition> action)
36+
{
37+
MarkMethodActions.Add (action);
38+
}
39+
}
40+
}

‎src/linker/Linker.Steps/MarkSubStepsDispatcher.cs‎

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
using System.Collections.Generic;
6+
7+
using Mono.Cecil;
8+
using Mono.Collections.Generic;
9+
10+
namespace Mono.Linker.Steps
11+
{
12+
//
13+
// Dispatcher for SubSteps which only need to run on marked assemblies.
14+
// This simplifies the implementation of linker custom steps, in the same
15+
// way that SubStepsDispatcher does, but it implements IMarkHandler
16+
// and registers a callback that gets invoked during MarkStep when an
17+
// assembly gets marked.
18+
//
19+
public abstract class MarkSubStepsDispatcher : IMarkHandler
20+
{
21+
readonly List<ISubStep> substeps;
22+
23+
List<ISubStep> on_assemblies;
24+
List<ISubStep> on_types;
25+
List<ISubStep> on_fields;
26+
List<ISubStep> on_methods;
27+
List<ISubStep> on_properties;
28+
List<ISubStep> on_events;
29+
30+
protected MarkSubStepsDispatcher (IEnumerable<ISubStep> subSteps)
31+
{
32+
substeps = new List<ISubStep> (subSteps);
33+
}
34+
35+
public virtual void Initialize (LinkContext context, MarkContext markContext)
36+
{
37+
InitializeSubSteps (context);
38+
markContext.RegisterMarkAssemblyAction (assembly => BrowseAssembly (assembly));
39+
}
40+
41+
static bool HasSubSteps (List<ISubStep> substeps) => substeps?.Count > 0;
42+
43+
void BrowseAssembly (AssemblyDefinition assembly)
44+
{
45+
CategorizeSubSteps (assembly);
46+
47+
if (HasSubSteps (on_assemblies))
48+
DispatchAssembly (assembly);
49+
50+
if (!ShouldDispatchTypes ())
51+
return;
52+
53+
BrowseTypes (assembly.MainModule.Types);
54+
}
55+
56+
bool ShouldDispatchTypes ()
57+
{
58+
return HasSubSteps (on_types)
59+
|| HasSubSteps (on_fields)
60+
|| HasSubSteps (on_methods)
61+
|| HasSubSteps (on_properties)
62+
|| HasSubSteps (on_events);
63+
}
64+
65+
void BrowseTypes (Collection<TypeDefinition> types)
66+
{
67+
foreach (TypeDefinition type in types) {
68+
DispatchType (type);
69+
70+
if (type.HasFields && HasSubSteps (on_fields)) {
71+
foreach (FieldDefinition field in type.Fields)
72+
DispatchField (field);
73+
}
74+
75+
if (type.HasMethods && HasSubSteps (on_methods)) {
76+
foreach (MethodDefinition method in type.Methods)
77+
DispatchMethod (method);
78+
}
79+
80+
if (type.HasProperties && HasSubSteps (on_properties)) {
81+
foreach (PropertyDefinition property in type.Properties)
82+
DispatchProperty (property);
83+
}
84+
85+
if (type.HasEvents && HasSubSteps (on_events)) {
86+
foreach (EventDefinition @event in type.Events)
87+
DispatchEvent (@event);
88+
}
89+
90+
if (type.HasNestedTypes)
91+
BrowseTypes (type.NestedTypes);
92+
}
93+
}
94+
95+
void DispatchAssembly (AssemblyDefinition assembly)
96+
{
97+
foreach (var substep in on_assemblies) {
98+
substep.ProcessAssembly (assembly);
99+
}
100+
}
101+
102+
void DispatchType (TypeDefinition type)
103+
{
104+
foreach (var substep in on_types) {
105+
substep.ProcessType (type);
106+
}
107+
}
108+
109+
void DispatchField (FieldDefinition field)
110+
{
111+
foreach (var substep in on_fields) {
112+
substep.ProcessField (field);
113+
}
114+
}
115+
116+
void DispatchMethod (MethodDefinition method)
117+
{
118+
foreach (var substep in on_methods) {
119+
substep.ProcessMethod (method);
120+
}
121+
}
122+
123+
void DispatchProperty (PropertyDefinition property)
124+
{
125+
foreach (var substep in on_properties) {
126+
substep.ProcessProperty (property);
127+
}
128+
}
129+
130+
void DispatchEvent (EventDefinition @event)
131+
{
132+
foreach (var substep in on_events) {
133+
substep.ProcessEvent (@event);
134+
}
135+
}
136+
137+
void InitializeSubSteps (LinkContext context)
138+
{
139+
foreach (var substep in substeps)
140+
substep.Initialize (context);
141+
}
142+
143+
void CategorizeSubSteps (AssemblyDefinition assembly)
144+
{
145+
on_assemblies = new List<ISubStep> ();
146+
on_types = new List<ISubStep> ();
147+
on_fields = new List<ISubStep> ();
148+
on_methods = new List<ISubStep> ();
149+
on_properties = new List<ISubStep> ();
150+
on_events = new List<ISubStep> ();
151+
152+
foreach (var substep in substeps)
153+
CategorizeSubStep (substep, assembly);
154+
}
155+
156+
void CategorizeSubStep (ISubStep substep, AssemblyDefinition assembly)
157+
{
158+
if (!substep.IsActiveFor (assembly))
159+
return;
160+
161+
CategorizeTarget (substep, SubStepTargets.Assembly, on_assemblies);
162+
CategorizeTarget (substep, SubStepTargets.Type, on_types);
163+
CategorizeTarget (substep, SubStepTargets.Field, on_fields);
164+
CategorizeTarget (substep, SubStepTargets.Method, on_methods);
165+
CategorizeTarget (substep, SubStepTargets.Property, on_properties);
166+
CategorizeTarget (substep, SubStepTargets.Event, on_events);
167+
}
168+
169+
static void CategorizeTarget (ISubStep substep, SubStepTargets target, List<ISubStep> list)
170+
{
171+
if (!Targets (substep, target))
172+
return;
173+
174+
list.Add (substep);
175+
}
176+
177+
static bool Targets (ISubStep substep, SubStepTargets target) => (substep.Targets & target) == target;
178+
}
179+
}

0 commit comments

Comments
 (0)