Infinities as floating point endpoints

curiousleo · curiousleo · commit ef2ee10d63ff · 2020-06-25T11:50:58.000+02:00
diff --git a/src/System/Random/Internal.hs b/src/System/Random/Internal.hs
@@ -423,7 +423,7 @@ runStateGen g f = runState (f StateGenM) g
 --
 -- >>> import System.Random.Stateful
 -- >>> let pureGen = mkStdGen 137
--- >>> runStateGen_ pureGen  randomM :: Int
+-- >>> runStateGen_ pureGen randomM :: Int
 -- 7879794327570578227
 --
 -- @since 1.2.0
@@ -772,6 +772,9 @@ instance UniformRange Bool where
 instance UniformRange Double where
   uniformRM (l, h) g
     | l == h = return l
+    | isInfinite l && isInfinite h = return (0/0) -- NaN
+    | isInfinite l = return l
+    | isInfinite h = return h
     | otherwise = do
       x <- uniformDouble01M g
       return $ x * l + (1 -x) * h
@@ -806,6 +809,9 @@ uniformDoublePositive01M g = (+ d) <$> uniformDouble01M g
 instance UniformRange Float where
   uniformRM (l, h) g
     | l == h = return l
+    | isInfinite l && isInfinite h = return (0/0) -- NaN
+    | isInfinite l = return l
+    | isInfinite h = return h
     | otherwise = do
       x <- uniformFloat01M g
       return $ x * l + (1 - x) * h
diff --git a/test/Spec.hs b/test/Spec.hs
@@ -126,13 +126,21 @@ integralSpec px =
 
 floatingSpec ::
      forall a.
-     (SC.Serial IO a, Typeable a, Num a, Ord a, Random a, UniformRange a, Show a)
+     (SC.Serial IO a, Typeable a, Num a, Ord a, Random a, UniformRange a, Read a, Show a)
   => Proxy a -> TestTree
 floatingSpec px =
   testGroup ("(" ++ showsType px ")")
   [ SC.testProperty "uniformR" $ seeded $ Range.uniformRangeWithin px
+  , testCase "r = +inf, x = 0" $ positiveInf @?= fst (uniformR (0, positiveInf) (ConstGen 0))
+  , testCase "r = +inf, x = 1" $ positiveInf @?= fst (uniformR (0, positiveInf) (ConstGen 1))
+  , testCase "l = -inf, x = 0" $ negativeInf @?= fst (uniformR (negativeInf, 0) (ConstGen 0))
+  , testCase "l = -inf, x = 1" $ negativeInf @?= fst (uniformR (negativeInf, 0) (ConstGen 1))
   -- TODO: Add more tests
   ]
+  where
+    positiveInf, negativeInf :: a
+    positiveInf = read "Infinity"
+    negativeInf = read "-Infinity"
 
 runSpec :: TestTree
 runSpec = testGroup "runGenState_ and runPrimGenIO_"
@@ -141,3 +149,9 @@ runSpec = testGroup "runGenState_ and runPrimGenIO_"
 -- | Create a StdGen instance from an Int and pass it to the given function.
 seeded :: (StdGen -> a) -> Int -> a
 seeded f = f . mkStdGen
+
+newtype ConstGen = ConstGen Word64
+
+instance RandomGen ConstGen where
+  genWord64 g@(ConstGen c) = (c, g)
+  split g = (g, g)
