From 8d6813a11244301bd89c1a38225b6e333abb8bfc Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Mon, 2 Jul 2018 18:54:33 +0300
Subject: [PATCH] checkout: use link_state properly

Signed-off-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/project.py      | 14 ++++++++++++--
 dvc/remote/local.py | 10 +++++-----
 dvc/state.py        | 16 ++++++++++------
 3 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/dvc/project.py b/dvc/project.py
index 09106c2d0..c71a0b2b7 100644
--- a/dvc/project.py
+++ b/dvc/project.py
@@ -241,14 +241,24 @@ class Project(object):
                 raise ReproductionError(stages[n].relpath, ex)
         return result
 
+    def _cleanup_unused_links(self, all_stages):
+        used = []
+        for stage in all_stages:
+            for out in stage.outs:
+                used.append(out.path)
+        self.link_state.remove_unused(used)
+
     def checkout(self, target=None):
+        all_stages = self.stages()
+
         if target:
             if not Stage.is_stage_file(target):
                 raise NotDvcFileError(target)
             stages = [Stage.load(self, target)]
         else:
-            self.link_state.remove_all()
-            stages = self.stages()
+            stages = all_stages
+
+        self._cleanup_unused_links(all_stages)
 
         for stage in stages:
             stage.checkout()
diff --git a/dvc/remote/local.py b/dvc/remote/local.py
index 7ddb6ab95..3aa7ae519 100644
--- a/dvc/remote/local.py
+++ b/dvc/remote/local.py
@@ -85,7 +85,7 @@ class RemoteLOCAL(RemoteBase):
 
         return False
 
-    def link(self, md5, path, dump=True):
+    def link(self, md5, path):
         cache = self.get(md5)
         if not cache or not os.path.exists(cache) or self.changed(md5):
             if cache:
@@ -105,7 +105,6 @@ class RemoteLOCAL(RemoteBase):
         for typ in types:
             try:
                 self.CACHE_TYPE_MAP[typ](cache, path)
-                self.link_state.update(path, dump=dump)
                 return
             except Exception as exc:
                 msg = 'Cache type \'{}\' is not supported'.format(typ)
@@ -152,7 +151,8 @@ class RemoteLOCAL(RemoteBase):
         Logger.debug(msg.format(os.path.relpath(path), md5))
 
         if not self.is_dir_cache(cache):
-            self.link(md5, path, dump=True)
+            self.link(md5, path)
+            self.link_state.update(path)
             return
 
         # Create dir separately so that dir is created
@@ -164,8 +164,8 @@ class RemoteLOCAL(RemoteBase):
             md5 = entry[self.PARAM_MD5]
             relpath = entry[self.PARAM_RELPATH]
             p = os.path.join(path, relpath)
-            self.link(md5, p, dump=False)
-        self.link_state.dump()
+            self.link(md5, p)
+        self.link_state.update(path)
 
     def _move(self, inp, outp):
         # moving in two stages to make last the move atomic in
diff --git a/dvc/state.py b/dvc/state.py
index 4f172e2c7..22c5b8fd4 100644
--- a/dvc/state.py
+++ b/dvc/state.py
@@ -209,11 +209,15 @@ class LinkState(State):
                 with self._lock_file:
                     self.dump()
 
-    def _do_remove_all(self):
-        for p, s in self._db.items():
+    def _do_remove_unused(self, used):
+        items = self._db.copy().items()
+        for p, s in items:
             path = os.path.join(self.root_dir, p)
             state = LinkStateEntry.loadd(s)
 
+            if path in used:
+                continue
+
             if not os.path.exists(path):
                 continue
 
@@ -221,12 +225,12 @@ class LinkState(State):
             mtime = self.mtime(path)
 
             if inode == state.inode and mtime == state.mtime:
+                Logger.debug('Removing \'{}\' as unused link.'.format(path))
                 remove(path)
+                del self._db[p]
 
-        self._db = {}
-
-    def remove_all(self):
+    def remove_unused(self, used):
         with self._lock:
             with self._lock_file:
-                self._do_remove_all()
+                self._do_remove_unused(used)
                 self.dump()