Utils.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Math.Float.Utils
  ( FPBits(..)
  , nextAfter, succFP, prevFP
  )
  where

-- base
import Data.Bits
  ( Bits((.&.), shiftL) )
import Data.Word
  ( Word32, Word64 )
import GHC.Float
  ( castFloatToWord32 , castWord32ToFloat
  , castDoubleToWord64, castWord64ToDouble
  )

--------------------------------------------------------------------------------

class ( RealFloat f, Num b, Bits b ) => FPBits f b | f -> b, b -> f where
  toBits   :: f -> b
  fromBits :: b -> f
  -- | Size in bytes.
  sizeOf   :: Int

instance FPBits Float Word32 where
  toBits   = castFloatToWord32
  fromBits = castWord32ToFloat
  sizeOf   = 4

instance FPBits Double Word64 where
  toBits   = castDoubleToWord64
  fromBits = castWord64ToDouble
  sizeOf   = 8


{-# SPECIALISE nextAfter @Float  #-}
{-# SPECIALISE nextAfter @Double #-}
{-# INLINEABLE nextAfter #-}
-- | @nextAfter a b@ computes the next floating-point value after @a@
-- in the direction of @b@.
nextAfter :: forall f b. FPBits f b => f -> f -> f
nextAfter a b
  | isNaN a
  = a
  | isNaN b
  = b
  | a == b
  = b
  | otherwise
  = let !res_bits
          | a == 0
          , let !sgn_mask = 1 `shiftL` ( sizeOf @f * 8 - 1 )
          = ( toBits b .&. sgn_mask ) + 1
          | ( a < b ) == ( a > 0 )
          = toBits a + 1
          | otherwise
          = toBits a - 1
    in fromBits res_bits

{-# SPECIALISE succFP @Float  #-}
{-# SPECIALISE succFP @Double #-}
{-# INLINEABLE succFP #-}
-- | The next floating-point number.
succFP :: forall f b. FPBits f b => f -> f
succFP x = nextAfter x (1/0)


{-# SPECIALISE prevFP @Float  #-}
{-# SPECIALISE prevFP @Double #-}
{-# INLINEABLE prevFP #-}
-- | The previous floating-point number.
prevFP :: forall f b. FPBits f b => f -> f
prevFP x = nextAfter x (-1/0)