Utils.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Distribution.Client.Utils
  ( MergeResult (..)
  , mergeBy
  , duplicates
  , duplicatesBy
  , readMaybe
  , withEnv
  , withEnvOverrides
  , logDirChange
  , withExtraPathEnv
  , determineNumJobs
  , numberOfProcessors
  , withTempFileName
  , makeAbsoluteToCwd
  , makeRelativeToCwd
  , makeRelativeToDir
  , makeRelativeToDirS
  , makeRelativeCanonical
  , filePathToByteString
  , byteStringToFilePath
  , tryCanonicalizePath
  , canonicalizePathNoThrow
  , moreRecentFile
  , existsAndIsMoreRecentThan
  , tryReadAddSourcePackageDesc
  , tryReadGenericPackageDesc
  , relaxEncodingErrors
  , ProgressPhase (..)
  , progressMessage
  , pvpize
  , incVersion
  , getCurrentYear
  , listFilesRecursive
  , listFilesInside
  , safeRead
  , hasElem
  , concatMapM
  , occursOnlyOrBefore
  , giveRTSWarning
  ) where

import Distribution.Client.Compat.Prelude
import Prelude ()

import qualified Control.Exception as Exception
  ( finally
  )
import qualified Control.Exception.Safe as Safe
  ( bracket
  )
import Control.Monad
  ( zipWithM_
  )
import Data.Bits
  ( shiftL
  , shiftR
  , (.|.)
  )
import qualified Data.ByteString.Lazy as BS
import Data.List
  ( elemIndex
  , groupBy
  )
import Distribution.Client.Errors
import Distribution.Compat.Environment
import Distribution.Compat.Time (getModTime)
import Distribution.Simple.Setup (Flag, pattern Flag, pattern NoFlag)
import Distribution.Simple.Utils (dieWithException, findPackageDesc, noticeNoWrap, removeFileForcibly)
import Distribution.Utils.Path
  ( CWD
  , FileOrDir (..)
  , Pkg
  , RelativePath
  , SymbolicPath
  , getSymbolicPath
  , makeSymbolicPath
  , relativeSymbolicPath
  , sameDirectory
  , symbolicPathRelative_maybe
  )
import Distribution.Version

import System.Directory
  ( canonicalizePath
  , doesDirectoryExist
  , doesFileExist
  , listDirectory
  )
import qualified System.Directory as Directory
import System.FilePath
import System.IO
  ( Handle
  , hClose
  , hGetEncoding
  , hSetEncoding
  , openTempFile
  )
import System.IO.Unsafe (unsafePerformIO)

import qualified Data.Set as Set
import Data.Time (utcToLocalTime)
import Data.Time.Calendar (toGregorian)
import Data.Time.Clock.POSIX (getCurrentTime)
import Data.Time.LocalTime (getCurrentTimeZone, localDay)
import Distribution.Simple.PackageDescription (readGenericPackageDescription)
import Distribution.Types.GenericPackageDescription (GenericPackageDescription)
import GHC.Conc.Sync (getNumProcessors)
import GHC.IO.Encoding
  ( TextEncoding (TextEncoding)
  , recover
  )
import GHC.IO.Encoding.Failure
  ( CodingFailureMode (TransliterateCodingFailure)
  , recoverEncode
  )
import qualified System.Directory as Dir
import qualified System.IO.Error as IOError

-- | Generic merging utility. For sorted input lists this is a full outer join.
mergeBy :: forall a b. (a -> b -> Ordering) -> [a] -> [b] -> [MergeResult a b]
mergeBy cmp = merge
  where
    merge :: [a] -> [b] -> [MergeResult a b]
    merge [] ys = [OnlyInRight y | y <- ys]
    merge xs [] = [OnlyInLeft x | x <- xs]
    merge (x : xs) (y : ys) =
      case x `cmp` y of
        GT -> OnlyInRight y : merge (x : xs) ys
        EQ -> InBoth x y : merge xs ys
        LT -> OnlyInLeft x : merge xs (y : ys)

data MergeResult a b = OnlyInLeft a | InBoth a b | OnlyInRight b

duplicates :: Ord a => [a] -> [NonEmpty a]
duplicates = duplicatesBy compare

duplicatesBy :: forall a. (a -> a -> Ordering) -> [a] -> [NonEmpty a]
duplicatesBy cmp = mapMaybe moreThanOne . groupBy eq . sortBy cmp
  where
    eq :: a -> a -> Bool
    eq a b = case cmp a b of
      EQ -> True
      _ -> False
    moreThanOne (x : xs@(_ : _)) = Just (x :| xs)
    moreThanOne _ = Nothing

