distinguish between target and key in Alias

fjetter · fjetter · commit 60df5d92356e · 2024-08-28T11:25:41.000+02:00
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
@@ -325,26 +325,24 @@ def resolve_aliases(dsk: dict, keys: set, dependents: dict) -> dict:
         raise ValueError("No keys provided")
     dsk = dict(dsk)
     work = list(keys)
-    seen = set()
     while work:
         k = work.pop()
-        if k in seen or k not in dsk:
+        if k not in dsk:
             continue
-        seen.add(k)
         t = dsk[k]
-        if (
-            isinstance(t, Alias)
-            and t.key not in keys
-            and t.key != k
-            and t.key in dsk
-            and len(dependents[t.key]) == 1
-        ):
-            t = dsk[k] = dsk.pop(t.key).copy()
-            seen.discard(k)
-            if isinstance(t, Alias):
-                work.append(k)
-            else:
+        if isinstance(t, Alias):
+            target_key = t.target.key
+
+            if (
+                target_key not in keys
+                and target_key != k
+                and target_key in dsk
+                and len(dependents[target_key]) == 1
+            ):
+                t = dsk[k] = dsk.pop(target_key).copy()
                 t.key = k
+                if isinstance(t, Alias):
+                    work.append(k)
 
         work.extend(t.dependencies)
     return dsk
@@ -417,19 +415,26 @@ def __sizeof__(self) -> int:
 
 class Alias(BaseTask):
     __weakref__: Any = None
+    target: KeyRef
     __slots__ = tuple(__annotations__)
 
-    def __init__(self, key: KeyRef | KeyType):
-        if isinstance(key, KeyRef):
-            key = key.key
+    def __init__(self, key: KeyType, target: KeyRef | KeyType | None = None):
         self.key = key
+        if target is None:
+            target = key
+        if not isinstance(target, KeyRef):
+            target = KeyRef(target)
+        self.target = target
         self.dependencies = {key}
 
+    def ref(self):
+        return self.target
+
     def copy(self):
-        return Alias(self.key)
+        return Alias(self.key, self.target)
 
     def __reduce__(self) -> str | tuple[Any, ...]:
-        return Alias, (self.key,)
+        return Alias, (self.key, self.target)
 
     def __call__(self, values=()):
         self._verify_values(values)
@@ -448,7 +453,7 @@ def inline(self, dsk) -> BaseTask:
             return self
 
     def __repr__(self):
-        return f"Alias({self.key})"
+        return f"Alias(key={self.key}, target={self.target})"
 
     def __eq__(self, value: object) -> bool:
         if not isinstance(value, Alias):
