Interceptor for identity resolution when attaching graphs (#28246)

src/EFCore.InMemory/Storage/Internal/InMemoryTable.cs

@@ -332,7 +332,7 @@ public virtual void BumpValueGenerators(object?[] row)
     }
 
     private TKey CreateKey(IUpdateEntry entry)
-        => _keyValueFactory.CreateFromCurrentValues(entry);
+        => _keyValueFactory.CreateFromCurrentValues(entry)!;
 
     private static object? SnapshotValue(IProperty property, ValueComparer? comparer, IUpdateEntry entry)
     {

src/EFCore/ChangeTracking/IPrincipalKeyValueFactory.cs

@@ -45,7 +45,7 @@ public interface IPrincipalKeyValueFactory<TKey>
     /// </summary>
     /// <param name="entry">The entry tracking an entity instance.</param>
     /// <returns>The key value.</returns>
-    TKey CreateFromCurrentValues(IUpdateEntry entry);
+    TKey? CreateFromCurrentValues(IUpdateEntry entry);
 
     /// <summary>
     ///     Finds the first null key value in the given entry and returns the associated <see cref="IProperty" />.
@@ -59,7 +59,7 @@ public interface IPrincipalKeyValueFactory<TKey>
     /// </summary>
     /// <param name="entry">The entry tracking an entity instance.</param>
     /// <returns>The key value.</returns>
-    TKey CreateFromOriginalValues(IUpdateEntry entry);
+    TKey? CreateFromOriginalValues(IUpdateEntry entry);
 
     /// <summary>
     ///     Creates a key object from the relationship snapshot key values in the given entry.

src/EFCore/ChangeTracking/Internal/DependentsMap.cs

@@ -126,7 +126,7 @@ private bool TryCreateFromCurrentValues(IUpdateEntry entry, [NotNullWhen(true)]
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public virtual IEnumerable<IUpdateEntry> GetDependents(IUpdateEntry principalEntry)
-        => _map.TryGetValue(_principalKeyValueFactory.CreateFromCurrentValues(principalEntry), out var dependents)
+        => _map.TryGetValue(_principalKeyValueFactory.CreateFromCurrentValues(principalEntry)!, out var dependents)
             ? dependents
             : Enumerable.Empty<IUpdateEntry>();
 

src/EFCore/ChangeTracking/Internal/EntityGraphAttacher.cs

@@ -12,6 +12,7 @@ namespace Microsoft.EntityFrameworkCore.ChangeTracking.Internal;
 public class EntityGraphAttacher : IEntityGraphAttacher
 {
     private readonly IEntityEntryGraphIterator _graphIterator;
+    private HashSet<object>? _visited;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -49,10 +50,12 @@ public virtual void AttachGraph(
                     null),
                 PaintAction);
 
+            _visited = null;
             rootEntry.StateManager.CompleteAttachGraph();
         }
         catch
         {
+            _visited = null;
             rootEntry.StateManager.AbortAttachGraph();
             throw;
         }
@@ -84,22 +87,25 @@ await _graphIterator.TraverseGraphAsync(
                 PaintActionAsync,
                 cancellationToken).ConfigureAwait(false);
 
+            _visited = null;
             rootEntry.StateManager.CompleteAttachGraph();
         }
         catch
         {
+            _visited = null;
             rootEntry.StateManager.AbortAttachGraph();
             throw;
         }
     }
 
