Skip to content

Commit

Permalink
Reduce allocations in SolutionCompilationState.CreateCompilationTrack…
Browse files Browse the repository at this point in the history 
…erMap (#72596)

Reduce allocations due to resizes in SolutionCompilationState.CreateCompilationTrackerMap. We now use an ImmutableSegmentedDictionary builder to both simplify the code and allocate less. The builder will only create a new dictionary if it's modified, and it does so in a way that doesn't require multiple resizes (which is where the allocation gains from this change are realized).

Additionally, use of the builder allowed a simplification where the modification callback can be changed to a simple action. The call to CreateCompilationTrackerMap now passes in an optimization flag indicating whether the callback needs to occur for empty collections.
  • Loading branch information
@ToddGrun
ToddGrun authored Mar 21, 2024
1 parent 5bc6de4 commit 92e8fa3
Showing 1 changed file with 39 additions and 46 deletions.
85 changes: 39 additions & 46 deletions src/Workspaces/Core/Portable/Workspace/Solution/SolutionCompilationState.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -179,15 +179,13 @@ private SolutionCompilationState ForceForkProject(
if (trackerMap.TryGetValue(arg.projectId, out var tracker))
{
if (!arg.forkTracker)
return trackerMap.Remove(arg.projectId);

trackerMap[arg.projectId] = tracker.Fork(arg.newProjectState, arg.translate);
return true;
trackerMap.Remove(arg.projectId);
else
trackerMap[arg.projectId] = tracker.Fork(arg.newProjectState, arg.translate);
}

return false;
},
(translate, forkTracker, projectId, newProjectState));
(translate, forkTracker, projectId, newProjectState),
skipEmptyCallback: true);

