From 5cfe61928e3dbd877898e7c553389c884bc812ab Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Tue, 12 Oct 2021 05:39:10 +0300
Subject: [PATCH] tree: use oids instead of objects (#6785)

---
 dvc/dependency/repo.py              |  8 ++----
 dvc/fs/dvc.py                       | 12 ++++-----
 dvc/objects/__init__.py             |  4 +--
 dvc/objects/checkout.py             |  9 +++----
 dvc/objects/diff.py                 | 39 ++++++++++++++---------------
 dvc/objects/stage.py                | 15 ++---------
 dvc/objects/status.py               |  8 +++---
 dvc/objects/transfer.py             |  4 +--
 dvc/objects/tree.py                 | 39 +++++++++++++----------------
 dvc/output.py                       | 13 +++++-----
 dvc/repo/diff.py                    |  4 +--
 tests/func/objects/db/test_index.py |  8 +++---
 tests/func/test_checkout.py         |  4 +--
 tests/func/test_data_cloud.py       |  4 +--
 tests/func/test_remote.py           |  8 +++---
 tests/unit/objects/test_tree.py     | 27 ++++++++++----------
 16 files changed, 91 insertions(+), 115 deletions(-)

diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py
index 21d96a889..f758eb58b 100644
--- a/dvc/dependency/repo.py
+++ b/dvc/dependency/repo.py
@@ -148,9 +148,7 @@ class RepoDependency(Dependency):
             self._staged_objs[rev] = staged_obj
             used_obj_ids[staging].add(staged_obj.hash_info)
             if isinstance(staged_obj, Tree):
-                used_obj_ids[staging].update(
-                    entry.hash_info for _, _, entry in staged_obj
-                )
+                used_obj_ids[staging].update(oid for _, _, oid in staged_obj)
             return used_obj_ids, staged_obj
 
     def _check_circular_import(self, odb, obj_ids):
@@ -166,9 +164,7 @@ class RepoDependency(Dependency):
             for hash_info in obj_ids:
                 if hash_info.isdir:
                     tree = Tree.load(odb, hash_info)
-                    yield from (
-                        odb.get(entry.hash_info) for _, _, entry in tree
-                    )
+                    yield from (odb.get(oid) for _, _, oid in tree)
                 else:
                     yield odb.get(hash_info)
 
diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py
index fe8b5f15d..90199d1d2 100644
--- a/dvc/fs/dvc.py
+++ b/dvc/fs/dvc.py
@@ -58,9 +58,9 @@ class DvcFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
         out.get_dir_cache(remote=remote)
         if out.obj is None:
             raise FileNotFoundError
-        (_, obj) = out.obj.trie.get(key) or (None, None)
-        if obj:
-            return obj.hash_info
+        (_, oid) = out.obj.trie.get(key) or (None, None)
+        if oid:
+            return oid
         raise FileNotFoundError
 
     def _get_fs_path(self, path: PathInfo, remote=None):
@@ -250,10 +250,10 @@ class DvcFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
         elif meta.part_of_output:
             (out,) = meta.outs
             key = path_info.relative_to(out.path_info).parts
-            (obj_meta, obj) = out.obj.trie.get(key) or (None, None)
-            if obj:
+            (obj_meta, oid) = out.obj.trie.get(key) or (None, None)
+            if oid:
                 ret["size"] = obj_meta.size if obj_meta else 0
-                ret[obj.hash_info.name] = obj.hash_info.value
+                ret[oid.name] = oid.value
 
         return ret
 
diff --git a/dvc/objects/__init__.py b/dvc/objects/__init__.py
index a57b8f3c3..ad45ba490 100644
--- a/dvc/objects/__init__.py
+++ b/dvc/objects/__init__.py
@@ -14,8 +14,8 @@ logger = logging.getLogger(__name__)
 
 def check(odb: "ObjectDB", obj: "HashFile", **kwargs):
     if isinstance(obj, Tree):
-        for _, _, entry in obj:
-            odb.check(entry.hash_info, **kwargs)
+        for _, _, oid in obj:
+            odb.check(oid, **kwargs)
 
     odb.check(obj.hash_info, **kwargs)
 
diff --git a/dvc/objects/checkout.py b/dvc/objects/checkout.py
index 834959459..020df6a6a 100644
--- a/dvc/objects/checkout.py
+++ b/dvc/objects/checkout.py
@@ -17,7 +17,6 @@ from dvc.objects.db.slow_link_detection import (  # type: ignore[attr-defined]
 from dvc.objects.diff import ROOT
 from dvc.objects.diff import diff as odiff
 from dvc.objects.stage import stage
-from dvc.objects.tree import Tree
 from dvc.types import Optional
 
 logger = logging.getLogger(__name__)
@@ -131,8 +130,8 @@ def _checkout_file(
 ):
     """The file is changed we need to checkout a new copy"""
     modified = False
-    cache_info = cache.hash_to_path_info(change.new.obj.hash_info.value)
-    if change.old.obj:
+    cache_info = cache.hash_to_path_info(change.new.oid.value)
+    if change.old.oid:
         if relink:
             if fs.iscopy(path_info) and _cache_is_copy(cache, path_info):
                 cache.unprotect(path_info)
@@ -148,7 +147,7 @@ def _checkout_file(
         modified = True
 
     if state:
-        state.save(path_info, fs, change.new.obj.hash_info)
+        state.save(path_info, fs, change.new.oid)
 
     if progress_callback:
         progress_callback(str(path_info))
@@ -213,7 +212,7 @@ def _checkout(
             if change.new.key != ROOT
             else path_info
         )
-        if isinstance(change.new.obj, Tree):
+        if change.new.oid.isdir:
             fs.makedirs(entry_path)
             continue
 
diff --git a/dvc/objects/diff.py b/dvc/objects/diff.py
index 62f60047f..52ee299b7 100644
--- a/dvc/objects/diff.py
+++ b/dvc/objects/diff.py
@@ -2,6 +2,8 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
 if TYPE_CHECKING:
+    from dvc.hash_info import HashInfo
+
     from .file import HashFile
 
 ADD = "add"
@@ -14,22 +16,19 @@ UNCHANGED = "unchanged"
 class TreeEntry:
     in_cache: bool
     key: Tuple[str]
-    obj: Optional["HashFile"] = field(default=None)
+    oid: Optional["HashInfo"] = field(default=None)
 
     def __bool__(self):
-        return bool(self.obj)
+        return bool(self.oid)
 
     def __eq__(self, other):
         if not isinstance(other, TreeEntry):
             return False
 
-        if self.key != other.key or bool(self.obj) != bool(other.obj):
+        if self.key != other.key:
             return False
 
-        if not self.obj:
-            return False
-
-        return self.obj.hash_info == other.obj.hash_info
+        return self.oid == other.oid
 
 
 @dataclass
@@ -89,33 +88,33 @@ def diff(
     old_keys = set(_get_keys(old))
     new_keys = set(_get_keys(new))
 
-    def _get_obj(obj, key):
+    def _get_oid(obj, key):
         if not obj or key == ROOT:
-            return obj
+            return obj.hash_info if obj else None
 
-        return obj.get(key)
+        entry_obj = obj.get(cache, key)
+        return entry_obj.hash_info if entry_obj else None
 
-    def _in_cache(obj, cache):
-        from . import check
+    def _in_cache(oid, cache):
         from .errors import ObjectFormatError
 
-        if not obj:
+        if not oid:
             return False
 
         try:
-            check(cache, obj)
+            cache.check(oid)
             return True
         except (FileNotFoundError, ObjectFormatError):
             return False
 
     ret = DiffResult()
     for key in old_keys | new_keys:
-        old_obj = _get_obj(old, key)
-        new_obj = _get_obj(new, key)
+        old_oid = _get_oid(old, key)
+        new_oid = _get_oid(new, key)
 
         change = Change(
-            old=TreeEntry(_in_cache(old_obj, cache), key, old_obj),
-            new=TreeEntry(_in_cache(new_obj, cache), key, new_obj),
+            old=TreeEntry(_in_cache(old_oid, cache), key, old_oid),
+            new=TreeEntry(_in_cache(new_oid, cache), key, new_oid),
         )
 
         if change.typ == ADD:
@@ -126,8 +125,8 @@ def diff(
             ret.deleted.append(change)
         else:
             assert change.typ == UNCHANGED
-            if not change.new.in_cache and not isinstance(
-                change.new.obj, Tree
+            if not change.new.in_cache and not (
+                change.new.oid and change.new.oid.isdir
             ):
                 ret.modified.append(change)
             else:
diff --git a/dvc/objects/stage.py b/dvc/objects/stage.py
index 8ed4f0f3e..2458c0530 100644
--- a/dvc/objects/stage.py
+++ b/dvc/objects/stage.py
@@ -90,6 +90,7 @@ def _stage_file(path_info, fs, name, odb=None, upload_odb=None, dry_run=False):
     else:
         odb.add(path_info, fs, hash_info, move=False)
         obj = odb.get(hash_info)
+
     return path_info, meta, obj
 
 
@@ -146,7 +147,7 @@ def _build_tree(path_info, fs, name, **kwargs):
         #
         # Yes, this is a BUG, as long as we permit "/" in
         # filenames on Windows and "\" on Unix
-        tree.add(file_info.relative_to(path_info).parts, meta, obj)
+        tree.add(file_info.relative_to(path_info).parts, meta, obj.hash_info)
 
         tree_meta.size += meta.size
         tree_meta.nfiles += 1
@@ -162,12 +163,6 @@ def _stage_tree(path_info, fs, fs_info, name, odb=None, **kwargs):
         hash_info = HashInfo(name, value)
         try:
             tree = Tree.load(odb, hash_info)
-            # NOTE: loaded entries are naive objects with hash_infos but no
-            # path_info. For staging trees, obj.path_info should be relative
-            # to the staging src `path_info` and src fs
-            for key, _, entry in tree:
-                entry.fs = fs
-                entry.path_info = path_info.joinpath(*key)
             return Meta(nfiles=len(tree)), tree
         except FileNotFoundError:
             pass
@@ -239,12 +234,6 @@ def _load_from_state(odb, staging, path_info, fs, name):
                     check(odb_, obj, check_hash=False)
                     if isinstance(obj, Tree):
                         meta.nfiles = len(obj)
-                        for key, _, entry in obj:
-                            entry.fs = fs
-                            entry.path_info = path_info.joinpath(*key)
-                    else:
-                        obj.fs = fs
-                        obj.path_info = path_info
                     assert obj.hash_info.name == name
                     return odb_, meta, obj
                 except (ObjectFormatError, FileNotFoundError):
diff --git a/dvc/objects/status.py b/dvc/objects/status.py
index f79908b6f..d0f7f163a 100644
--- a/dvc/objects/status.py
+++ b/dvc/objects/status.py
@@ -56,7 +56,7 @@ def _indexed_dir_hashes(odb, index, dir_objs, name, cache_odb):
                 tree = Tree.load(cache_odb, HashInfo(name, dir_hash))
             except FileNotFoundError:
                 continue
-        file_hashes = [entry.hash_info.value for _, _, entry in tree]
+        file_hashes = [oid.value for _, _, oid in tree]
         if dir_hash not in index:
             logger.debug(
                 "Indexing new .dir '%s' with '%s' nested files",
@@ -102,9 +102,9 @@ def status(
                 tree = None
             else:
                 tree = Tree.load(cache_odb, hash_info)
-                for _, _, entry in tree:
-                    assert entry.hash_info and entry.hash_info.value
-                    hash_infos[entry.hash_info.value] = entry.hash_info
+                for _, _, oid in tree:
+                    assert oid and oid.value
+                    hash_infos[oid.value] = oid
             if index:
                 dir_objs[hash_info.value] = tree
         hash_infos[hash_info.value] = hash_info
diff --git a/dvc/objects/transfer.py b/dvc/objects/transfer.py
index 95aad3b70..00078d1ae 100644
--- a/dvc/objects/transfer.py
+++ b/dvc/objects/transfer.py
@@ -74,7 +74,7 @@ def _do_transfer(
         dir_obj = find_tree_by_obj_id([cache_odb, src], dir_hash)
         assert dir_obj
 
-        entry_ids = {entry.hash_info for _, _, entry in dir_obj}
+        entry_ids = {oid for _, _, oid in dir_obj}
         bound_file_ids = all_file_ids & entry_ids
         all_file_ids -= entry_ids
 
@@ -117,7 +117,7 @@ def _do_transfer(
     # index successfully pushed dirs
     if dest_index:
         for dir_obj in succeeded_dir_objs:
-            file_hashes = {entry.hash_info.value for _, _, entry in dir_obj}
+            file_hashes = {oid.value for _, _, oid in dir_obj}
             logger.debug(
                 "Indexing pushed dir '%s' with '%s' nested files",
                 dir_obj.hash_info,
diff --git a/dvc/objects/tree.py b/dvc/objects/tree.py
index e7be2e191..0914c882f 100644
--- a/dvc/objects/tree.py
+++ b/dvc/objects/tree.py
@@ -22,7 +22,7 @@ class Tree(HashFile):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._dict: Dict[Tuple[str], Tuple["Meta", "HashFile"]] = {}
+        self._dict: Dict[Tuple[str], Tuple["Meta", "HashInfo"]] = {}
 
     @cached_property
     def trie(self):
@@ -30,9 +30,9 @@ class Tree(HashFile):
 
         return Trie(self._dict)
 
-    def add(self, key, meta, obj):
+    def add(self, key: Tuple[str], meta: "Meta", oid: "HashInfo"):
         self.__dict__.pop("trie", None)
-        self._dict[key] = (meta, obj)
+        self._dict[key] = (meta, oid)
 
     def digest(self, hash_info: Optional["HashInfo"] = None):
         from dvc.fs.memory import MemoryFileSystem
@@ -72,10 +72,10 @@ class Tree(HashFile):
                 {
                     # NOTE: not using hash_info.to_dict() because we don't want
                     # size/nfiles fields at this point.
-                    obj.hash_info.name: obj.hash_info.value,
+                    oid.name: oid.value,
                     self.PARAM_RELPATH: posixpath.sep.join(parts),
                 }
-                for parts, _, obj in self  # noqa: B301
+                for parts, _, oid in self  # noqa: B301
             ),
             key=itemgetter(self.PARAM_RELPATH),
         )
@@ -93,8 +93,7 @@ class Tree(HashFile):
             relpath = entry.pop(cls.PARAM_RELPATH)
             parts = tuple(relpath.split(posixpath.sep))
             hash_info = HashInfo.from_dict(entry)
-            obj = HashFile(None, None, hash_info)
-            tree.add(parts, None, obj)
+            tree.add(parts, None, hash_info)
         return tree
 
     @classmethod
@@ -117,8 +116,6 @@ class Tree(HashFile):
         tree = cls.from_list(raw)
         tree.path_info = obj.path_info
         tree.fs = obj.fs
-        for _, _, entry_obj in tree:
-            entry_obj.fs = obj.fs
         tree.hash_info = hash_info
 
         return tree
@@ -134,26 +131,26 @@ class Tree(HashFile):
         """
         tree = Tree(self.path_info, self.fs, self.hash_info)
         try:
-            for key, (meta, obj) in self.trie.items(prefix):
-                tree.add(key, meta, obj)
+            for key, (meta, oid) in self.trie.items(prefix):
+                tree.add(key, meta, oid)
         except KeyError:
             pass
         return tree
 
-    def get(self, prefix: Tuple[str]) -> Optional[HashFile]:
+    def get(self, odb, prefix: Tuple[str]) -> Optional[HashFile]:
         """Return object at the specified prefix in this tree.
 
         Returns None if no object exists at the specified prefix.
         """
-        _, obj = self._dict.get(prefix) or (None, None)
-        if obj:
-            return obj
+        _, oid = self._dict.get(prefix) or (None, None)
+        if oid:
+            return odb.get(oid)
 
         tree = Tree(None, None, None)
         depth = len(prefix)
         try:
-            for key, (meta, obj) in self.trie.items(prefix):
-                tree.add(key[depth:], meta, obj)
+            for key, (meta, entry_oid) in self.trie.items(prefix):
+                tree.add(key[depth:], meta, entry_oid)
         except KeyError:
             return None
         tree.digest()
@@ -163,8 +160,8 @@ class Tree(HashFile):
 def du(odb, tree):
     try:
         return sum(
-            odb.fs.getsize(odb.hash_to_path_info(obj.hash_info.value))
-            for _, _, obj in tree
+            odb.fs.getsize(odb.hash_to_path_info(oid.value))
+            for _, _, oid in tree
         )
     except FileNotFoundError:
         return None
@@ -225,8 +222,8 @@ def merge(odb, ancestor_info, our_info, their_info):
     merged_dict = _merge(ancestor.as_dict(), our.as_dict(), their.as_dict())
 
     merged = Tree(None, None, None)
-    for key, (meta, obj) in merged_dict.items():
-        merged.add(key, meta, obj)
+    for key, (meta, oid) in merged_dict.items():
+        merged.add(key, meta, oid)
     merged.digest()
 
     return merged
diff --git a/dvc/output.py b/dvc/output.py
index b8f734945..814b51c6b 100644
--- a/dvc/output.py
+++ b/dvc/output.py
@@ -623,12 +623,11 @@ class Output:
             dvcignore=self.dvcignore,
         )
         save_obj = save_obj.filter(prefix)
-        checkout_obj = save_obj.get(prefix)
+        checkout_obj = save_obj.get(self.odb, prefix)
         otransfer(
             staging,
             self.odb,
-            {save_obj.hash_info}
-            | {entry.hash_info for _, _, entry in save_obj},
+            {save_obj.hash_info} | {oid for _, _, oid in save_obj},
             shallow=True,
             move=True,
         )
@@ -719,7 +718,7 @@ class Output:
 
         if filter_info and filter_info != self.path_info:
             prefix = filter_info.relative_to(self.path_info).parts
-            obj = obj.get(prefix)
+            obj = obj.get(self.odb, prefix)
 
         return obj
 
@@ -956,9 +955,9 @@ class Output:
         obj.hash_info.obj_name = name
         oids = {obj.hash_info}
         if isinstance(obj, Tree):
-            for key, _, entry_obj in obj:
-                entry_obj.hash_info.obj_name = self.fs.sep.join([name, *key])
-                oids.add(entry_obj.hash_info)
+            for key, _, oid in obj:
+                oid.obj_name = self.fs.sep.join([name, *key])
+                oids.add(oid)
         return oids
 
     def get_used_external(
diff --git a/dvc/repo/diff.py b/dvc/repo/diff.py
index 7a947a86c..f1ea7b635 100644
--- a/dvc/repo/diff.py
+++ b/dvc/repo/diff.py
@@ -164,13 +164,13 @@ def _output_paths(repo, targets):
 
 
 def _dir_output_paths(path_info, obj, targets=None):
-    for key, _, entry_obj in obj:
+    for key, _, oid in obj:
         fname = path_info.joinpath(*key)
         if targets is None or any(
             fname.isin_or_eq(target) for target in targets
         ):
             # pylint: disable=no-member
-            yield str(fname), entry_obj.hash_info.value
+            yield str(fname), oid.value
 
 
 def _filter_missing(repo_fs, paths):
diff --git a/tests/func/objects/db/test_index.py b/tests/func/objects/db/test_index.py
index 7e79c90c7..ead6fe732 100644
--- a/tests/func/objects/db/test_index.py
+++ b/tests/func/objects/db/test_index.py
@@ -16,7 +16,7 @@ def index(dvc, local_remote, mocker):
 def test_indexed_on_status(tmp_dir, dvc, index):
     foo = tmp_dir.dvc_gen({"foo": "foo content"})[0].outs[0]
     bar = tmp_dir.dvc_gen({"bar": {"baz": "baz content"}})[0].outs[0]
-    baz_hash = bar.obj.trie.get(("baz",))[1].hash_info
+    baz_hash = bar.obj.trie.get(("baz",))[1]
     clean_staging()
     dvc.push()
     index.clear()
@@ -30,7 +30,7 @@ def test_indexed_on_status(tmp_dir, dvc, index):
 def test_indexed_on_push(tmp_dir, dvc, index):
     foo = tmp_dir.dvc_gen({"foo": "foo content"})[0].outs[0]
     bar = tmp_dir.dvc_gen({"bar": {"baz": "baz content"}})[0].outs[0]
-    baz_hash = bar.obj.trie.get(("baz",))[1].hash_info
+    baz_hash = bar.obj.trie.get(("baz",))[1]
     clean_staging()
 
     dvc.push()
@@ -60,8 +60,8 @@ def test_clear_on_download_err(tmp_dir, dvc, index, mocker):
     out = tmp_dir.dvc_gen({"dir": {"foo": "foo content"}})[0].outs[0]
     dvc.push()
 
-    for _, _, entry in out.obj:
-        remove(dvc.odb.local.get(entry.hash_info).path_info)
+    for _, _, oid in out.obj:
+        remove(dvc.odb.local.get(oid).path_info)
     remove(out.path_info)
 
     assert list(index.hashes())
diff --git a/tests/func/test_checkout.py b/tests/func/test_checkout.py
index e4f374f3d..b9caff4ea 100644
--- a/tests/func/test_checkout.py
+++ b/tests/func/test_checkout.py
@@ -101,9 +101,9 @@ class TestCheckoutCorruptedCacheDir(TestDvc):
         # NOTE: modifying cache file for one of the files inside the directory
         # to check if dvc will detect that the cache is corrupted.
         obj = load(self.dvc.odb.local, out.hash_info)
-        _, _, entry_obj = list(obj)[0]
+        _, _, entry_oid = list(obj)[0]
         cache = os.fspath(
-            self.dvc.odb.local.hash_to_path_info(entry_obj.hash_info.value)
+            self.dvc.odb.local.hash_to_path_info(entry_oid.value)
         )
 
         os.chmod(cache, 0o644)
diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py
index d08aa7ca3..a62663c84 100644
--- a/tests/func/test_data_cloud.py
+++ b/tests/func/test_data_cloud.py
@@ -57,9 +57,7 @@ def test_cloud(tmp_dir, dvc, remote):  # pylint:disable=unused-argument
     out_dir = stage_dir.outs[0]
     cache_dir = out_dir.cache_path
     dir_hash = out_dir.hash_info
-    dir_hashes = {dir_hash} | {
-        entry_obj.hash_info for _, _, entry_obj in out_dir.obj
-    }
+    dir_hashes = {dir_hash} | {oid for _, _, oid in out_dir.obj}
 
     def _check_status(status, **kwargs):
         for key in ("ok", "missing", "new", "deleted"):
diff --git a/tests/func/test_remote.py b/tests/func/test_remote.py
index e0072600e..9c2533915 100644
--- a/tests/func/test_remote.py
+++ b/tests/func/test_remote.py
@@ -274,7 +274,7 @@ def test_push_order(tmp_dir, dvc, tmp_path_factory, mocker, local_remote):
     # foo .dir file should be uploaded after bar
     odb = dvc.cloud.get_remote_odb("upstream")
     foo_path = odb.hash_to_path_info(foo.hash_info.value)
-    bar_path = odb.hash_to_path_info(foo.obj.trie[("bar",)][1].hash_info.value)
+    bar_path = odb.hash_to_path_info(foo.obj.trie[("bar",)][1].value)
     paths = [args[1] for args, _ in mocked_upload.call_args_list]
     assert paths.index(foo_path) > paths.index(bar_path)
 
@@ -427,12 +427,12 @@ def test_push_incomplete_dir(tmp_dir, dvc, mocker, local_remote):
     file_objs = [entry_obj for _, _, entry_obj in out.obj]
 
     # remove one of the cache files for directory
-    remove(odb.hash_to_path_info(file_objs[0].hash_info.value))
+    remove(odb.hash_to_path_info(file_objs[0].value))
 
     dvc.push()
     assert not remote_odb.exists(out.hash_info)
-    assert not remote_odb.exists(file_objs[0].hash_info)
-    assert remote_odb.exists(file_objs[1].hash_info)
+    assert not remote_odb.exists(file_objs[0])
+    assert remote_odb.exists(file_objs[1])
 
 
 def test_upload_exists(tmp_dir, dvc, local_remote):
diff --git a/tests/unit/objects/test_tree.py b/tests/unit/objects/test_tree.py
index cc6d1783d..953d3d8a2 100644
--- a/tests/unit/objects/test_tree.py
+++ b/tests/unit/objects/test_tree.py
@@ -3,7 +3,6 @@ from operator import itemgetter
 import pytest
 
 from dvc.hash_info import HashInfo
-from dvc.objects.file import HashFile
 from dvc.objects.meta import Meta
 from dvc.objects.tree import Tree, _merge
 
@@ -20,10 +19,10 @@ from dvc.objects.tree import Tree, _merge
                 {"md5": "456", "relpath": "bar"},
             ],
             {
-                ("zzz",): (None, HashFile(None, None, HashInfo("md5", "def"))),
-                ("foo",): (None, HashFile(None, None, HashInfo("md5", "123"))),
-                ("bar",): (None, HashFile(None, None, HashInfo("md5", "456"))),
-                ("aaa",): (None, HashFile(None, None, HashInfo("md5", "abc"))),
+                ("zzz",): (None, HashInfo("md5", "def")),
+                ("foo",): (None, HashInfo("md5", "123")),
+                ("bar",): (None, HashInfo("md5", "456")),
+                ("aaa",): (None, HashInfo("md5", "abc")),
             },
         ),
         (
@@ -41,30 +40,30 @@ from dvc.objects.tree import Tree, _merge
             {
                 ("dir", "b"): (
                     None,
-                    HashFile(None, None, HashInfo("md5", "123")),
+                    HashInfo("md5", "123"),
                 ),
                 ("dir", "z"): (
                     None,
-                    HashFile(None, None, HashInfo("md5", "456")),
+                    HashInfo("md5", "456"),
                 ),
                 ("dir", "a"): (
                     None,
-                    HashFile(None, None, HashInfo("md5", "789")),
+                    HashInfo("md5", "789"),
                 ),
-                ("b",): (None, HashFile(None, None, HashInfo("md5", "abc"))),
-                ("a",): (None, HashFile(None, None, HashInfo("md5", "def"))),
-                ("z",): (None, HashFile(None, None, HashInfo("md5", "ghi"))),
+                ("b",): (None, HashInfo("md5", "abc")),
+                ("a",): (None, HashInfo("md5", "def")),
+                ("z",): (None, HashInfo("md5", "ghi")),
                 ("dir", "subdir", "b"): (
                     None,
-                    HashFile(None, None, HashInfo("md5", "jkl")),
+                    HashInfo("md5", "jkl"),
                 ),
                 ("dir", "subdir", "z"): (
                     None,
-                    HashFile(None, None, HashInfo("md5", "mno")),
+                    HashInfo("md5", "mno"),
                 ),
                 ("dir", "subdir", "a"): (
                     None,
-                    HashFile(None, None, HashInfo("md5", "pqr")),
+                    HashInfo("md5", "pqr"),
                 ),
             },
         ),