sizes-2.4.4: Sizes.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}

module Sizes (
    -- * Main entry point
    sizesMain,

    -- * Types
    SizesOpts (..),
    EntryInfo (..),
    ReportEntries,
    SeenInodes,

    -- * Lenses
    entryPath,
    entryCount,
    entryAllocSize,
    entryIsDir,

    -- * Functions
    humanReadable,
    reportEntryP,
    newEntry,
    emptyReportEntries,
    reportEntriesToList,
    combineEntryResults,
    crossesFileSystemBoundary,
) where

import Control.Concurrent.ParallelIO
import Control.DeepSeq
import Control.Exception hiding (catch)
import Control.Lens
import Control.Monad
import Control.Monad.Catch (catch)
import Control.Monad.State.Strict
import qualified Data.Foldable as Foldable
import Data.Function
import qualified Data.List as L
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import Data.Text as T hiding (chunksOf, filter, map)
import Filesystem (isFile, listDirectory)
import Filesystem.Path.CurrentOS
import GHC.Conc
import Stat
import System.Console.CmdArgs
import System.Environment (getArgs, withArgs)
import System.Exit (ExitCode (ExitFailure), exitWith)
import System.IO (hPutStrLn, stderr)
import System.Posix.Files hiding (fileBlockSize)
import System.Posix.Types (DeviceID, FileID)
import Text.Printf
import Text.Regex.PCRE
import Unsafe.Coerce
import Prelude hiding (FilePath, sequence)

default (Integer, Text)

version :: String
version = "2.4.4"

copyright :: String
copyright = "2012-2026"

sizesSummary :: String
sizesSummary = "sizes v" ++ version ++ ", (C) John Wiegley " ++ copyright

data SizesOpts = SizesOpts
    { jobs :: Int
    , byCount :: Bool
    , annex :: Bool
    , apparent :: Bool
    , baseTen :: Bool
    , exclude :: [String]
    , minSize :: Int
    , minCount :: Int
    , blockSize :: Int
    , smalls :: Bool
    , dedupeLinks :: Bool
    , oneFileSystem :: Bool
    , -- , dirsOnly  :: Bool
      depth :: Int
    , dirs :: [String]
    }
    deriving (Data, Typeable, Show, Eq)

sizesOpts :: SizesOpts
sizesOpts =
    SizesOpts
        { jobs =
            def
                &= name "j"
                &= typ "INT"
                &= help "Run INT concurrent finds at once (default: 2)"
        , byCount =
            def
                &= name "c"
                &= typ "BOOL"
                &= help "Sort output by count (default: by size)"
        , annex =
            def
                &= name "A"
                &= typ "BOOL"
                &= help "Be mindful of how git-annex stores files"
        , apparent =
            def
                &= typ "BOOL"
                &= help "Print apparent sizes, rather than disk usage"
        , baseTen =
            def
                &= name "H"
                &= typ "BOOL"
                &= help "Print amounts divided by 1000 rather than 1024"
        , exclude =
            def
                &= name "x"
                &= typ "REGEX"
                &= help "Exclude files whose path matches any of the REGEXes (repeatable)"
        , minSize =
            def
                &= name "m"
                &= typ "INT"
                &= help "Smallest size to show, in MB (default: 10)"
        , minCount =
            def
                &= name "M"
                &= typ "INT"
                &= help "Smallest count to show (default: 100)"
        , blockSize =
            def
                &= name "B"
                &= typ "INT"
                &= help "Size of blocks on disk (default: 512)"
        , smalls =
            def
                &= name "s"
                &= typ "BOOL"
                &= help "Also show small (<1M && <100 files) entries"
        , dedupeLinks =
            def
                &= name "L"
                &= typ "BOOL"
                &= help "Deduplicate hard links (count each inode only once)"
        , oneFileSystem =
            def
                &= explicit
                &= name "one-file-system"
                &= name "X"
                &= typ "BOOL"
                &= help "Do not descend into directories on other filesystems"
        , -- , dirsOnly = def &= typ "BOOL"
          --                  &= help "Show directories only"
          depth =
            def
                &= typ "INT"
                &= help "Report entries to a depth of INT (default: 1)"
        , dirs = def &= args &= typ "DIRS..."
        }
        &= summary sizesSummary
        &= program "sizes"
        &= help "Calculate amount of disk used by the given directories"

data EntryInfo = EntryInfo
    { _entryPath :: FilePath
    , _entryCount :: !Int
    , _entryAllocSize :: !Int
    , _entryIsDir :: Bool
    }
    deriving (Show, Eq)

makeLenses ''EntryInfo

-- Track (DeviceID, FileID) pairs to detect hard links
type SeenInodes = Set.Set (DeviceID, FileID)

-- | Stack-safe collection of entries retained for the final report.
newtype ReportEntries = ReportEntries (Seq EntryInfo)

-- | An empty report-entry collection.
emptyReportEntries :: ReportEntries
emptyReportEntries = ReportEntries Seq.empty

