Add 'instance UniformRange Natural'

System/Random/Internal.hs

@@ -75,11 +75,12 @@ import Foreign.Ptr (plusPtr)
 import Foreign.Storable (peekByteOff, pokeByteOff)
 import GHC.Exts
 import GHC.ForeignPtr
+import GHC.IO (IO(..))
+import GHC.Word
+import Numeric.Natural (Natural)
 import System.IO.Unsafe (unsafePerformIO)
 import qualified System.Random.SplitMix as SM
 import qualified System.Random.SplitMix32 as SM32
-import GHC.Word
-import GHC.IO (IO(..))
 
 -- | 'RandomGen' is an interface to pure pseudo-random number generators.
 --
@@ -457,7 +458,10 @@ class UniformRange a where
   uniformRM :: MonadRandom g s m => (a, a) -> g s -> m a
 
 instance UniformRange Integer where
-  uniformRM = uniformIntegerM
+  uniformRM = uniformIntegralM
+
+instance UniformRange Natural where
+  uniformRM = uniformIntegralM
 
 instance Uniform Int8 where
   uniformM = fmap (fromIntegral :: Word8 -> Int8) . uniformWord8
@@ -751,75 +755,75 @@ randomIvalInteger (l,h) rng
                         (x,g') = next g
                         v' = (v * b + (fromIntegral x - fromIntegral genlo))
 
--- | Generate an 'Integer' in the range @[l, h]@ if @l <= h@ and @[h, l]@
+-- | Generate an integral in the range @[l, h]@ if @l <= h@ and @[h, l]@
 -- otherwise.
-uniformIntegerM :: (MonadRandom g s m) => (Integer, Integer) -> g s -> m Integer
-uniformIntegerM (l, h) gen = case l `compare` h of
+uniformIntegralM :: (Bits a, Integral a, MonadRandom g s m) => (a, a) -> g s -> m a
+uniformIntegralM (l, h) gen = case l `compare` h of
   LT -> do
     let limit = h - l
     let limitAsWord64 :: Word64 = fromIntegral limit
     bounded <-
-      if (toInteger limitAsWord64) == limit
+      if (fromIntegral limitAsWord64) == limit
         -- Optimisation: if 'limit' fits into 'Word64', generate a bounded
         -- 'Word64' and then convert to 'Integer'
-        then toInteger <$> unsignedBitmaskWithRejectionM uniformWord64 limitAsWord64 gen
-        else boundedExclusiveIntegerM (limit + 1) gen
+        then fromIntegral <$> unsignedBitmaskWithRejectionM uniformWord64 limitAsWord64 gen
+        else boundedExclusiveIntegralM (limit + 1) gen
     return $ l + bounded
-  GT -> uniformIntegerM (h, l) gen
+  GT -> uniformIntegralM (h, l) gen
   EQ -> pure l
-{-# INLINE uniformIntegerM #-}
+{-# INLINEABLE uniformIntegralM #-}
 
--- | Generate an 'Integer' in the range @[0, s)@ using a variant of Lemire's
+-- | Generate an integral in the range @[0, s)@ using a variant of Lemire's
 -- multiplication method.
 --
 -- Daniel Lemire. 2019. Fast Random Integer Generation in an Interval. In ACM
 -- Transactions on Modeling and Computer Simulation
 -- https://doi.org/10.1145/3230636
 --
 -- PRECONDITION (unchecked): s > 0
-boundedExclusiveIntegerM :: (MonadRandom g s m) => Integer -> g s -> m Integer
-boundedExclusiveIntegerM s gen = go
+boundedExclusiveIntegralM :: (Bits a, Integral a, Ord a, MonadRandom g s m) => a -> g s -> m a
+boundedExclusiveIntegralM (s :: a) gen = go
   where
-    n = integerWordSize s
+    n = integralWordSize s
     -- We renamed 'L' from the paper to 'k' here because 'L' is not a valid
     -- variable name in Haskell and 'l' is already used in the algorithm.
     k = WORD_SIZE_IN_BITS * n
-    twoToK = (1::Integer) `shiftL` k
+    twoToK = (1::a) `shiftL` k
     modTwoToKMask = twoToK - 1
 
     t = (twoToK - s) `mod` s
     go = do
-      x <- uniformIntegerWords n gen
+      x <- uniformIntegralWords n gen
       let m = x * s
       -- m .&. modTwoToKMask == m `mod` twoToK
       let l = m .&. modTwoToKMask
       if l < t
         then go
         -- m `shiftR` k == m `quot` twoToK
         else return $ m `shiftR` k
-{-# INLINE boundedExclusiveIntegerM #-}
+{-# INLINE boundedExclusiveIntegralM #-}
 
--- | @integerWordSize i@ returns that least @w@ such that
+-- | @integralWordSize i@ returns that least @w@ such that
 -- @i <= WORD_SIZE_IN_BITS^w@.
-integerWordSize :: Integer -> Int
-integerWordSize = go 0
+integralWordSize :: (Bits a, Num a) => a -> Int
+integralWordSize = go 0
   where
     go !acc i
       | i == 0 = acc
       | otherwise = go (acc + 1) (i `shiftR` WORD_SIZE_IN_BITS)
-{-# INLINE integerWordSize #-}
+{-# INLINE integralWordSize #-}
 
--- | @uniformIntegerWords n@ is a uniformly pseudo-random 'Integer' in the range
+-- | @uniformIntegralWords n@ is a uniformly pseudo-random integral in the range
 -- @[0, WORD_SIZE_IN_BITS^n)@.
-uniformIntegerWords :: (MonadRandom g s m) => Int -> g s -> m Integer
-uniformIntegerWords n gen = go 0 n
+uniformIntegralWords :: (Bits a, Integral a, MonadRandom g s m) => Int -> g s -> m a
+uniformIntegralWords n gen = go 0 n
   where
     go !acc i
       | i == 0 = return acc
       | otherwise = do
         (w :: Word) <- uniformM gen
         go ((acc `shiftL` WORD_SIZE_IN_BITS) .|. (fromIntegral w)) (i - 1)
-{-# INLINE uniformIntegerWords #-}
+{-# INLINE uniformIntegralWords #-}
 
 -- | Uniformly generate an 'Integral' in an inclusive-inclusive range.
 --

bench/Main.hs

@@ -13,6 +13,7 @@ import Data.Typeable
 import Data.Word
 import Foreign.C.Types
 import Gauge.Main
+import Numeric.Natural (Natural)
 import System.Random.SplitMix as SM
 
 import System.Random
@@ -176,6 +177,9 @@ main = do
           , let !i = (10 :: Integer) ^ (100 :: Integer)
                 !range = (-i - 1, i + 1)
             in pureUniformRBench @Integer range sz
+          , let !n = (10 :: Natural) ^ (100 :: Natural)
+                !range = (1, n - 1)
+            in pureUniformRBench @Natural range sz
           ]
         ]
       ]

test/Spec.hs

@@ -9,15 +9,16 @@
 module Main (main) where
 
 import Data.Coerce
-import Data.Word
 import Data.Int
+import Data.Typeable
+import Data.Word
+import Foreign.C.Types
+import Numeric.Natural (Natural)
 import System.Random
+import Test.SmallCheck.Series as SC
 import Test.Tasty
 import Test.Tasty.HUnit
 import Test.Tasty.SmallCheck as SC
-import Test.SmallCheck.Series as SC
-import Data.Typeable
-import Foreign.C.Types
 
 #include "HsBaseConfig.h"
 
@@ -66,6 +67,7 @@ main =
     , integralSpec @CIntMax
     , integralSpec @CUIntMax
     , integralSpec @Integer
+    , integralSpec @Natural
     -- , bitmaskSpec @Word8
     -- , bitmaskSpec @Word16
     -- , bitmaskSpec @Word32
@@ -103,7 +105,7 @@ showsType = showsTypeRep (typeRep (Proxy :: Proxy t))
 
 rangeSpec ::
      forall a.
-     (SC.Serial IO a, Typeable a, Ord a, Random a, UniformRange a, Show a)
+     (SC.Serial IO a, Typeable a, Ord a, UniformRange a, Show a)
   => TestTree
 rangeSpec =
   testGroup ("Range (" ++ showsType @a ")")
@@ -112,7 +114,7 @@ rangeSpec =
 
 integralSpec ::
      forall a.
-     (SC.Serial IO a, Typeable a, Ord a, Random a, UniformRange a, Show a)
+     (SC.Serial IO a, Typeable a, Ord a, UniformRange a, Show a)
   => TestTree
 integralSpec  =
   testGroup ("(" ++ showsType @a ")")

test/Spec/Range.hs

@@ -8,20 +8,20 @@ module Spec.Range
 
 import System.Random.Monad
 
-symmetric :: (RandomGen g, Random a, Eq a) => g -> (a, a) -> Bool
-symmetric g (l, r) = fst (randomR (l, r) g) == fst (randomR (r, l) g)
+symmetric :: (RandomGen g, UniformRange a, Eq a) => g -> (a, a) -> Bool
+symmetric g (l, r) = fst (uniformR (l, r) g) == fst (uniformR (r, l) g)
 
-bounded :: (RandomGen g, Random a, Ord a) => g -> (a, a) -> Bool
+bounded :: (RandomGen g, UniformRange a, Ord a) => g -> (a, a) -> Bool
 bounded g (l, r) = bottom <= result && result <= top
   where
     bottom = min l r
     top = max l r
-    result = fst (randomR (l, r) g)
+    result = fst (uniformR (l, r) g)
 
-singleton :: (RandomGen g, Random a, Eq a) => g -> a -> Bool
+singleton :: (RandomGen g, UniformRange a, Eq a) => g -> a -> Bool
 singleton g x = result == x
   where
-    result = fst (randomR (x, x) g)
+    result = fst (uniformR (x, x) g)
 
 uniformRangeWithin :: (RandomGen g, UniformRange a, Ord a) => g -> (a, a) -> Bool
 uniformRangeWithin gen (l, r) =