Interceptor for identity resolution when attaching graphs

Part of #626
Fixes #20124

This is essentially a call back that happens when an instance withe the same key as an existing tracked instance is attached. The existing instance is always retained, but the interceptor gets the chance to change property values of the tracked instance. This will typically be used to copy over values from the duplicate instance (i.e. painting from a detached entity) or just skip the duplicate. Interceptor implementations are provided to do these two things.

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,32 @@ 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.ResolveToExistingEntry(node.InboundNavigation, node.SourceEntry?.GetInfrastructure()))
+        {
+            (_visited ??= new HashSet<object>()).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 +148,21 @@ 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.ResolveToExistingEntry(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/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/InternalEntityEntry.cs

@@ -428,6 +428,17 @@ private void SetServiceProperties(EntityState oldState, EntityState newState)
         }
     }
 
+    /// <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 bool ResolveToExistingEntry(
+        INavigationBase? navigation,
+        InternalEntityEntry? referencedFromEntry)
+        => StateManager.ResolveToExistingEntry(this, navigation, 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/NavigationFixer.cs

@@ -955,14 +955,24 @@ private void DelayedFixup(
                             FixupToDependent(entry, referencedEntry, navigation.ForeignKey, setModified, fromQuery);
                         }
                     }
-                    else if (referencedEntry.Entity == navigationValue)
+                    else
                     {
-                        FixupToDependent(entry, referencedEntry, navigation.ForeignKey, setModified, fromQuery);
+                        FixupToDependent(
+                            entry, 
+                            referencedEntry, 
+                            navigation.ForeignKey, 
+                            referencedEntry.Entity == navigationValue && setModified, 
+                            fromQuery);
                     }
                 }
-                else if (referencedEntry.Entity == navigationValue)
+                else
                 {
-                    FixupToPrincipal(entry, referencedEntry, navigation.ForeignKey, setModified, fromQuery);
+                    FixupToPrincipal(
+                        entry, 
+                        referencedEntry, 
+                        navigation.ForeignKey, 
+                        referencedEntry.Entity == navigationValue && setModified, 
+                        fromQuery);
 
                     FixupSkipNavigations(entry, navigation.ForeignKey, fromQuery);
                 }