-- | Materialize report entries in traversal order.
reportEntriesToList :: ReportEntries -> [EntryInfo]
reportEntriesToList (ReportEntries entries) = Foldable.toList entries

newEntry :: FilePath -> Bool -> EntryInfo
newEntry p = EntryInfo p 0 0

instance Semigroup EntryInfo where
    x <> y =
        let !count' = x ^. entryCount + y ^. entryCount
            !allocSize' = x ^. entryAllocSize + y ^. entryAllocSize
         in x
                & entryCount .~ count'
                & entryAllocSize .~ allocSize'

instance Monoid EntryInfo where
    mempty = newEntry "" False

instance NFData EntryInfo where
    rnf entry =
        rnf (toTextIgnore (entry ^. entryPath)) `seq`
            rnf (entry ^. entryCount) `seq`
                rnf (entry ^. entryAllocSize) `seq`
                    rnf (entry ^. entryIsDir)

-- | Add one child aggregate and its retained descendants to a directory.
combineEntryResults :: Bool -> (EntryInfo, ReportEntries) -> (EntryInfo, ReportEntries) -> (EntryInfo, ReportEntries)
combineEntryResults keepReports (total, ReportEntries reports) (child, ReportEntries childReports) =
    let total' = total <> child
        reports' =
            if keepReports
                then ReportEntries ((reports Seq.|> child) Seq.>< childReports)
                else emptyReportEntries
     in total' `seq` reports' `seq` (total', reports')

sizesMain :: IO ()
sizesMain = do
    mainArgs <- getArgs
    opts <-
        withArgs
            (if L.null mainArgs then [] else mainArgs)
            (cmdArgs sizesOpts)
    GHC.Conc.setNumCapabilities $ case jobs opts of 0 -> 2; x -> x
    runSizes $ case depth opts of 0 -> opts{depth = 1}; _ -> opts

runSizes :: SizesOpts -> IO ()
runSizes opts = do
    let dirsOpt = dirs opts
        directories = if L.null dirsOpt then ["."] else dirsOpt
    reportSizes opts $ Prelude.map (fromText . pack) directories
    stopGlobalPool

reportEntryP :: SizesOpts -> EntryInfo -> Bool
reportEntryP opts entry =
    smalls opts
        || entry ^. entryAllocSize >= minSize'
        || entry ^. entryCount >= minCount'
  where
    minSize' =
        (if minSize opts == 0 then 10 else minSize opts)
            * (if baseTen opts then 1000 else 1024) ^ (2 :: Int)
    minCount' = if minCount opts == 0 then 100 else minCount opts

