improve floating point random generation and include tree for stdgen

1 files changed, 26 insertions, 20 deletions

diff --git a/src/PPL/Internal.hs b/src/PPL/Internal.hs
index d09d8ff..3c54078 100644
--- a/src/PPL/Internal.hs
+++ b/src/PPL/Internal.hs
@@ -17,8 +17,10 @@ module PPL.Internal
     sample,
     memoize,
     samples,
-    newTree,
     HashMap,
+    random,
+    randoms,
+    newTree,
     Tree (..),
   )
 where
@@ -27,17 +29,20 @@ import Control.Monad
 import Control.Monad.Trans.Class
 import Control.Monad.Trans.Writer
 import Data.Bifunctor
+import Data.Bits (countTrailingZeros, shiftL, shiftR)
 import Data.IORef
 import Data.Map qualified as Q
 import Data.Monoid
 import Data.Vector.Hashtables qualified as H
 import Data.Vector.Mutable qualified as M
 import Data.Vector.Unboxed.Mutable qualified as UM
+import Data.Word (Word64)
 import Language.Haskell.TH.Syntax qualified as TH
 import Numeric.Log
 import System.IO.Unsafe
-import System.Random hiding (split, uniform)
+import System.Random hiding (random, randoms, split, uniform)
 import System.Random qualified as R
+import Unsafe.Coerce
 
 type HashMap k v = H.Dictionary (H.PrimState IO) M.MVector k UM.MVector v
 
@@ -48,24 +53,25 @@ data Tree = Tree
     split :: (Tree, Tree)
   }
 
-{-# INLINE newTree #-}
-newTree :: IORef (HashMap Integer Double, StdGen) -> Tree
-newTree s = go 1
+-- Let's draw doubles more uniformly than the standard generator
+-- based on https://www.corsix.org/content/higher-quality-random-floats
+random :: StdGen -> (Double, StdGen)
+random g0 =
+  let (r1 :: Word64, g1) = R.random g0
+      (r2 :: Word64, g2) = R.random g1
+      e = let r = countTrailingZeros r1 - 11 in if r >= 0 then countTrailingZeros r2 else r
+      dbl = (1 + (r1 `shiftR` 11)) `shiftR` 1 - ((unsafeCoerce e - 1011) `shiftL` 52)
+   in (unsafeCoerce dbl, g2)
+
+randoms :: StdGen -> [Double]
+randoms g =
+  let (x, g') = random g
+   in x : randoms g'
+
+newTree :: StdGen -> Tree
+newTree g0 = Tree x (newTree g1, newTree g2)
   where
-    go :: Integer -> Tree
-    go id =
-      Tree
-        ( unsafePerformIO $ do
-            (m, g) <- readIORef s
-            H.lookup m id >>= \case
-              Nothing -> do
-                let (x, g') = R.random g
-                H.insert m id x
-                writeIORef s (m, g')
-                pure x
-              Just x -> pure x
-        )
-        (go (2 * id), go (2 * id + 1))
+    (x, R.split -> (g1, g2)) = random g0
 
 newtype Prob a = Prob {runProb :: Tree -> a}
 
@@ -106,7 +112,7 @@ samples (Meas m) = map (second getProduct) . runProb f
     f = runWriterT m >>= \x -> (x :) <$> f
 
 {-# NOINLINE memoize #-}
-memoize :: Ord a => (a -> Prob b) -> Prob (a -> b)
+memoize :: (Ord a) => (a -> Prob b) -> Prob (a -> b)
 memoize f = unsafePerformIO $ do
   ref <- newIORef mempty
   pure $ Prob $ \t -> \x -> unsafePerformIO $ do