content: {-# LANGUAGE RecordWildCards #-}
  
  {-| This module is for detecting which SIMD instruction sets your CPU supports. In particular, it can detect SSE4.2, AVX and AVX2. -} 
  module SDR.CPUID (
      -- * Raw CPUID
      cpuid,
      cpuidExtended,
  
      -- * High level CPU capabilities
      CPUInfo(..),
      getCPUInfo,
  
      -- * Features
      hasSSE42,
      hasAVX,
      hasAVX2,
  
      -- * Convenience functions
      featureSelect
      ) where
  
  import Data.Word
  import Data.Bits
  import Foreign.Ptr
  import Foreign.Marshal.Alloc
  import Foreign.Storable
  import Data.List
  import Data.Maybe
  import Control.Applicative
  
  foreign import ccall unsafe "cpuid"
      cpuid_c :: Word32 -> Ptr Word32 -> Ptr Word32 -> Ptr Word32 -> Ptr Word32 -> IO ()
  
  foreign import ccall unsafe "cpuid_extended"
      cpuidExtended_c :: Word32 -> Word32 -> Ptr Word32 -> Ptr Word32 -> Ptr Word32 -> Ptr Word32 -> IO ()
  
  -- | Execute the CPUID instruction
  cpuid :: Word32                              -- ^ Operation (EAX)
        -> IO (Word32, Word32, Word32, Word32) -- ^ Result (EAX, EBX, ECX, EDX)
  cpuid x = 
      alloca $ \p1 -> 
      alloca $ \p2 -> 
      alloca $ \p3 -> 
      alloca $ \p4 -> do
          cpuid_c x p1 p2 p3 p4
          (,,,) <$> peek p1 <*> peek p2 <*> peek p3 <*> peek p4
  
  -- | Execute the CPUID instruction setting ECX as well
  cpuidExtended :: Word32                              -- ^ Operation (EAX)
                -> Word32                              -- ^ ECX
                -> IO (Word32, Word32, Word32, Word32) -- ^ Result (EAX, EBX, ECX, EDX)
  cpuidExtended x y = 
      alloca $ \p1 -> 
      alloca $ \p2 -> 
      alloca $ \p3 -> 
      alloca $ \p4 -> do
          cpuidExtended_c x y p1 p2 p3 p4
          (,,,) <$> peek p1 <*> peek p2 <*> peek p3 <*> peek p4
  
  -- | Information about the features supported by your CPU
  data CPUInfo = CPUInfo {
      features         :: Word32,
      extendedFeatures :: Maybe Word32
  }
  
  -- | Get a `CPUInfo`
  getCPUInfo :: IO CPUInfo
  getCPUInfo = do
      (x, _, _, _) <- cpuid 0
      (_, _, f, _) <- cpuid 1
      if x < 7 then
          return $ CPUInfo f Nothing
      else do
          (_, e, _, _) <- cpuidExtended 7 0
          return $ CPUInfo f (Just e)
  
  -- | Feature bit for SSE4.2
  sse42 = 20
  
  -- | Feature bit for AVX
  avx   = 28
  
  -- | Extended feature bit for AVX2
  avx2  = 5
  
  -- | Check if the CPU supports SSE4.2
  hasSSE42 :: CPUInfo -> Bool
  hasSSE42 CPUInfo{..} = testBit features sse42
  
  -- | Check if the CPU supports AVX
  hasAVX   :: CPUInfo -> Bool
  hasAVX CPUInfo{..} = testBit features avx 
  
  -- | Check if the CPU supports AVX2
  hasAVX2  :: CPUInfo -> Bool
  hasAVX2 (CPUInfo _ Nothing)  = False
  hasAVX2 (CPUInfo _ (Just f)) = testBit f avx2
  
  -- | Convenience function for selecting a function based on the features that the CPU supports
  featureSelect :: CPUInfo                -- ^ The CPU features
                -> a                      -- ^ Default implementation
                -> [(CPUInfo -> Bool, a)] -- ^ List of (feature, implementation) pairs
                -> a                      -- ^ The selected implementation
  featureSelect info def list = maybe def snd $ find (($ info) . fst) list
  
  