-- | A variant of 'withTempFile' that only gives us the file name, and while
-- it will clean up the file afterwards, it's lenient if the file is
-- moved\/deleted.
withTempFileName
  :: FilePath
  -> String
  -> (FilePath -> IO a)
  -> IO a
withTempFileName tmpDir template action =
  Safe.bracket
    (openTempFile tmpDir template)
    (\(name, _) -> removeFileForcibly name)
    (\(name, h) -> hClose h >> action name)

-- | Executes the action with an environment variable set to some
-- value.
--
-- Warning: This operation is NOT thread-safe, because current
-- environment is a process-global concept.
withEnv :: String -> String -> IO a -> IO a
withEnv k v m = do
  mb_old <- lookupEnv k
  setEnv k v
  m `Exception.finally` setOrUnsetEnv k mb_old

-- | Executes the action with a list of environment variables and
-- corresponding overrides, where
--
-- * @'Just' v@ means \"set the environment variable's value to @v@\".
-- * 'Nothing' means \"unset the environment variable\".
--
-- Warning: This operation is NOT thread-safe, because current
-- environment is a process-global concept.
withEnvOverrides :: [(String, Maybe FilePath)] -> IO a -> IO a
withEnvOverrides overrides m = do
  mb_olds <- traverse lookupEnv envVars
  traverse_ (uncurry setOrUnsetEnv) overrides
  m `Exception.finally` zipWithM_ setOrUnsetEnv envVars mb_olds
  where
    envVars :: [String]
    envVars = map fst overrides

setOrUnsetEnv :: String -> Maybe String -> IO ()
setOrUnsetEnv var Nothing = unsetEnv var
setOrUnsetEnv var (Just val) = setEnv var val

-- | Executes the action, increasing the PATH environment
-- in some way
--
-- Warning: This operation is NOT thread-safe, because the
-- environment variables are a process-global concept.
withExtraPathEnv :: [FilePath] -> IO a -> IO a
withExtraPathEnv paths m = do
  oldPathSplit <- getSearchPath
  let newPath :: String
      newPath = mungePath $ intercalate [searchPathSeparator] (paths ++ oldPathSplit)
      oldPath :: String
      oldPath = mungePath $ intercalate [searchPathSeparator] oldPathSplit
      -- TODO: This is a horrible hack to work around the fact that
      -- setEnv can't take empty values as an argument
      mungePath p
        | p == "" = "/dev/null"
        | otherwise = p
  setEnv "PATH" newPath
  m `Exception.finally` setEnv "PATH" oldPath

-- | Log directory change in 'make' compatible syntax
logDirChange :: (String -> IO ()) -> Maybe FilePath -> IO a -> IO a
logDirChange _ Nothing m = m
logDirChange l (Just d) m = do
  l $ "cabal: Entering directory '" ++ d ++ "'\n"
  m
    `Exception.finally` l ("cabal: Leaving directory '" ++ d ++ "'\n")

