|
// 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 System.Collections.Immutable;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.PooledObjects;
using Microsoft.CodeAnalysis.Shared.Extensions;
namespace Microsoft.CodeAnalysis.Shared.Utilities;
internal sealed class AnnotatedSymbolMapping(
ImmutableDictionary<ISymbol, SyntaxAnnotation> symbolToDeclarationAnnotationMap,
Solution annotatedSolution,
ImmutableDictionary<DocumentId, ImmutableArray<ISymbol>> documentIdsToSymbolMap,
SyntaxAnnotation typeNodeAnnotation)
{
/// <summary>
/// Used to map a symbol to the annotation that was added at the beginning of it's definition. Use the
/// annotation the symbol declaration again across edits.
/// </summary>
public ImmutableDictionary<ISymbol, SyntaxAnnotation> SymbolToDeclarationAnnotationMap { get; } = symbolToDeclarationAnnotationMap;
/// <summary>
/// The original solution that modifications made to annotate the symbol declarations
/// </summary>
public Solution AnnotatedSolution { get; } = annotatedSolution;
/// <summary>
/// A map of the document ids that were used and what symbols are in them.
/// </summary>
public ImmutableDictionary<DocumentId, ImmutableArray<ISymbol>> DocumentIdsToSymbolMap { get; } = documentIdsToSymbolMap;
/// <summary>
/// The annotation added to the type declaration that was passed in to create the mapping
/// </summary>
public SyntaxAnnotation TypeNodeAnnotation { get; } = typeNodeAnnotation;
/// <summary>
/// Creates a <see cref="AnnotatedSymbolMapping"/> where the first token of each symbol is annotated
/// and added to a map to keep track. This allows modification of the trees and later lookup of symbols
/// based on the original annotations added. Assumes each symbol only has one location.
/// </summary>
public static async Task<AnnotatedSymbolMapping> CreateAsync(
IEnumerable<ISymbol> symbols,
Solution solution,
SyntaxNode typeNode,
CancellationToken cancellationToken)
{
using var _ = PooledDictionary<ISymbol, SyntaxAnnotation>.GetInstance(out var symbolToDeclarationAnnotationMap);
using var _1 = PooledDictionary<SyntaxTree, SyntaxNode>.GetInstance(out var currentRoots);
using var _2 = PooledDictionary<DocumentId, List<ISymbol>>.GetInstance(out var documentIdToSymbolsMap);
var typeNodeRoot = await typeNode.SyntaxTree.GetRootAsync(cancellationToken).ConfigureAwait(false);
var typeNodeAnnotation = new SyntaxAnnotation();
// CurrentRoots will keep track of the root node with annotations added
currentRoots[typeNode.SyntaxTree] = typeNodeRoot.ReplaceNode(typeNode, typeNode.WithAdditionalAnnotations(typeNodeAnnotation));
// DocumentIds will track all of the documents where annotations were added since
// it's not guaranteed that all of the symbols are in the same document
documentIdToSymbolsMap.Add(solution.GetRequiredDocument(typeNode.SyntaxTree).Id, []);
foreach (var symbol in symbols)
{
var location = symbol.Locations.Single();
var tree = location.SourceTree!;
var id = solution.GetRequiredDocument(tree).Id;
// If there's not currently an entry for this tree then make sure to add it
if (!currentRoots.TryGetValue(tree, out var root))
{
root = await tree.GetRootAsync(cancellationToken).ConfigureAwait(false);
documentIdToSymbolsMap.Add(id, []);
}
var token = root.FindToken(location.SourceSpan.Start);
var annotation = new SyntaxAnnotation();
// Add the instance of the annotation used to annotate this symbol so it
// can be retrieved later
symbolToDeclarationAnnotationMap.Add(symbol, annotation);
// Add the symbol to list of symbols contained in the document
var symbolsInDocument = documentIdToSymbolsMap[id];
symbolsInDocument.Add(symbol);
// Store the updated root node with the annotation added
currentRoots[tree] = root.ReplaceToken(token, token.WithAdditionalAnnotations(annotation));
}
// Make sure each document is updated with the annotated root
var annotatedSolution = solution;
foreach (var root in currentRoots)
{
var document = annotatedSolution.GetRequiredDocument(root.Key);
annotatedSolution = document.WithSyntaxRoot(root.Value).Project.Solution;
}
var immutableDocumentIdToSymbolsMap = documentIdToSymbolsMap.ToImmutableDictionary(
keySelector: (kvp) => kvp.Key,
elementSelector: (kvp) => kvp.Value.ToImmutableArray());
return new AnnotatedSymbolMapping(symbolToDeclarationAnnotationMap.ToImmutableDictionary(), annotatedSolution, immutableDocumentIdToSymbolsMap, typeNodeAnnotation);
}
}
|