Only image hash on file hash collision

filesmap.go

@@ -17,6 +17,10 @@ type FilesMap struct {
 	Images      []*imageEntry
 	FilesBySize map[int64]string
 	FilesByHash map[string][]string
+	// fileSizes records the size of every file found during the walk, so image
+	// candidates can be hashed without re-stat'ing the files. It is not
+	// serialized to the results file.
+	fileSizes       map[string]int64
 	FilesHashing    chan *fileEntry
 	FilesHashed     chan *fileEntry
 	ImagesHashing   chan *imageEntry
@@ -32,15 +36,18 @@ func newFilesMap() *FilesMap {
 	return &FilesMap{
 		FilesBySize:   map[int64]string{},
 		FilesByHash:   map[string][]string{},
-		FilesHashed:   make(chan *fileEntry, 100000),
+		fileSizes:     map[string]int64{},
 		FilesHashing:  make(chan *fileEntry),
-		ImagesHashed:  make(chan *imageEntry, 100000),
+		FilesHashed:   make(chan *fileEntry, 100000),
 		ImagesHashing: make(chan *imageEntry),
+		ImagesHashed:  make(chan *imageEntry, 100000),
 		progress:      mpb.New(mpb.WithWidth(64)),
 	}
 }
 
+// Stage 2: compute the file hash of every file sent over FilesHashing.
 func (fm *FilesMap) FileHashingWorker(wg *sync.WaitGroup) {
+	defer wg.Done()
 	for file := range fm.FilesHashing {
 		if *verbose {
 			fmt.Fprintf(fm.progress, "Hashing file %s\n", file.path)
@@ -56,10 +63,25 @@ func (fm *FilesMap) FileHashingWorker(wg *sync.WaitGroup) {
 		file.hash = hash
 		fm.FilesHashed <- file
 	}
-	wg.Done()
 }
 
+// Stage 2: collect the file hashes into FilesByHash once FilesHashed is closed.
+func (fm *FilesMap) fileHashMergeWorker(done chan bool) {
+	for file := range fm.FilesHashed {
+		if *verbose {
+			fmt.Printf("Finishing '%s' with hash '%s'\n", file.path, file.hash)
+		}
+
+		fm.lock.Lock()
+		fm.FilesByHash[file.hash] = append(fm.FilesByHash[file.hash], file.path)
+		fm.lock.Unlock()
+	}
+	done <- true
+}
+
+// Stage 4: compute the image hash of every candidate sent over ImagesHashing.
 func (fm *FilesMap) ImageHashingWorker(wg *sync.WaitGroup) {
+	defer wg.Done()
 	for file := range fm.ImagesHashing {
 		if *verbose {
 			fmt.Fprintf(fm.progress, "Hashing image %s\n", file.path)
@@ -78,20 +100,10 @@ func (fm *FilesMap) ImageHashingWorker(wg *sync.WaitGroup) {
 		file.imageHash = hash
 		fm.ImagesHashed <- file
 	}
-	wg.Done()
 }
 
-func (fm *FilesMap) HashedWorker(done chan bool) {
+// Stage 4: collect the image hashes into Images once ImagesHashed is closed.
-	for file := range fm.FilesHashed {
+func (fm *FilesMap) imageHashMergeWorker(done chan bool) {
-		if *verbose {
-			fmt.Printf("Finishing '%s' with hash '%s'\n", file.path, file.hash)
-		}
-
-		fm.lock.Lock()
-		fm.FilesByHash[file.hash] = append(fm.FilesByHash[file.hash], file.path)
-		fm.lock.Unlock()
-	}
-
 	for file := range fm.ImagesHashed {
 		if *verbose {
 			fmt.Println("Finishing", file.path)
@@ -101,10 +113,13 @@ func (fm *FilesMap) HashedWorker(done chan bool) {
 		fm.Images = append(fm.Images, file)
 		fm.lock.Unlock()
 	}
-
 	done <- true
 }
 
+// Stage 1: walk the trees and deduplicate by file size.
+// A file is only file-hashed once another file with the same size has been
+// seen; a unique size cannot be an exact duplicate. It returns the number of
+// files found.
 func (fm *FilesMap) WalkDirectories() int64 {
 	var countFiles int64 = 0
 	sumSize := int64(0)
@@ -122,7 +137,7 @@ func (fm *FilesMap) WalkDirectories() int64 {
 			countFiles++
 			fm.incomingBar.SetTotal(int64(countFiles), false)
 
-			fm.hashImage(path, size)
+			fm.fileSizes[path] = size
 			count := fm.hashFile(path, size)
 			if count > 0 {
 				sumSize += size * count
@@ -134,10 +149,11 @@ func (fm *FilesMap) WalkDirectories() int64 {
 
 	fm.incomingBar.SetTotal(int64(countFiles), true)
 	close(fm.FilesHashing)
-	close(fm.ImagesHashing)
 	return countFiles
 }
 
+// hashFile applies the file size step for a single file. It returns how many
+// files were queued to be file-hashed.
 func (fm *FilesMap) hashFile(path string, size int64) int64 {
 	prevFile, ok := fm.FilesBySize[size]
 	if !ok {
@@ -160,10 +176,6 @@ func (fm *FilesMap) hashFile(path string, size int64) int64 {
 	return 1
 }
 
-func (fm *FilesMap) hashImage(path string, size int64) {
-	fm.ImagesHashing <- &imageEntry{path, size, 0}
-}
-
 type imageEntry struct {
 	path      string
 	size      int64

main.go

@@ -86,22 +86,62 @@ func main() {
 				decor.CurrentKibiByte("%5d"),
 			),
 		)
-		done := make(chan bool)
+		// Stage 2: compute the file hash of every file found.
-		wg := sync.WaitGroup{}
+		fileHashingDone := make(chan bool)
+		fileWG := sync.WaitGroup{}
 		for i := 0; i < runtime.GOMAXPROCS(0); i++ {
-			wg.Add(2)
+			fileWG.Add(1)
-			go filesMap.ImageHashingWorker(&wg)
+			go filesMap.FileHashingWorker(&fileWG)
-			go filesMap.FileHashingWorker(&wg)
 		}
-
+		go filesMap.fileHashMergeWorker(fileHashingDone)
-		go filesMap.HashedWorker(done)
 
 		countFiles = filesMap.WalkDirectories()
 
-		wg.Wait()
+		fileWG.Wait()
 		close(filesMap.FilesHashed)
+		<-fileHashingDone
+
+		// Stage 3: image hashes are only needed for files with no file-hash
+		// collision: either they were never file-hashed (unique size) or their
+		// file hash did not match any other file.
+		var imageCandidates []*imageEntry
+
+		// Files whose size is unique were never file-hashed, so they have no
+		// collision but may still be visually similar images.
+		for size, file := range filesMap.FilesBySize {
+			if file == "" {
+				continue
+			}
+			imageCandidates = append(imageCandidates, &imageEntry{path: file, size: size})
+		}
+
+		// Files that were file-hashed without colliding also need an image hash.
+		for _, files := range filesMap.FilesByHash {
+			if len(files) != 1 {
+				continue
+			}
+			imageCandidates = append(imageCandidates, &imageEntry{path: files[0], size: filesMap.fileSizes[files[0]]})
+		}
+
+		// Stage 4: compute the image hash of every candidate.
+		go func() {
+			for _, candidate := range imageCandidates {
+				filesMap.ImagesHashing <- candidate
+			}
+			close(filesMap.ImagesHashing)
+		}()
+
+		imageWG := sync.WaitGroup{}
+		imageHashingDone := make(chan bool)
+		for i := 0; i < runtime.GOMAXPROCS(0); i++ {
+			imageWG.Add(1)
+			go filesMap.ImageHashingWorker(&imageWG)
+		}
+		go filesMap.imageHashMergeWorker(imageHashingDone)
+
+		imageWG.Wait()
 		close(filesMap.ImagesHashed)
-		<-done
+		<-imageHashingDone
 	}
 
 	if *toFile != "" && *fromFile == "" {