return this.Branch(
newSolutionState,
Expand All @@ -204,10 +202,11 @@ private SolutionCompilationState ForceForkProject(
private ImmutableSegmentedDictionary<ProjectId, ICompilationTracker> CreateCompilationTrackerMap<TArg>(
ProjectId changedProjectId,
ProjectDependencyGraph dependencyGraph,
Func<Dictionary<ProjectId, ICompilationTracker>, TArg, bool> modifyNewTrackerInfo,
TArg arg)
Action<ImmutableSegmentedDictionary<ProjectId, ICompilationTracker>.Builder, TArg> modifyNewTrackerInfo,
TArg arg,
bool skipEmptyCallback)
{
return CreateCompilationTrackerMap(CanReuse, (changedProjectId, dependencyGraph), modifyNewTrackerInfo, arg);
return CreateCompilationTrackerMap(CanReuse, (changedProjectId, dependencyGraph), modifyNewTrackerInfo, arg, skipEmptyCallback);

// Returns true if 'tracker' can be reused for project 'id'
static bool CanReuse(ProjectId id, (ProjectId changedProjectId, ProjectDependencyGraph dependencyGraph) arg)
Expand All @@ -231,10 +230,11 @@ static bool CanReuse(ProjectId id, (ProjectId changedProjectId, ProjectDependenc
private ImmutableSegmentedDictionary<ProjectId, ICompilationTracker> CreateCompilationTrackerMap<TArg>(
ImmutableArray<ProjectId> changedProjectIds,
ProjectDependencyGraph dependencyGraph,
Func<Dictionary<ProjectId, ICompilationTracker>, TArg, bool> modifyNewTrackerInfo,
TArg arg)
Action<ImmutableSegmentedDictionary<ProjectId, ICompilationTracker>.Builder, TArg> modifyNewTrackerInfo,
TArg arg,
bool skipEmptyCallback)
{
return CreateCompilationTrackerMap(CanReuse, (changedProjectIds, dependencyGraph), modifyNewTrackerInfo, arg);
return CreateCompilationTrackerMap(CanReuse, (changedProjectIds, dependencyGraph), modifyNewTrackerInfo, arg, skipEmptyCallback);

// Returns true if 'tracker' can be reused for project 'id'
static bool CanReuse(ProjectId id, (ImmutableArray<ProjectId> changedProjectIds, ProjectDependencyGraph dependencyGraph) arg)
Expand Down Expand Up @@ -262,44 +262,39 @@ static bool CanReuse(ProjectId id, (ImmutableArray<ProjectId> changedProjectIds,
private ImmutableSegmentedDictionary<ProjectId, ICompilationTracker> CreateCompilationTrackerMap<TArgCanReuse, TArgModifyNewTrackerInfo>(
Func<ProjectId, TArgCanReuse, bool> canReuse,
TArgCanReuse argCanReuse,
Func<Dictionary<ProjectId, ICompilationTracker>, TArgModifyNewTrackerInfo, bool> modifyNewTrackerInfo,
TArgModifyNewTrackerInfo argModifyNewTrackerInfo)
Action<ImmutableSegmentedDictionary<ProjectId, ICompilationTracker>.Builder, TArgModifyNewTrackerInfo> modifyNewTrackerInfo,
TArgModifyNewTrackerInfo argModifyNewTrackerInfo,
bool skipEmptyCallback)
{
using var _ = PooledDictionary<ProjectId, ICompilationTracker>.GetInstance(out var newTrackerInfo);

// Keep _projectIdToTrackerMap in a local as it can change during the execution of this method
var projectIdToTrackerMap = _projectIdToTrackerMap;

#if NETCOREAPP
newTrackerInfo.EnsureCapacity(projectIdToTrackerMap.Count);
#endif
// Avoid allocating the builder if the map is empty and the callback doesn't need
// to be called with empty collections.
if (projectIdToTrackerMap.Count == 0 && skipEmptyCallback)
return projectIdToTrackerMap;

var allReused = true;
var projectIdToTrackerMapBuilder = projectIdToTrackerMap.ToBuilder();
foreach (var (id, tracker) in projectIdToTrackerMap)
{
var localTracker = tracker;
if (!canReuse(id, argCanReuse))
{
localTracker = tracker.Fork(tracker.ProjectState, translate: null);
allReused = false;
}
var localTracker = tracker.Fork(tracker.ProjectState, translate: null);

newTrackerInfo.Add(id, localTracker);
projectIdToTrackerMapBuilder[id] = localTracker;
}
}

var isModified = modifyNewTrackerInfo(newTrackerInfo, argModifyNewTrackerInfo);

if (allReused && !isModified)
return projectIdToTrackerMap;
modifyNewTrackerInfo(projectIdToTrackerMapBuilder, argModifyNewTrackerInfo);

return ImmutableSegmentedDictionary.CreateRange(newTrackerInfo);
return projectIdToTrackerMapBuilder.ToImmutable();
}

/// <inheritdoc cref="SolutionState.AddProject(ProjectInfo)"/>
public SolutionCompilationState AddProject(ProjectInfo projectInfo)
{
var newSolutionState = this.SolutionState.AddProject(projectInfo);
var newTrackerMap = CreateCompilationTrackerMap(projectInfo.Id, newSolutionState.GetProjectDependencyGraph(), static (_, _) => false, /* unused */ 0);
var newTrackerMap = CreateCompilationTrackerMap(projectInfo.Id, newSolutionState.GetProjectDependencyGraph(), static (_, _) => { }, /* unused */ 0, skipEmptyCallback: true);

return Branch(
newSolutionState,
Expand All @@ -315,9 +310,10 @@ public SolutionCompilationState RemoveProject(ProjectId projectId)
newSolutionState.GetProjectDependencyGraph(),
static (trackerMap, projectId) =>
{
return trackerMap.Remove(projectId);
trackerMap.Remove(projectId);
},
projectId);
projectId,
skipEmptyCallback: true);

return this.Branch(
newSolutionState,
Expand Down Expand Up @@ -1014,27 +1010,27 @@ public SolutionCompilationState WithoutFrozenSourceGeneratedDocuments()
if (FrozenSourceGeneratedDocumentStates == null)
return this;

var projectIdsToUnfreeze = FrozenSourceGeneratedDocumentStates.Value.States.Values.Select(static state => state.Identity.DocumentId.ProjectId).Distinct();
var projectIdsToUnfreeze = FrozenSourceGeneratedDocumentStates.Value.States.Values
.Select(static state => state.Identity.DocumentId.ProjectId)
.Distinct()
.ToImmutableArray();

// Since we previously froze documents in these projects, we should have a CompilationTracker entry for it, and it should be a
// GeneratedFileReplacingCompilationTracker. To undo the operation, we'll just restore the original CompilationTracker.
var newTrackerMap = CreateCompilationTrackerMap(
projectIdsToUnfreeze.ToImmutableArray(),
projectIdsToUnfreeze,
this.SolutionState.GetProjectDependencyGraph(),
static (trackerMap, projectIdsToUnfreeze) =>
{
var mapModified = false;
foreach (var projectId in projectIdsToUnfreeze)
{
Contract.ThrowIfFalse(trackerMap.TryGetValue(projectId, out var existingTracker));
var replacingItemTracker = (GeneratedFileReplacingCompilationTracker)existingTracker;
trackerMap[projectId] = replacingItemTracker.UnderlyingTracker;
mapModified = true;
}

return mapModified;
},
projectIdsToUnfreeze);
projectIdsToUnfreeze,
skipEmptyCallback: projectIdsToUnfreeze.Length == 0);

// We pass the same solution state, since this change is only a change of the generated documents -- none of the core
// documents or project structure changes in any way.
Expand Down Expand Up @@ -1102,7 +1098,6 @@ public SolutionCompilationState WithFrozenSourceGeneratedDocuments(
this.SolutionState.GetProjectDependencyGraph(),
static (trackerMap, arg) =>
{
var mapModified = false;
foreach (var (projectId, documentStatesForProject) in arg.documentStatesByProjectId)
{
// We want to create a new snapshot with a new compilation tracker that will do this replacement.
Expand All @@ -1114,12 +1109,10 @@ public SolutionCompilationState WithFrozenSourceGeneratedDocuments(
}

trackerMap[projectId] = new GeneratedFileReplacingCompilationTracker(existingTracker, new(documentStatesForProject));
mapModified = true;
}

return mapModified;
},
(documentStatesByProjectId, this.SolutionState));
(documentStatesByProjectId, this.SolutionState),
skipEmptyCallback: false);

// We pass the same solution state, since this change is only a change of the generated documents -- none of the core
// documents or project structure changes in any way.
Expand Down

0 comments on commit 92e8fa3

Please sign in to comment.