-- Copyright 2018 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

-- | Get info about CPUs.
module Util.Cpu where
import qualified Data.Set as Set
import qualified Data.Text as Text
import           Data.Text (Text)
import qualified Data.Text.IO as Text.IO

import qualified System.Environment as Environment
import qualified System.Info
import qualified System.Process as Process

import qualified Text.Read as Read


-- | Get number of physical cores.  This can be overidden with a CPUS
-- environment variable.  This is useful if you are running on a VM in CI
-- and don't agree with how many cores it claims to have.
physicalCores :: IO Int
physicalCores :: IO Int
physicalCores = forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO Int
getPhysicalCores forall (m :: * -> *) a. Monad m => a -> m a
return forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO (Maybe Int)
envCores

getPhysicalCores :: IO Int
getPhysicalCores :: IO Int
getPhysicalCores = case String
System.Info.os of
    String
"darwin" -> forall a. Read a => String -> a
read forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
        String -> [String] -> String -> IO String
Process.readProcess String
"/usr/sbin/sysctl" [String
"-n", String
"hw.physicalcpu"] String
""
    String
"linux" -> Text -> Int
linuxPhysicalCores forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO Text
Text.IO.readFile String
"/proc/cpuinfo"
    String
_ -> forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"unknown platform: " forall a. [a] -> [a] -> [a]
++ String
System.Info.os

envCores :: IO (Maybe Int)
envCores :: IO (Maybe Int)
envCores = (forall a. Read a => String -> Maybe a
Read.readMaybe =<<) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Maybe String)
Environment.lookupEnv String
"CPUS"

-- | Parse /proc/cpuinfo for physical cpu count.
linuxPhysicalCores :: Text -> Int
linuxPhysicalCores :: Text -> Int
linuxPhysicalCores =
    forall (t :: * -> *) a. Foldable t => t a -> Int
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> [a]
unique forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Bool
null) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map Text -> [Text]
cpu forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => Text -> Text -> [Text]
Text.splitOn Text
"\n\n"
    where
    -- unique pairs of (physical id, core id)
    cpu :: Text -> [Text]
cpu = forall a. (a -> Bool) -> [a] -> [a]
filter (\Text
s -> forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Text -> Text -> Bool
`Text.isPrefixOf` Text
s) [Text
"physical id", Text
"core id"])
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text]
Text.lines

unique :: Ord a => [a] -> [a]
unique :: forall a. Ord a => [a] -> [a]
unique = forall a. Set a -> [a]
Set.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => [a] -> Set a
Set.fromList