hasktorch-CudaHashMap-0.1.0.0: Setup.hs

{-# LANGUAGE CPP #-}

import Control.Monad (filterM, unless, when)
import Distribution.Simple
import Distribution.Simple.Setup (ConfigFlags (..))
import System.Directory
    ( doesDirectoryExist
    , doesFileExist
    , makeAbsolute
    , removePathForcibly
    )
import System.Process (callProcess)

#if MIN_VERSION_Cabal(3,14,0)
import Distribution.Utils.Path (makeSymbolicPath)
#else
makeSymbolicPath :: a -> a
makeSymbolicPath = id
#endif

main :: IO ()
main = defaultMainWithHooks simpleUserHooks
    { confHook = \(gpd, hbi) flags -> do
        -- Use cbits/build if it was created manually;
        -- otherwise use dist-newstyle/cbits-build, which is removed automatically by 'cabal clean'.
        manualBuildExists <- doesFileExist "cbits/build/CMakeCache.txt"
        let buildDir = if manualBuildExists
                         then "cbits/build"
                         else "dist-newstyle/cbits-build"

        cmakeCacheExists <- doesFileExist (buildDir ++ "/CMakeCache.txt")
        unless cmakeCacheExists $
            callProcess "cmake" ["-S", "cbits", "-B", buildDir]
        callProcess "cmake" ["--build", buildDir, "-j4"]

        absBuildDir <- makeAbsolute buildDir
        cudaDirs <- filterM doesDirectoryExist
            [ "/usr/local/cuda/lib64"
            , "/usr/lib/x86_64-linux-gnu"
            ]
        let updatedFlags = flags
                { configExtraLibDirs =
                    map makeSymbolicPath (absBuildDir : cudaDirs)
                    ++ configExtraLibDirs flags
                }
        confHook simpleUserHooks (gpd, hbi) updatedFlags

    , cleanHook = \pd z uh cf -> do
        exists <- doesDirectoryExist "cbits/build"
        when exists $ removePathForcibly "cbits/build"
        cleanHook simpleUserHooks pd z uh cf
    }