-    private static bool PaintAction(
+    private bool PaintAction(
         EntityEntryGraphNode<(EntityState TargetState, EntityState StoreGenTargetState, bool Force)> node)
     {
         SetReferenceLoaded(node);
 
         var internalEntityEntry = node.GetInfrastructure();
-        if (internalEntityEntry.EntityState != EntityState.Detached)
+        if (internalEntityEntry.EntityState != EntityState.Detached
+            || (_visited != null && _visited.Contains(internalEntityEntry.Entity)))
         {
             return false;
         }
@@ -108,24 +114,34 @@ private static bool PaintAction(
 
         var (isGenerated, isSet) = internalEntityEntry.IsKeySet;
 
-        internalEntityEntry.SetEntityState(
-            isSet
-                ? (isGenerated ? storeGenTargetState : targetState)
-                : EntityState.Added, // Key can only be not-set if it is store-generated
-            acceptChanges: true,
-            forceStateWhenUnknownKey: force ? targetState : null);
+        if (internalEntityEntry.StateManager.ResolveToExistingEntry(
+                internalEntityEntry,
+                node.InboundNavigation, node.SourceEntry?.GetInfrastructure()))
+        {
+            (_visited ??= new HashSet<object>(ReferenceEqualityComparer.Instance)).Add(internalEntityEntry.Entity);
+        }
+        else
+        {
+            internalEntityEntry.SetEntityState(
+                isSet
+                    ? (isGenerated ? storeGenTargetState : targetState)
+                    : EntityState.Added, // Key can only be not-set if it is store-generated
+                acceptChanges: true,
+                forceStateWhenUnknownKey: force ? targetState : null);
+        }
 
         return true;
     }
 
-    private static async Task<bool> PaintActionAsync(
+    private async Task<bool> PaintActionAsync(
         EntityEntryGraphNode<(EntityState TargetState, EntityState StoreGenTargetState, bool Force)> node,
         CancellationToken cancellationToken)
     {
         SetReferenceLoaded(node);
 
         var internalEntityEntry = node.GetInfrastructure();
-        if (internalEntityEntry.EntityState != EntityState.Detached)
+        if (internalEntityEntry.EntityState != EntityState.Detached
+            || (_visited != null && _visited.Contains(internalEntityEntry.Entity)))
         {
             return false;
         }
@@ -134,14 +150,23 @@ private static async Task<bool> PaintActionAsync(
 
         var (isGenerated, isSet) = internalEntityEntry.IsKeySet;
 
-        await internalEntityEntry.SetEntityStateAsync(
-                isSet
-                    ? (isGenerated ? storeGenTargetState : targetState)
-                    : EntityState.Added, // Key can only be not-set if it is store-generated
-                acceptChanges: true,
-                forceStateWhenUnknownKey: force ? targetState : null,
-                cancellationToken: cancellationToken)
-            .ConfigureAwait(false);
+        if (internalEntityEntry.StateManager.ResolveToExistingEntry(
+                internalEntityEntry,
+                node.InboundNavigation, node.SourceEntry?.GetInfrastructure()))
+        {
+            (_visited ??= new HashSet<object>()).Add(internalEntityEntry.Entity);
+        }
+        else
+        {
+            await internalEntityEntry.SetEntityStateAsync(
+                    isSet
+                        ? (isGenerated ? storeGenTargetState : targetState)
+                        : EntityState.Added, // Key can only be not-set if it is store-generated
+                    acceptChanges: true,
+                    forceStateWhenUnknownKey: force ? targetState : null,
+                    cancellationToken: cancellationToken)
+                .ConfigureAwait(false);
+        }
 
         return true;
     }

src/EFCore/ChangeTracking/Internal/IIdentityMap.cs

@@ -35,6 +35,14 @@ public interface IIdentityMap
     /// </summary>
     bool Contains(IForeignKey foreignKey, in ValueBuffer valueBuffer);
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    InternalEntityEntry? TryGetEntry(InternalEntityEntry entry);
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in

src/EFCore/ChangeTracking/Internal/IInternalEntityEntryNotifier.cs

@@ -33,6 +33,14 @@ public interface IInternalEntityEntryNotifier
     /// </summary>
     void StateChanged(InternalEntityEntry entry, EntityState oldState, bool fromQuery);
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    void FixupResolved(InternalEntityEntry entry, InternalEntityEntry duplicateEntry);
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in

src/EFCore/ChangeTracking/Internal/INavigationFixer.cs

@@ -105,4 +105,14 @@ void StateChanged(
         InternalEntityEntry entry,
         EntityState oldState,
         bool fromQuery);
+
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    void FixupResolved(
+        InternalEntityEntry entry,
+        InternalEntityEntry duplicateEntry);
 }

src/EFCore/ChangeTracking/Internal/IStateManager.cs

@@ -252,6 +252,29 @@ void RecordReferencedUntrackedEntity(
         INavigationBase navigation,
         InternalEntityEntry referencedFromEntry);
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    void UpdateReferencedUntrackedEntity(
+        object referencedEntity,
+        object newReferencedEntity,
+        INavigationBase navigation,
+        InternalEntityEntry referencedFromEntry);
+
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    bool ResolveToExistingEntry(
+        InternalEntityEntry newEntry,
+        INavigationBase? navigation,
+        InternalEntityEntry? referencedFromEntry);
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in

src/EFCore/ChangeTracking/Internal/IdentityMap.cs

@@ -83,6 +83,20 @@ public virtual bool Contains(IForeignKey foreignKey, in ValueBuffer valueBuffer)
         => foreignKey.GetDependentKeyValueFactory<TKey>()!.TryCreateFromBuffer(valueBuffer, out var key)
             && _identityMap.ContainsKey(key);
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public virtual InternalEntityEntry? TryGetEntry(InternalEntityEntry entry)
+    {
+        var key = PrincipalKeyValueFactory.CreateFromCurrentValues(entry);
+        return key != null && _identityMap.TryGetValue(key, out var existingEntry)
+            ? existingEntry
+            : null;
+    }
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -194,7 +208,7 @@ public virtual bool Contains(IForeignKey foreignKey, in ValueBuffer valueBuffer)
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public virtual void AddOrUpdate(InternalEntityEntry entry)
-        => Add(PrincipalKeyValueFactory.CreateFromCurrentValues(entry), entry, updateDuplicate: true);
+        => Add(PrincipalKeyValueFactory.CreateFromCurrentValues(entry)!, entry, updateDuplicate: true);
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -203,7 +217,7 @@ public virtual void AddOrUpdate(InternalEntityEntry entry)
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public virtual void Add(InternalEntityEntry entry)
-        => Add(PrincipalKeyValueFactory.CreateFromCurrentValues(entry), entry);
+        => Add(PrincipalKeyValueFactory.CreateFromCurrentValues(entry)!, entry);
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -371,7 +385,7 @@ public virtual void Clear()
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public virtual void Remove(InternalEntityEntry entry)
-        => Remove(PrincipalKeyValueFactory.CreateFromCurrentValues(entry), entry);
+        => Remove(PrincipalKeyValueFactory.CreateFromCurrentValues(entry)!, entry);
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to

src/EFCore/ChangeTracking/Internal/InternalEntityEntryNotifier.cs

@@ -55,6 +55,17 @@ public virtual void StateChanged(InternalEntityEntry entry, EntityState oldState
         _localViewListener.StateChanged(entry, oldState, fromQuery);
     }
 
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public virtual void FixupResolved(
+        InternalEntityEntry entry,
+        InternalEntityEntry duplicateEntry)
+        => _navigationFixer.FixupResolved(entry, duplicateEntry);
+
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
     ///     the same compatibility standards as public APIs. It may be changed or removed without notice in