perceptual-hash-0.1.4.2: app/Parallel.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Parallel ( pathMaps ) where

import           Control.Concurrent.STM      (atomically)
import           Control.Concurrent.STM.TVar (TVar, modifyTVar', newTVarIO, readTVarIO)
import           Data.Functor                (($>))
import           Data.List.NonEmpty          (NonEmpty (..), (<|))
import qualified Data.Map                    as M
import           Data.Word                   (Word64)
import           PerceptualHash              (fileHash)
import           System.Directory.Parallel   (parTraverseAll)
import           System.FilePath             (takeExtension)

imgExtension :: String -> Bool
imgExtension ".jpg"  = True
imgExtension ".jpeg" = True
imgExtension ".png"  = True
imgExtension ".gif"  = True
imgExtension ".hdr"  = True
imgExtension ".pic"  = True
imgExtension ".bmp"  = True
imgExtension ".TGA"  = True
imgExtension ".tga"  = True
imgExtension ".tif"  = True
imgExtension ".tiff" = True
imgExtension ".webp" = True
imgExtension _       = False

insertHash :: FilePath -> IO (M.Map Word64 (NonEmpty FilePath) -> M.Map Word64 (NonEmpty FilePath))
insertHash fp = do
    hash <- fileHash fp
    case hash of
        Right x ->
            pure $ \hashes ->
                let go (Just others) = Just (fp <| others)
                    go Nothing       = Just (fp :| [])
                    in M.alter go x hashes
        Left err -> putStrLn ("WARNING: skipping " ++ fp ++ "\n" ++ err) $> id

stepMap :: TVar (M.Map Word64 (NonEmpty FilePath)) -> FilePath -> IO ()
stepMap var fp = do
    mod' <- insertHash fp
    atomically $ modifyTVar' var mod'

pathMaps :: [FilePath] -> IO (M.Map Word64 (NonEmpty FilePath))
pathMaps fps = do
    total <- newTVarIO mempty
    parTraverseAll (stepMap total) fileFilter (\_ -> pure True) fps
    readTVarIO total

    where fileFilter = pure . imgExtension . takeExtension