reportSizes :: SizesOpts -> [FilePath] -> IO ()
reportSizes opts xs = do
    excludes <- compileExcludes (exclude opts)
    entryInfos <- parallel $ Prelude.map (reportSizesForDir excludes) xs
    let infos =
            Prelude.map fst entryInfos
                ++ Prelude.concatMap (reportEntriesToList . snd) entryInfos
        sorted =
            L.sortBy
                ( (compare `on`) $
                    if byCount opts
                        then (^. entryCount)
                        else (^. entryAllocSize)
                )
                infos
    mapM_
        (reportEntry (baseTen opts))
        (Prelude.filter (reportEntryP opts) sorted)
  where
    reportSizesForDir excludes dir =
        -- fsStatus <- getFilesystemStatus (E.encodeUtf8 (toTextIgnore dir))
        let fsBlkSize = statBlockSize -- filesystemBlockSize fsStatus
            opts' =
                if blockSize opts == 0
                    then opts{blockSize = fromIntegral fsBlkSize}
                    else opts
         in fst <$> runStateT (gatherSizes opts' excludes Nothing 0 dir) Set.empty

{- | Compile the exclusion patterns once, up front.  An invalid pattern is a
fatal startup error rather than a silent per-path failure.
-}
compileExcludes :: [String] -> IO [Regex]
compileExcludes = mapM compileOne
  where
    compileOne regex = do
        compiled <- try (makeRegexM regex)
        case compiled of
            Left ex -> do
                hPutStrLn stderr $ "invalid exclusion regex " ++ Prelude.show regex ++ ": " ++ Prelude.show (ex :: SomeException)
                exitWith (ExitFailure 1)
            Right r -> pure r

humanReadable :: Int -> Int -> String
humanReadable x d
    | x < d = printf "%db" x
    | x < d ^ (2 :: Int) = printf "%.0fK" (fromIntegral x / (fromIntegral d :: Double))
    | x < d ^ (3 :: Int) = printf "%.1fM" (fromIntegral x / (fromIntegral d ^ (2 :: Int) :: Double))
    | x < d ^ (4 :: Int) = printf "%.2fG" (fromIntegral x / (fromIntegral d ^ (3 :: Int) :: Double))
    | x < d ^ (5 :: Int) = printf "%.3fT" (fromIntegral x / (fromIntegral d ^ (4 :: Int) :: Double))
    | x < d ^ (6 :: Int) = printf "%.3fP" (fromIntegral x / (fromIntegral d ^ (5 :: Int) :: Double))
    | x < d ^ (7 :: Int) = printf "%.3fX" (fromIntegral x / (fromIntegral d ^ (6 :: Int) :: Double))
    | otherwise = printf "%db" x

reportEntry :: Bool -> EntryInfo -> IO ()
reportEntry bTen entry =
    let path = T.unpack (toTextIgnore (entry ^. entryPath))
     in printf
            (T.unpack "%10s %10d  %s%s\n")
            (humanReadable (entry ^. entryAllocSize) (if bTen then 1000 else 1024))
            (entry ^. entryCount)
            path
            ( T.unpack $
                if entry ^. entryIsDir && L.last path /= '/'
                    then "/"
                    else ""
            )

toTextIgnore :: FilePath -> Text
toTextIgnore = either id id . toText

returnEmpty :: FilePath -> StateT SeenInodes IO (EntryInfo, ReportEntries)
returnEmpty path = return (newEntry path False, emptyReportEntries)

{- | Decide whether an entry should be skipped because it lies on a different
filesystem than the traversal root.  The second argument is the device of
the traversal root, or 'Nothing' at the top-level argument itself (where the
boundary is merely established, so nothing is skipped).
-}
crossesFileSystemBoundary :: Bool -> Maybe DeviceID -> DeviceID -> Bool
crossesFileSystemBoundary False _ _ = False
crossesFileSystemBoundary True Nothing _ = False
crossesFileSystemBoundary True (Just rootDev) dev = dev /= rootDev

-- | Decide whether a path matches any of the compiled exclusion patterns.
excludePath :: [Regex] -> String -> Bool
excludePath patterns path = L.any (`match` path) patterns

gatherSizes :: SizesOpts -> [Regex] -> Maybe DeviceID -> Int -> FilePath -> StateT SeenInodes IO (EntryInfo, ReportEntries)
gatherSizes opts excludes mRootDev curDepth path =
    if excludePath excludes path'
        then returnEmpty path
        else
            ( do
                status <-
                    liftIO
                        ( if curDepth == 0
                            then getFileStatus path'
                            else getSymbolicLinkStatus path'
                        )
                if crossesFileSystemBoundary (oneFileSystem opts) mRootDev (deviceID status)
                    then returnEmpty path
                    else go status
            )
                `catch` ( \e -> do
                            liftIO $ putStrLn $ path' ++ ": " ++ Prelude.show (e :: IOException)
                            returnEmpty path
                        )
  where
    pathT = toTextIgnore path
    path' = T.unpack pathT

    go status
        | isDirectory status =
            foldM
                ( \(y, ys) x -> do
                    child <- gatherSizes opts excludes (Just (deviceID status)) (curDepth + 1) (collapse x)
                    return $! combineEntryResults (curDepth < depth opts) (y, ys) child
                )
                (newEntry path True, emptyReportEntries)
                =<< liftIO (listDirectory path)
        | ( isRegularFile status
                && not (annex opts && ".git/annex/" `isInfixOf` pathT)
          )
            || (annex opts && isSymbolicLink status) = do
            status' <-
                -- If status is for a symbolic link, it must be a Git-annex'd file
                if isSymbolicLink status
                    then do
                        destPath <- liftIO $ readSymbolicLink path'
                        if ".git/annex/" `L.isInfixOf` destPath
                            then do
                                let destFilePath = fromText (T.pack destPath)
                                    destPath' =
                                        if relative destFilePath
                                            then
                                                T.unpack . toTextIgnore $
                                                    parent path </> destFilePath
                                            else destPath
                                    destFilePath' = fromText (T.pack destPath')
                                exists <- liftIO $ isFile destFilePath'
                                if exists
                                    then liftIO $ getFileStatus destPath'
                                    else return status
                            else return status
                    else return status

            -- Check for hard link deduplication
            let dev = deviceID status'
                ino = fileID status'
                numLinks = linkCount status'
                shouldDedupe = dedupeLinks opts && numLinks > 1

            -- If deduplication is enabled and this has multiple hard links, check if we've seen it
            alreadySeen <-
                if shouldDedupe
                    then do
                        seen <- get
                        if Set.member (dev, ino) seen
                            then return True
                            else do
                                put $ Set.insert (dev, ino) seen
                                return False
                    else return False

            let fsize = fileSize status'
                blksize = fileBlockSize (unsafeCoerce status')
                allocSize
                    | alreadySeen = 0 -- Don't count size if we've already seen this inode
                    | apparent opts = fromIntegral fsize
                    | otherwise = fromIntegral blksize * blockSize opts

            return
                ( EntryInfo
                    { _entryPath = path
                    , _entryCount = if alreadySeen then 0 else 1
                    , _entryAllocSize = allocSize
                    , _entryIsDir = False
                    }
                , emptyReportEntries
                )
        | otherwise = returnEmpty path

-- Sizes.hs ends here