Skip to content

Fix type map assembly target traversal in ILC #124247

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jkoritzinsky merged 4 commits into from
Feb 14, 2026
Merged

Fix type map assembly target traversal in ILC #124247

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a16c526
Type map assembly target references should only be followed for the p…
jkoritzinsky Feb 9, 2026
d4cfc39
Add negative test
jkoritzinsky Feb 12, 2026
90e75ec
Fix attribute to check
jkoritzinsky Feb 12, 2026
653585b
Copilot feedback
jkoritzinsky Feb 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
203 changes: 173 additions & 30 deletions src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/TypeMapMetadata.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ protected override int CompareToImpl(MethodDesc other, TypeSystemComparer compar

private readonly Dictionary<TypeDesc, TypeDesc> _associatedTypeMap = [];
private readonly Dictionary<string, (TypeDesc type, TypeDesc trimmingTarget)> _externalTypeMap = [];
private readonly List<ModuleDesc> _targetModules = [];
private ThrowingMethodStub _externalTypeMapExceptionStub;
private ThrowingMethodStub _associatedTypeMapExceptionStub;

Expand Down Expand Up @@ -104,6 +105,57 @@ public void SetAssociatedTypeMapException(ModuleDesc stubModule, TypeSystemExcep
_associatedTypeMapExceptionStub ??= new ThrowingMethodStub(stubModule.GetGlobalModuleType(), TypeMapGroup, externalTypeMap: false, exception);
}

public void MergePendingMap(Map pendingMap)
{
// Don't waste time adding entries from the pending map if we already have an exception stub,
// as the exception stub means the map is invalid and the entries won't be used anyway.
if (_associatedTypeMapExceptionStub is null)
{
if (pendingMap._associatedTypeMapExceptionStub is not null)
{
_associatedTypeMapExceptionStub = pendingMap._associatedTypeMapExceptionStub;
}
else
{
foreach (KeyValuePair<TypeDesc, TypeDesc> kvp in pendingMap._associatedTypeMap)
{
AddAssociatedTypeMapEntry(kvp.Key, kvp.Value);
}
}
}
else if (_associatedTypeMapExceptionStub.Exception is not TypeSystemException.FileNotFoundException &&
pendingMap._associatedTypeMapExceptionStub?.Exception is TypeSystemException.FileNotFoundException)
{
// If the pending map has a FileNotFoundException, it takes precedence over our existing exception stub, so use it instead.
_associatedTypeMapExceptionStub = pendingMap._associatedTypeMapExceptionStub;
}

// Don't waste time adding entries from the pending map if we already have an exception stub,
// as the exception stub means the map is invalid and the entries won't be used anyway.
if (_externalTypeMapExceptionStub is null)
{
if (pendingMap._externalTypeMapExceptionStub is not null)
{
_externalTypeMapExceptionStub = pendingMap._externalTypeMapExceptionStub;
}
else
{
foreach (KeyValuePair<string, (TypeDesc type, TypeDesc trimmingTarget)> kvp in pendingMap._externalTypeMap)
{
AddExternalTypeMapEntry(kvp.Key, kvp.Value.type, kvp.Value.trimmingTarget);
}
}
}
else if (_externalTypeMapExceptionStub.Exception is not TypeSystemException.FileNotFoundException &&
pendingMap._externalTypeMapExceptionStub?.Exception is TypeSystemException.FileNotFoundException)
{
// If the pending map has a FileNotFoundException, it takes precedence over our existing exception stub, so use it instead.
_externalTypeMapExceptionStub = pendingMap._externalTypeMapExceptionStub;
}

_targetModules.AddRange(pendingMap._targetModules);
}

public IExternalTypeMapNode GetExternalTypeMapNode()
{
if (_externalTypeMapExceptionStub is not null)
Expand All @@ -121,6 +173,16 @@ public IProxyTypeMapNode GetProxyTypeMapNode()
}
return new ProxyTypeMapNode(TypeMapGroup, _associatedTypeMap);
}

public void AddTargetModule(ModuleDesc targetModule)
{
_targetModules.Add(targetModule);
}

