Fixed #248

dotnet · Jul 30, 2024 · 15dd715 · 15dd715
1 parent 6fe0f20
commit 15dd715
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 10 deletions.
diff --git a/src/DotNext.Tests/BasicExtensionsTests.cs b/src/DotNext.Tests/BasicExtensionsTests.cs
@@ -69,6 +69,19 @@ public static void CopyDataStorage2()
         Equal(42L, obj1.GetUserData().Get(slot));
     }
 
+    [Fact]
+    public static void RegressionIssue248()
+    {
+        var slot = new UserDataSlot<long>();
+        var str1 = new string('a', 3);
+        var str2 = new string('b', 3);
+        NotSame(str1, str2);
+        str2.GetUserData().Set(slot, 42L);
+        str1.GetUserData().CopyTo(str2);
+
+        False(str2.GetUserData().TryGet(slot, out _));
+    }
+
     [Fact]
     public static void UserDataStorageGetOrSet()
     {

diff --git a/src/DotNext/UserDataStorage.BackingStorage.cs b/src/DotNext/UserDataStorage.BackingStorage.cs
@@ -282,7 +282,7 @@ static UserDataStorage()
         Partitions = new ConditionalWeakTable<object, BackingStorage>?[size];
     }
 
-    private static ref ConditionalWeakTable<object, BackingStorage>? GetStorage(object source)
+    private static ref ConditionalWeakTable<object, BackingStorage>? GetPartition(object source)
     {
         Debug.Assert(BitOperations.IsPow2(Partitions.Length));
 
@@ -292,9 +292,9 @@ static UserDataStorage()
         return ref Unsafe.Add(ref MemoryMarshal.GetArrayDataReference(Partitions), bucketIndex);
     }
 
-    private static ConditionalWeakTable<object, BackingStorage> GetOrCreateStorage(object source)
+    private static ConditionalWeakTable<object, BackingStorage> GetOrCreatePartition(object source)
     {
-        ref var partition = ref GetStorage(source);
+        ref var partition = ref GetPartition(source);
         ConditionalWeakTable<object, BackingStorage> newStorage;
         return Volatile.Read(ref partition) ?? Interlocked.CompareExchange(ref partition, newStorage = [], null) ?? newStorage;
     }

diff --git a/src/DotNext/UserDataStorage.cs b/src/DotNext/UserDataStorage.cs
@@ -122,15 +122,15 @@ internal UserDataStorage(object source)
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private BackingStorage? GetStorage()
         => source is not null &&
-           (source is BackingStorage storage || (Volatile.Read(in GetStorage(source))?.TryGetValue(source, out storage!) ?? false))
+           (source is BackingStorage storage || (Volatile.Read(in GetPartition(source))?.TryGetValue(source, out storage!) ?? false))
             ? storage
             : null;
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private BackingStorage GetOrCreateStorage()
     {
         if (source is not BackingStorage storage)
-            storage = GetOrCreateStorage(source).GetOrCreateValue(source);
+            storage = GetOrCreatePartition(source).GetOrCreateValue(source);
 
         return storage;
     }
@@ -380,12 +380,17 @@ public void CopyTo(object obj)
                 break;
         }
 
-        if (GetStorage() is { } source)
+        if (GetStorage() is not { } source)
         {
-            if (obj is BackingStorage destination)
-                source.CopyTo(destination);
-            else
-                GetOrCreateStorage(obj).AddOrUpdate(obj, source.Copy());
+            Volatile.Read(in GetPartition(obj))?.Remove(obj);
+        }
+        else if (obj is BackingStorage destination)
+        {
+            source.CopyTo(destination);
+        }
+        else
+        {
+            GetOrCreatePartition(obj).AddOrUpdate(obj, source.Copy());
         }
     }