--- a/contrib/hg-relink	Wed Mar 21 14:06:25 2007 -0700
+++ b/contrib/hg-relink	Fri Mar 23 14:11:31 2007 -0700
@@ -23,20 +23,13 @@
             if not os.path.exists(os.path.join(d, '.hg')):
                 raise ConfigError("%s: not a mercurial repository" % d)
 
-try:
-    cfg = Config(sys.argv)
-except ConfigError, inst:
-    print str(inst)
-    usage()
-    sys.exit(1)
-
 def collect(src):
     seplen = len(os.path.sep)
     candidates = []
     for dirpath, dirnames, filenames in os.walk(src):
         relpath = dirpath[len(src) + seplen:]
         for filename in filenames:
-            if not (filename.endswith('.i') or filename.endswith('.d')):
+            if not filename.endswith('.i'):
                 continue
             st = os.stat(os.path.join(dirpath, filename))
             candidates.append((os.path.join(relpath, filename), st))
@@ -44,25 +37,56 @@
     return candidates
 
 def prune(candidates, dst):
+    def getdatafile(path):
+        if not path.endswith('.i'):
+            return None, None
+        df = path[:-1] + 'd'
+        try:
+            st = os.stat(df)
+        except OSError:
+            return None, None
+        return df, st
+
+    def linkfilter(dst, st):
+        try:
+            ts = os.stat(dst)
+        except OSError:
+            # Destination doesn't have this file?
+            return False
+        if st.st_ino == ts.st_ino:
+            return False
+        if st.st_dev != ts.st_dev:
+            # No point in continuing
+            raise Exception('Source and destination are on different devices')
+        if st.st_size != ts.st_size:
+            # TODO: compare revlog heads
+            return False
+        return st
+
     targets = []
     for fn, st in candidates:
         tgt = os.path.join(dst, fn)
-        try:
-            ts = os.stat(tgt)
-        except OSError:
-            # Destination doesn't have this file?
-            continue
-        if st.st_ino == ts.st_ino:
-            continue
-        if st.st_dev != ts.st_dev:
-            raise Exception('Source and destination are on different devices')
-        if st.st_size != ts.st_size:
+        ts = linkfilter(tgt, st)
+        if not ts:
             continue
         targets.append((fn, ts.st_size))
+        df, ts = getdatafile(tgt)
+        if df:
+            targets.append((fn[:-1] + 'd', ts.st_size))
 
     return targets
 
 def relink(src, dst, files):
+    def relinkfile(src, dst):
+        bak = dst + '.bak'
+        os.rename(dst, bak)
+        try:
+            os.link(src, dst)
+        except OSError:
+            os.rename(bak, dst)
+            raise
+        os.remove(bak)
+
     CHUNKLEN = 65536
     relinked = 0
     savedbytes = 0
@@ -81,21 +105,22 @@
         if sin:
             continue
         try:
-            os.rename(tgt, tgt + '.bak')
-            try:
-                os.link(source, tgt)
-            except OSError:
-                os.rename(tgt + '.bak', tgt)
-                raise
+            relinkfile(source, tgt)
             print 'Relinked %s' % f
             relinked += 1
             savedbytes += sz
-            os.remove(tgt + '.bak')
         except OSError, inst:
             print '%s: %s' % (tgt, str(inst))
 
     print 'Relinked %d files (%d bytes reclaimed)' % (relinked, savedbytes)
 
+try:
+    cfg = Config(sys.argv)
+except ConfigError, inst:
+    print str(inst)
+    usage()
+    sys.exit(1)
+
 src = os.path.join(cfg.src, '.hg')
 dst = os.path.join(cfg.dst, '.hg')
 candidates = collect(src)