/// <summary>
/// The modules targeted with TypeMapAssemblyTarget attributes for this type map group. This is only populated when TypeMapMetadata is created with TypeMapAssemblyTargetsMode.Record. When TypeMapMetadata is created with TypeMapAssemblyTargetsMode.Traverse, this will be empty as the target assemblies will be traversed to include their type maps instead of just being recorded as targets.
/// </summary>
public IEnumerable<ModuleDesc> TargetModules => _targetModules;
}

public static readonly TypeMapMetadata Empty = new TypeMapMetadata(new Dictionary<TypeDesc, Map>(), "No type maps");
Expand All @@ -141,22 +203,55 @@ private TypeMapMetadata(IReadOnlyDictionary<TypeDesc, Map> states, string diagno

public string DiagnosticName { get; }

public static TypeMapMetadata CreateFromAssembly(EcmaAssembly assembly, CompilerTypeSystemContext typeSystemContext)
public static TypeMapMetadata CreateFromAssembly(EcmaAssembly assembly, ModuleDesc throwHelperEmitModule, TypeMapAssemblyTargetsMode assemblyTargetsMode)
{
Dictionary<TypeDesc, Map> typeMapStates = [];
HashSet<EcmaAssembly> scannedAssemblies = [];
// The pendingMaps collection represents assemblies that have been scanned, but the provided assembly
// has not been added as a target yet for the specified type map group.
// This can occur when an assembly is the target of a TypeMapAssemblyTarget attribute for a different group.
Dictionary<(EcmaAssembly assembly, TypeDesc typeMapGroup), (TypeDesc scannedDuringGroup, Map map)> pendingMaps = [];
HashSet<(EcmaAssembly assembly, TypeDesc typeMapGroup)> scannedAssemblies = [];

Queue<EcmaAssembly> assembliesToScan = new Queue<EcmaAssembly>();
assembliesToScan.Enqueue(assembly);
Queue<(EcmaAssembly assembly, TypeDesc typeMapGroup)> assembliesToScan = [];
assembliesToScan.Enqueue((assembly, null));

while (assembliesToScan.Count > 0)
{
EcmaAssembly currentAssembly = assembliesToScan.Dequeue();
if (scannedAssemblies.Contains(currentAssembly))
continue;
(EcmaAssembly currentAssembly, TypeDesc currentTypeMapGroup) = assembliesToScan.Dequeue();

scannedAssemblies.Add(currentAssembly);
// A null currentTypeMapGroup means we're looking at all type maps in the assembly, not just traversing one.
// Otherwise, we're searching for a specific one.
// Don't rescan specific assembly + type map group combos as the results will be identical.
if (currentTypeMapGroup is not null)
{
if (scannedAssemblies.Contains((currentAssembly, currentTypeMapGroup)))
{
if (pendingMaps.TryGetValue((currentAssembly, currentTypeMapGroup), out var pendingMap))
{
// We have already scanned this assembly, but we haven't added the results to the type map state for the specified type map group.
// Merge in the results.
if (!typeMapStates.TryGetValue(currentTypeMapGroup, out Map typeMapState))
{
typeMapStates[currentTypeMapGroup] = typeMapState = new Map(currentTypeMapGroup);
}
typeMapState.MergePendingMap(pendingMap.map);
foreach (ModuleDesc targetModule in pendingMap.map.TargetModules)
{
Debug.Assert(assemblyTargetsMode == TypeMapAssemblyTargetsMode.Traverse, "We should only have pending maps with target modules when we're traversing for type map groups, as opposed to just recording targets.");
// If the pending map has target modules,
// then we need to ensure that they also get included in the metadata.
assembliesToScan.Enqueue(((EcmaAssembly)targetModule, currentTypeMapGroup));
}
pendingMaps.Remove((currentAssembly, currentTypeMapGroup));
}
// We've already scanned this assembly for this type map group.
// There's no more work to do for this case.
continue;
}
}

// Either we haven't seen this assembly + group combo before
// or we are scanning all type map groups in the assembly.
foreach (CustomAttributeHandle attrHandle in currentAssembly.MetadataReader.GetCustomAttributes(EntityHandle.AssemblyDefinition))
{
CustomAttribute attr = currentAssembly.MetadataReader.GetCustomAttribute(attrHandle);
Expand Down Expand Up @@ -185,20 +280,55 @@ public static TypeMapMetadata CreateFromAssembly(EcmaAssembly assembly, Compiler

TypeDesc typeMapGroup = type.Instantiation[0];

Map typeMapState;

if (currentTypeMapGroup is null || typeMapGroup == currentTypeMapGroup)
{
// This attribute is definitely included in the type map.
if (!typeMapStates.TryGetValue(typeMapGroup, out typeMapState))
{
typeMapStates[typeMapGroup] = typeMapState = new Map(typeMapGroup);
}
}
else
{
// This attribute may be included in the type map, but this type map group isn't the one that triggered us to scan this assembly,
// so we need to save it as a pending map in case its type map group is triggered later when we scan another assembly.
if (!pendingMaps.TryGetValue((currentAssembly, typeMapGroup), out var pendingMap))
{
typeMapState = new Map(typeMapGroup);
pendingMaps[(currentAssembly, typeMapGroup)] = (currentTypeMapGroup, typeMapState);
}
else
{
if (pendingMap.scannedDuringGroup != currentTypeMapGroup)
{
// We already scanned this assembly for this type map group while
// processing a different type map group.
// Skip processing it again to avoid hitting duplicate entries.
continue;
}
typeMapState = pendingMap.map;
}
}

// Mark this assembly + type map group as scanned to avoid redundant work in the future.
scannedAssemblies.Add((currentAssembly, typeMapGroup));

try
{
switch (attrKind)
{
case TypeMapAttributeKind.TypeMapAssemblyTarget:
ProcessTypeMapAssemblyTargetAttribute(attrValue);
ProcessTypeMapAssemblyTargetAttribute(attrValue, typeMapState);
break;

case TypeMapAttributeKind.TypeMap:
ProcessTypeMapAttribute(attrValue, typeMapGroup);
ProcessTypeMapAttribute(attrValue, typeMapState);
break;

case TypeMapAttributeKind.TypeMapAssociation:
ProcessTypeMapAssociationAttribute(attrValue, typeMapGroup);
ProcessTypeMapAssociationAttribute(attrValue, typeMapState);
break;

default:
Expand All @@ -216,17 +346,24 @@ public static TypeMapMetadata CreateFromAssembly(EcmaAssembly assembly, Compiler

if (attrKind is TypeMapAttributeKind.TypeMapAssemblyTarget or TypeMapAttributeKind.TypeMap)
{
value.SetExternalTypeMapException(typeSystemContext.GeneratedAssembly, ex);
value.SetExternalTypeMapException(throwHelperEmitModule, ex);
}

if (attrKind is TypeMapAttributeKind.TypeMapAssemblyTarget or TypeMapAttributeKind.TypeMapAssociation)
{
value.SetAssociatedTypeMapException(typeSystemContext.GeneratedAssembly, ex);
value.SetAssociatedTypeMapException(throwHelperEmitModule, ex);
}
}
}

void ProcessTypeMapAssemblyTargetAttribute(CustomAttributeValue<TypeDesc> attrValue)
if (currentTypeMapGroup is not null)
{
// Mark this assembly + type map group as scanned to avoid redundant work in the future.
// We can hit this case when the type map group has no entries in the current assembly.
scannedAssemblies.Add((currentAssembly, currentTypeMapGroup));
}

void ProcessTypeMapAssemblyTargetAttribute(CustomAttributeValue<TypeDesc> attrValue, Map typeMapState)
{
if (attrValue.FixedArguments is not [{ Value: string assemblyName }])
{
Expand All @@ -235,30 +372,29 @@ void ProcessTypeMapAssemblyTargetAttribute(CustomAttributeValue<TypeDesc> attrVa
}

EcmaAssembly targetAssembly = (EcmaAssembly)assembly.Context.ResolveAssembly(AssemblyNameInfo.Parse(assemblyName), throwIfNotFound: true);
typeMapState.AddTargetModule(targetAssembly);

assembliesToScan.Enqueue(targetAssembly);
if (assemblyTargetsMode == TypeMapAssemblyTargetsMode.Traverse
&& (currentTypeMapGroup is null || currentTypeMapGroup == typeMapState.TypeMapGroup))
{
// If we're traversing for the current type map group, enqueue the target to be scanned.
// Otherwise, we'll pull the targets to be scanned when the pending group is scanned.
assembliesToScan.Enqueue((targetAssembly, typeMapState.TypeMapGroup));
}
}

void ProcessTypeMapAttribute(CustomAttributeValue<TypeDesc> attrValue, TypeDesc typeMapGroup)
void ProcessTypeMapAttribute(CustomAttributeValue<TypeDesc> attrValue, Map typeMapState)
{
switch (attrValue.FixedArguments)
{
case [{ Value: string typeName }, { Value: TypeDesc targetType }]:
{
if (!typeMapStates.TryGetValue(typeMapGroup, out Map typeMapState))
{
typeMapStates[typeMapGroup] = typeMapState = new Map(typeMapGroup);
}
typeMapState.AddExternalTypeMapEntry(typeName, targetType, null);
break;
}

case [{ Value: string typeName }, { Value: TypeDesc targetType }, { Value: TypeDesc trimTargetType }]:
{
if (!typeMapStates.TryGetValue(typeMapGroup, out Map typeMapState))
{
typeMapStates[typeMapGroup] = typeMapState = new Map(typeMapGroup);
}
typeMapState.AddExternalTypeMapEntry(typeName, targetType, trimTargetType);
break;
}
Expand All @@ -269,7 +405,7 @@ void ProcessTypeMapAttribute(CustomAttributeValue<TypeDesc> attrValue, TypeDesc
}
}

void ProcessTypeMapAssociationAttribute(CustomAttributeValue<TypeDesc> attrValue, TypeDesc typeMapGroup)
void ProcessTypeMapAssociationAttribute(CustomAttributeValue<TypeDesc> attrValue, Map typeMapState)
{
// If attribute is TypeMapAssociationAttribute, we need to extract the generic argument (type map group)
// and process it.
Expand All @@ -279,16 +415,23 @@ void ProcessTypeMapAssociationAttribute(CustomAttributeValue<TypeDesc> attrValue
return;
}

if (!typeMapStates.TryGetValue(typeMapGroup, out Map typeMapState))
{
typeMapStates[typeMapGroup] = typeMapState = new Map(typeMapGroup);
}

typeMapState.AddAssociatedTypeMapEntry(type, associatedType);
}
}

return new TypeMapMetadata(typeMapStates, $"Type maps rooted at {assembly}");
}
}

public enum TypeMapAssemblyTargetsMode
{
/// <summary>
/// Traverse the TypeMapAssemblyTarget attributes and include type maps from the target assemblies.
/// </summary>
Traverse,
/// <summary>
/// Record the TypeMapAssemblyTarget attributes but do not traverse them to include type maps from the target assemblies.
/// </summary>
Record
}
}
4 changes: 2 additions & 2 deletions src/coreclr/tools/aot/ILCompiler.Trimming.Tests/TestCasesRunner/TrimmingDriver.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,12 +150,12 @@ public ILScanResults Trim(ILCompilerOptions options, TrimmingCustomizations? cus
if (options.TypeMapEntryAssembly is not null
&& typeSystemContext.ResolveAssembly(AssemblyNameInfo.Parse(options.TypeMapEntryAssembly), throwIfNotFound: true) is EcmaAssembly typeMapEntryAssembly)
{
typeMapManager = new UsageBasedTypeMapManager(TypeMapMetadata.CreateFromAssembly(typeMapEntryAssembly, typeSystemContext));
typeMapManager = new UsageBasedTypeMapManager(TypeMapMetadata.CreateFromAssembly(typeMapEntryAssembly, typeSystemContext.GeneratedAssembly, TypeMapAssemblyTargetsMode.Traverse));
}
else if (entrypointModule is { Assembly: EcmaAssembly entryAssembly })
{
// Pass null for typeMappingEntryAssembly to use default entry assembly behavior in tests
typeMapManager = new UsageBasedTypeMapManager(TypeMapMetadata.CreateFromAssembly(entryAssembly, typeSystemContext));
typeMapManager = new UsageBasedTypeMapManager(TypeMapMetadata.CreateFromAssembly(entryAssembly, typeSystemContext.GeneratedAssembly, TypeMapAssemblyTargetsMode.Traverse));
}

CompilationBuilder builder = new RyuJitCompilationBuilder(typeSystemContext, compilationGroup)
Expand Down
Loading
Loading