-- The number of processors is not going to change during the duration of the
-- program, so unsafePerformIO is safe here.
numberOfProcessors :: Int
numberOfProcessors = unsafePerformIO getNumProcessors
{-# NOINLINE numberOfProcessors #-}

-- | Determine the number of jobs to use given the value of the '-j' flag.
determineNumJobs :: Flag (Maybe Int) -> Int
determineNumJobs numJobsFlag =
  case numJobsFlag of
    NoFlag -> 1
    Flag Nothing -> numberOfProcessors
    Flag (Just n) -> n

-- | Given a relative path, make it absolute relative to the current
-- directory. Absolute paths are returned unmodified.
makeAbsoluteToCwd :: FilePath -> IO FilePath
makeAbsoluteToCwd path
  | isAbsolute path = return path
  | otherwise = do
      cwd <- Directory.getCurrentDirectory
      return $! cwd </> path

-- | Given a path (relative or absolute), make it relative to the current
-- directory, including using @../..@ if necessary.
makeRelativeToCwd :: FilePath -> IO FilePath
makeRelativeToCwd path =
  makeRelativeCanonical <$> canonicalizePath path <*> Directory.getCurrentDirectory

-- | Given a path (relative or absolute), make it relative to the given
-- directory, including using @../..@ if necessary.
makeRelativeToDir :: FilePath -> FilePath -> IO FilePath
makeRelativeToDir path dir =
  makeRelativeCanonical <$> canonicalizePath path <*> canonicalizePath dir

-- | makeRelativeToDir for SymbolicPath
makeRelativeToDirS :: Maybe (SymbolicPath CWD (Dir dir)) -> SymbolicPath CWD to -> IO (SymbolicPath dir to)
makeRelativeToDirS Nothing s = makeRelativeToDirS (Just sameDirectory) s
makeRelativeToDirS (Just root) p =
  case symbolicPathRelative_maybe p of
    -- TODO: Use AbsolutePath
    Nothing -> return $ makeSymbolicPath (getSymbolicPath p)
    Just rel_path ->
      makeSymbolicPath <$> makeRelativeToDir (getSymbolicPath root) (getSymbolicPath rel_path)

-- | Given a canonical absolute path and canonical absolute dir, make the path
-- relative to the directory, including using @../..@ if necessary. Returns
-- the original absolute path if it is not on the same drive as the given dir.
makeRelativeCanonical :: FilePath -> FilePath -> FilePath
makeRelativeCanonical path dir
  | takeDrive path /= takeDrive dir = path
  | otherwise = go (splitPath path) (splitPath dir)
  where
    go (p : ps) (d : ds) | p' == d' = go ps ds
      where
        (p', d') = (dropTrailingPathSeparator p, dropTrailingPathSeparator d)
    go [] [] = "./"
    go ps ds = joinPath (replicate (length ds) ".." ++ ps)

-- | Convert a 'FilePath' to a lazy 'ByteString'. Each 'Char' is
-- encoded as a little-endian 'Word32'.
filePathToByteString :: FilePath -> BS.ByteString
filePathToByteString p =
  BS.pack $ foldr conv [] codepts
  where
    codepts :: [Word32]
    codepts = map (fromIntegral . ord) p

    conv :: Word32 -> [Word8] -> [Word8]
    conv w32 rest = b0 : b1 : b2 : b3 : rest
      where
        b0 = fromIntegral w32
        b1 = fromIntegral $ w32 `shiftR` 8
        b2 = fromIntegral $ w32 `shiftR` 16
        b3 = fromIntegral $ w32 `shiftR` 24

-- | Reverse operation to 'filePathToByteString'.
byteStringToFilePath :: BS.ByteString -> FilePath
byteStringToFilePath bs
  | bslen `mod` 4 /= 0 = unexpected
  | otherwise = go 0
  where
    unexpected = "Distribution.Client.Utils.byteStringToFilePath: unexpected"
    bslen = BS.length bs

    go i
      | i == bslen = []
      | otherwise = (chr . fromIntegral $ w32) : go (i + 4)
      where
        w32 :: Word32
        w32 = b0 .|. (b1 `shiftL` 8) .|. (b2 `shiftL` 16) .|. (b3 `shiftL` 24)
        b0 = fromIntegral $ BS.index bs i
        b1 = fromIntegral $ BS.index bs (i + 1)
        b2 = fromIntegral $ BS.index bs (i + 2)
        b3 = fromIntegral $ BS.index bs (i + 3)

-- | Workaround for the inconsistent behaviour of 'canonicalizePath'. Always
-- throws an error if the path refers to a non-existent file.
tryCanonicalizePath :: FilePath -> IO FilePath
tryCanonicalizePath path = do
  ret <- canonicalizePath path
  exists <- liftM2 (||) (doesFileExist ret) (Dir.doesDirectoryExist ret)
  unless exists $
    IOError.ioError $
      IOError.mkIOError
        IOError.doesNotExistErrorType
        "canonicalizePath"
        Nothing
        (Just ret)
  return ret

-- | A non-throwing wrapper for 'canonicalizePath'. If 'canonicalizePath' throws
-- an exception, returns the path argument unmodified.
canonicalizePathNoThrow :: FilePath -> IO FilePath
canonicalizePathNoThrow path = do
  canonicalizePath path `catchIO` (\_ -> return path)

--------------------
-- Modification time

-- | Like Distribution.Simple.Utils.moreRecentFile, but uses getModTime instead
-- of getModificationTime for higher precision. We can't merge the two because
-- Distribution.Client.Time uses MIN_VERSION macros.
moreRecentFile :: FilePath -> FilePath -> IO Bool
moreRecentFile a b = do
  exists <- doesFileExist b
  if not exists
    then return True
    else do
      tb <- getModTime b
      ta <- getModTime a
      return (ta > tb)

-- | Like 'moreRecentFile', but also checks that the first file exists.
existsAndIsMoreRecentThan :: FilePath -> FilePath -> IO Bool
existsAndIsMoreRecentThan a b = do
  exists <- doesFileExist a
  if not exists
    then return False
    else a `moreRecentFile` b

-- | Sets the handler for encoding errors to one that transliterates invalid
-- characters into one present in the encoding (i.e., \'?\').
-- This is opposed to the default behavior, which is to throw an exception on
-- error. This function will ignore file handles that have a Unicode encoding
-- set. It's a no-op for versions of `base` less than 4.4.
relaxEncodingErrors :: Handle -> IO ()
relaxEncodingErrors handle = do
  maybeEncoding <- hGetEncoding handle
  case maybeEncoding of
    Just (TextEncoding name decoder encoder)
      | not ("UTF" `isPrefixOf` name) ->
          let relax x = x{recover = recoverEncode TransliterateCodingFailure}
           in hSetEncoding handle (TextEncoding name decoder (fmap relax encoder))
    _ ->
      return ()

-- | Like 'tryFindPackageDesc', but with error specific to add-source deps.
tryReadAddSourcePackageDesc
  :: Verbosity
  -> FilePath
  -> String
  -> IO GenericPackageDescription
tryReadAddSourcePackageDesc verbosity depPath err = do
  let pkgDir = makeSymbolicPath depPath
  pkgDescPath <-
    try_find_package_desc verbosity pkgDir $
      err
        ++ "\n"
        ++ "Failed to read cabal file of add-source dependency: "
        ++ depPath
  readGenericPackageDescription verbosity (Just pkgDir) (relativeSymbolicPath pkgDescPath)

-- | Try to read a @.cabal@ file, in directory @depPath@. Fails if one cannot be
--  found, with @err@ prefixing the error message. This function simply allows
--  us to give a more descriptive error than that provided by @findPackageDesc@.
tryReadGenericPackageDesc
  :: Verbosity
  -> SymbolicPath CWD (Dir Pkg)
  -> String
  -> IO GenericPackageDescription
tryReadGenericPackageDesc verbosity pkgDir err = do
  pkgDescPath <- try_find_package_desc verbosity pkgDir err
  readGenericPackageDescription verbosity (Just pkgDir) (relativeSymbolicPath pkgDescPath)

-- | Internal helper function for 'tryReadAddSourcePackageDesc' and 'tryReadGenericPackageDesc'.
try_find_package_desc
  :: Verbosity
  -> SymbolicPath CWD (Dir Pkg)
  -> String
  -> IO (RelativePath Pkg File)
try_find_package_desc verbosity pkgDir err = do
  errOrCabalFile <- findPackageDesc (Just pkgDir)
  case errOrCabalFile of
    Right file -> return file
    Left _ -> dieWithException verbosity $ TryFindPackageDescErr err

-- | Phase of building a dependency. Represents current status of package
-- dependency processing. See #4040 for details.
data ProgressPhase
  = ProgressDownloading
  | ProgressDownloaded
  | ProgressStarting
  | ProgressBuilding
  | ProgressHaddock
  | ProgressInstalling
  | ProgressCompleted

progressMessage :: Verbosity -> ProgressPhase -> String -> IO ()
progressMessage verbosity phase subject = do
  noticeNoWrap verbosity $ phaseStr ++ subject ++ "\n"
  where
    phaseStr = case phase of
      ProgressDownloading ->
        "Downloading  "
      ProgressDownloaded ->
        "Downloaded   "
      ProgressStarting ->
        "Starting     "
      ProgressBuilding ->
        "Building     "
      ProgressHaddock ->
        "Haddock      "
      ProgressInstalling ->
        "Installing   "
      ProgressCompleted ->
        "Completed    "

-- | Given a version, return an API-compatible (according to PVP) version range.
--
-- If the boolean argument denotes whether to use a desugared
-- representation (if 'True') or the new-style @^>=@-form (if
-- 'False').
--
-- Example: @pvpize True (mkVersion [0,4,1])@ produces the version range @>= 0.4 && < 0.5@ (which is the
-- same as @0.4.*@).
pvpize :: Bool -> Version -> VersionRange
pvpize False v = majorBoundVersion v
pvpize True v =
  orLaterVersion v'
    `intersectVersionRanges` earlierVersion (incVersion 1 v')
  where
    v' = alterVersion (take 2) v

-- | Increment the nth version component (counting from 0).
incVersion :: Int -> Version -> Version
incVersion n = alterVersion (incVersion' n)
  where
    incVersion' 0 [] = [1]
    incVersion' 0 (v : _) = [v + 1]
    incVersion' m [] = replicate m 0 ++ [1]
    incVersion' m (v : vs) = v : incVersion' (m - 1) vs

-- | Returns the current calendar year.
getCurrentYear :: IO Integer
getCurrentYear = do
  u <- getCurrentTime
  z <- getCurrentTimeZone
  let l = utcToLocalTime z u
      (y, _, _) = toGregorian $ localDay l
  return y

-- | From System.Directory.Extra
--   https://hackage.haskell.org/package/extra-1.7.9
listFilesInside :: (FilePath -> IO Bool) -> FilePath -> IO [FilePath]
listFilesInside test dir = ifNotM (test $ dropTrailingPathSeparator dir) (pure []) $ do
  (dirs, files) <- partitionM doesDirectoryExist =<< listContents dir
  rest <- concatMapM (listFilesInside test) dirs
  pure $ files ++ rest

-- | From System.Directory.Extra
--   https://hackage.haskell.org/package/extra-1.7.9
listFilesRecursive :: FilePath -> IO [FilePath]
listFilesRecursive = listFilesInside (const $ pure True)

listContents :: FilePath -> IO [FilePath]
listContents dir =
  map (dir </>) . sort <$> listDirectory dir

-- | From Control.Monad.Extra
--   https://hackage.haskell.org/package/extra-1.7.9
ifM :: Monad m => m Bool -> m a -> m a -> m a
ifM b t f = do b' <- b; if b' then t else f

-- | 'ifM' with swapped branches:
--   @ifNotM b t f = ifM (not <$> b) t f@
ifNotM :: Monad m => m Bool -> m a -> m a -> m a
ifNotM = flip . ifM

-- | From Control.Monad.Extra
--   https://hackage.haskell.org/package/extra-1.7.9
concatMapM :: Monad m => (a -> m [b]) -> [a] -> m [b]
{-# INLINE concatMapM #-}
concatMapM op = foldr f (pure [])
  where
    f x xs = do x' <- op x; if null x' then xs else do { xs' <- xs; pure $ x' ++ xs' }

-- | From Control.Monad.Extra
--   https://hackage.haskell.org/package/extra-1.7.9
partitionM :: Monad m => (a -> m Bool) -> [a] -> m ([a], [a])
partitionM _ [] = pure ([], [])
partitionM f (x : xs) = do
  res <- f x
  (as, bs) <- partitionM f xs
  pure ([x | res] ++ as, [x | not res] ++ bs)

safeRead :: Read a => String -> Maybe a
safeRead s
  | [(x, "")] <- reads s = Just x
  | otherwise = Nothing

-- | @hasElem xs x = elem x xs@ except that @xs@ is turned into a 'Set' first.
--   Use underapplied to speed up subsequent lookups, e.g. @filter (hasElem xs) ys@.
--   Only amortized when used several times!
--
--   Time complexity \(O((n+m) \log(n))\) for \(m\) lookups in a list of length \(n\).
--   (Compare this to 'elem''s \(O(nm)\).)
--
--   This is [Agda.Utils.List.hasElem](https://hackage.haskell.org/package/Agda-2.6.2.2/docs/Agda-Utils-List.html#v:hasElem).
hasElem :: Ord a => [a] -> a -> Bool
hasElem xs = (`Set.member` Set.fromList xs)

-- True if x occurs before y
occursOnlyOrBefore :: Eq a => [a] -> a -> a -> Bool
occursOnlyOrBefore xs x y = case (elemIndex x xs, elemIndex y xs) of
  (Just i, Just j) -> i < j
  (Just _, _) -> True
  _ -> False

giveRTSWarning :: String -> String
giveRTSWarning "run" =
  "Your RTS options are applied to cabal, not the "
    ++ "executable. Use '--' to separate cabal options from your "
    ++ "executable options. For example, use 'cabal run -- +RTS -N "
    ++ "to pass the '-N' RTS option to your executable."
giveRTSWarning "test" =
  "Some RTS options were found standalone, "
    ++ "which affect cabal and not the binary. "
    ++ "Please note that +RTS inside the --test-options argument "
    ++ "suffices if your goal is to affect the tested binary. "
    ++ "For example, use \"cabal test --test-options='+RTS -N'\" "
    ++ "to pass the '-N' RTS option to your binary."
giveRTSWarning "bench" =
  "Some RTS options were found standalone, "
    ++ "which affect cabal and not the binary. Please note "
    ++ "that +RTS inside the --benchmark-options argument "
    ++ "suffices if your goal is to affect the benchmarked "
    ++ "binary. For example, use \"cabal test --benchmark-options="
    ++ "'+RTS -N'\" to pass the '-N' RTS option to your binary."
giveRTSWarning _ =
  "Your RTS options are applied to cabal, not the "
    ++ "binary."