diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,11 @@ +# Changelog for lazyppl-singlesite + +## 0.1.0 + +- The single-site Metropolis-Hastings sampler `mh1` was part of `lazyppl` 1.0.1. From `lazyppl` 1.1.0 it is in this package, so that `lazyppl` no longer depends on `ghc-heap`. +- Import `LazyPPL.SingleSite` as before and depend on `lazyppl-singlesite`, which works with `lazyppl ==1.1.*`. +- `lazyppl-singlesite` depends on `ghc-heap >=9.6.1 && <9.15` and is tested with GHC 9.6.7, 9.12.4 and 9.14.1. +- `mh1` changes one random choice per step. It finds the random choices that a model has evaluated by inspecting GHC's heap, so garbage collection can change which ones it finds. +- `mh1` rejects a proposal when the reverse site is absent from the proposed tree's evaluated sites. Otherwise, while the current weight is zero, it accepts every proposal. +- `mh1` accepts a move when its uniform draw is below the acceptance ratio. It used to accept on equality, so a uniform draw of 0 accepted a move to weight 0. +- `mh1` draws its proposed labels from the open interval (0,1) and takes its generator from `newStdGen`. It raises an error when a total log weight overflows to positive infinity. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,7 @@ +Copyright 2021 Sam Staton and others + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/lazyppl-singlesite.cabal b/lazyppl-singlesite.cabal new file mode 100644 --- /dev/null +++ b/lazyppl-singlesite.cabal @@ -0,0 +1,86 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.39.6. +-- +-- see: https://github.com/sol/hpack + +name: lazyppl-singlesite +version: 0.1.0 +synopsis: Single-site Metropolis-Hastings for LazyPPL +description: Single-site Metropolis-Hastings for LazyPPL. Import LazyPPL.SingleSite + to use mh1. The sampler discovers evaluated random choices through GHC's heap. +category: Statistics +stability: experimental +homepage: https://lazyppl-team.github.io/ +bug-reports: https://github.com/lazyppl-team/lazyppl/issues +author: LazyPPL team +maintainer: sam.staton@cs.ox.ac.uk +copyright: 2021-2026 LazyPPL team +license: MIT +license-file: LICENSE +build-type: Simple +tested-with: + GHC == 9.6.7 + , GHC == 9.12.4 + , GHC == 9.14.1 +extra-source-files: + ChangeLog.md + +source-repository head + type: git + location: https://github.com/lazyppl-team/lazyppl + +library + exposed-modules: + LazyPPL.SingleSite + LazyPPL.SingleSite.Internal + hs-source-dirs: + src + build-depends: + base >=4.18.3.0 && <4.23 + , containers >=0.6.7 && <0.9 + , deepseq >=1.4.8.1 && <1.6 + , ghc-heap >=9.6.1 && <9.15 + , lazyppl ==1.1.* + , log-domain >=0.13.2 && <0.14 + , random >=1.3.1 && <1.4 + , transformers >=0.6.1.0 && <0.7 + default-language: Haskell2010 + +test-suite reverse-support-test + type: exitcode-stdio-1.0 + main-is: Api.hs + other-modules: + ReverseSupport + hs-source-dirs: + test + build-depends: + base >=4.18.3.0 && <4.23 + , containers >=0.6.7 && <0.9 + , deepseq >=1.4.8.1 && <1.6 + , ghc-heap >=9.6.1 && <9.15 + , lazyppl ==1.1.* + , lazyppl-singlesite + , log-domain >=0.13.2 && <0.14 + , random >=1.3.1 && <1.4 + , transformers >=0.6.1.0 && <0.7 + default-language: Haskell2010 + +test-suite site-discovery-test + type: exitcode-stdio-1.0 + main-is: SiteDiscovery.hs + hs-source-dirs: + test + ghc-options: -threaded -rtsopts "-with-rtsopts=-T -A128M" + build-depends: + QuickCheck >=2.14.2 && <2.20 + , base >=4.18.3.0 && <4.23 + , containers >=0.6.7 && <0.9 + , deepseq >=1.4.8.1 && <1.6 + , ghc-heap >=9.6.1 && <9.15 + , lazyppl ==1.1.* + , lazyppl-singlesite + , log-domain >=0.13.2 && <0.14 + , random >=1.3.1 && <1.4 + , transformers >=0.6.1.0 && <0.7 + default-language: Haskell2010 diff --git a/src/LazyPPL/SingleSite.hs b/src/LazyPPL/SingleSite.hs new file mode 100644 --- /dev/null +++ b/src/LazyPPL/SingleSite.hs @@ -0,0 +1,11 @@ +{-| Single-site [Metropolis-Hastings](https://en.wikipedia.org/wiki/Metropolis-Hastings_algorithm). + +@mh1 m@ gets out an infinite lazy stream of (result,weight) pairs from the +unnormalized measure @m@, the first of which is the initial draw. At each step +after that, we randomly change one site (a node in the random tree), and then +accept or reject the change. If rejected, we repeat the previous pair. + +Use 'drop' for burn-in and @LazyPPL.every@ for thinning. -} +module LazyPPL.SingleSite (mh1) where + +import LazyPPL.SingleSite.Internal (mh1) diff --git a/src/LazyPPL/SingleSite/Internal.hs b/src/LazyPPL/SingleSite/Internal.hs new file mode 100644 --- /dev/null +++ b/src/LazyPPL/SingleSite/Internal.hs @@ -0,0 +1,205 @@ +{-# LANGUAGE MagicHash, UnboxedTuples #-} +{-# LANGUAGE ScopedTypeVariables, CPP + #-} +module LazyPPL.SingleSite.Internal where + +import LazyPPL +import LazyPPL.Internal.Sampler (randomLabel, checkedLogWeight) +import Control.Monad.Trans.Writer + +import Data.Monoid +import Numeric.Log +import qualified Data.Map as M + +import System.Random hiding (uniform) + +import Control.DeepSeq +import Control.Exception (evaluate) + +import System.IO.Unsafe +#ifdef LAZYPPL_BROWSER +import Foreign.Marshal.Alloc (alloca) +import Foreign.Marshal.Array (pokeArray) +import Foreign.Ptr (castPtr) +import Foreign.Storable (peek, sizeOf) +#else +import Unsafe.Coerce +#endif + + +import Data.IORef +import GHC.Exts (Ptr(..), unpackClosure#, indexArray#, sizeofArray#, Int(I#)) +import GHC.Exts.Heap.Utils (dataConNames) +import GHC.Exts.Heap + + +{-- | Single-site Metropolis-Hastings --} + +-- | A partial tree, used for examining finite evaluated subtrees of the working infinite +-- | random tree. +data PTree = PTree (Maybe Double) [Maybe PTree] deriving Show + +type Site = [Int] +type Subst = M.Map Site Double + +getSites :: PTree -> [Site] +getSites p = getSites' p [] + +getSites' :: PTree -> Site -> [Site] +getSites' (PTree (Just v) ts) site = site : concat [ getSites' t (n:site) | (n, Just t) <- zip [0..] ts ] +getSites' (PTree Nothing ts) site = concat [ getSites' t (n:site) | (n, Just t) <- zip [0..] ts ] + +mutateNode :: Tree -> Site -> Double -> Tree +mutateNode (Tree _ ts) [] d = Tree d ts +mutateNode (Tree v ts) (n:ns) d = Tree v $ take n ts ++ mutateNode (ts!!n) ns d : drop (n+1) ts + +mutateNodes :: Tree -> Subst -> Tree +mutateNodes = M.foldrWithKey (\k d t -> mutateNode t k d) + +randomElement :: RandomGen g => g -> [a] -> (g, a) +randomElement g xs = if null xs then error "0d sample space" else (g', xs !! n) + where (n, g') = uniformR (0, length xs - 1) g + -- ^^^^^^^^ + -- Depending on the version of `random` this can be either uniformR or randomR. + +mh1 :: forall a. NFData a => Meas a -> IO [(a, Product (Log Double))] +mh1 (Meas m) = do + g <- newStdGen + let (gTree,g') = split g + tree = randomTree gTree + (x0, w0) = checkedSample (runProb (runWriterT m) tree) + p = unsafePerformIO $ do { evaluate (rnf x0); evaluate (rnf w0) ; trunc tree } + samples = map (\(_,_,_,_,s) -> s) $ iterate (mhStep (Meas m)) (gTree, g', M.empty, p, (x0, w0)) + return samples + +-- | One transition of 'mh1'. The state holds the seed of the initial tree, the +-- proposal seed, the accepted substitutions, the evaluated sites, and the +-- current sample with its weight. +mhStep :: forall a g. (NFData a, RandomGen g) => Meas a -> (g, g, Subst, PTree, (a, Product (Log Double))) -> (g, g, Subst, PTree, (a, Product (Log Double))) +mhStep (Meas m) (treeSeed, seed, sub, ptree, (x,w)) = + let (seed1, seed2) = split seed + sites = getSites ptree + -- A measure that reads no random site has a one-point sample space: + -- there is no site to mutate, so the chain keeps its sample. + in if null sites + then (treeSeed, seed2, sub, ptree, (x,w)) + else + let (seed1', randSite) = (\(x,y) -> (x, reverse y)) $ randomElement seed1 sites + (newNode :: Double, seed1'') = randomLabel seed1' + (u :: Double, _) = randomLabel seed1'' + sub' = M.insert randSite newNode sub + t' = mutateNodes (randomTree treeSeed) sub' + (x',w') = checkedSample (runProb (runWriterT m) t') + ptree' = unsafePerformIO $ do { evaluate (rnf x'); evaluate (rnf w'); trunc t' } + sites' = getSites ptree' + alpha = fromIntegral (length sites) / fromIntegral (length sites') + * exp (ln (getProduct w' / getProduct w)) + -- The reverse proposal cannot select a site absent from the proposed tree. + -- A move is accepted with probability min 1 alpha, so u < alpha; with + -- u <= alpha, u = 0 would accept a move to weight 0. u is a label, in + -- (0,1): random's 1 would reject a move with alpha = 1. + in if reverse randSite `elem` sites' && (getProduct w == 0 || u < alpha) + then (treeSeed, seed2, sub', ptree', (x',w')) + else (treeSeed, seed2, sub, ptree, (x,w)) + +-- A product of valid weights can still overflow to a logarithm of +Infinity. +checkedSample :: (a, Product (Log Double)) -> (a, Product (Log Double)) +checkedSample (v, Product l) = (v, Product (checkedLogWeight "mh1" l)) + +-- Functions for truncating a tree. +getGCClosureData inspect b = do + c <- inspect b + case c of + BlackholeClosure _ n -> getGCClosureData inspect n + _ -> return c + +helperT :: (Box -> IO Closure) -> Box -> IO PTree +helperT inspect b = do + c <- getGCClosureData inspect b + case c of + ConstrClosure _ [n,l] [] _ _ "Tree" -> + do n' <- getGCClosureData inspect n + l' <- getGCClosureData inspect l + let missing :: IO x + missing = return $ error $ "Missing case:\n" ++ show n' ++ "\n" ++ show l' + -- The label is a site once it is evaluated, and the children were + -- demanded once the list is evaluated; each field is read on its own, + -- so a node with both fields unevaluated has no site. + label <- case n' of +#ifdef LAZYPPL_BROWSER + ConstrClosure {dataArgs = ds, name = "D#"} -> Just <$> doubleFromWords ds +#else + ConstrClosure {dataArgs = [d], name = "D#"} -> return (Just (unsafeCoerce d)) +#endif + ThunkClosure {} -> return Nothing + SelectorClosure {} -> return Nothing + _ -> missing + children <- case l' of + ConstrClosure {name = ":"} -> helperB inspect l + ThunkClosure {} -> return [] + _ -> missing + return $ PTree label children + ThunkClosure {} -> return $ PTree Nothing [] + +#ifdef LAZYPPL_BROWSER +-- A Double occupies one 64-bit word or two 32-bit words, in host byte order. +doubleFromWords :: [Word] -> IO Double +doubleFromWords ws + | length ws * sizeOf (0 :: Word) /= sizeOf (0 :: Double) = + fail "Unexpected Double heap payload size" + | otherwise = alloca $ \ptr -> do + pokeArray (castPtr ptr) ws + peek ptr +#endif + +helperB :: (Box -> IO Closure) -> Box -> IO [Maybe PTree] +helperB inspect b = do + c <- getGCClosureData inspect b + case c of + ConstrClosure _ [n,l] [] _ _ ":" -> + do n' <- getGCClosureData inspect n + l' <- getGCClosureData inspect l + case (n',l') of + (ConstrClosure {name = "Tree"}, ConstrClosure {name = ":"}) -> + do n'' <- helperT inspect n + l'' <- helperB inspect l + return $ Just n'' : l'' + (ConstrClosure {name = "Tree"}, ThunkClosure {}) -> + do n'' <- helperT inspect n + return [Just n''] + (ThunkClosure {}, ConstrClosure {name = ":"}) -> + do l'' <- helperB inspect l + return $ Nothing : l'' + (ThunkClosure {}, ThunkClosure {}) -> + return [] -- alternatively, Nothing : [] + _ -> return $ error $ "Missing case:\n" ++ show n' ++ "\n" ++ show l' + ThunkClosure {} -> return [] + +trunc :: Tree -> IO PTree +trunc t = do + cache <- newIORef M.empty + helperT (getCachedClosureData cache) (asBox t) + + +-- Constructor descriptions are immutable during this traversal. +getCachedClosureData :: IORef (M.Map (Ptr StgInfoTable) (String, String, String)) -> Box -> IO Closure +getCachedClosureData cache (Box x) = + case unpackClosure# x of + (# infoAddress, heapRep, pointers #) -> do + let infoPointer = Ptr infoAddress + boxes = [case indexArray# pointers i of { (# p #) -> Box p } + | I# i <- [0 .. I# (sizeofArray# pointers) - 1]] + names = do + known <- readIORef cache + case M.lookup infoPointer known of + Just result -> return result + Nothing -> do + result <- dataConNames infoPointer + writeIORef cache (M.insert infoPointer result known) + return result + table <- peekItbl infoPointer + case tipe table of + TSO -> return (UnsupportedClosure table) + STACK -> return (UnsupportedClosure table) + _ -> getClosureDataFromHeapRepPrim names (const (return Nothing)) table heapRep boxes + diff --git a/test/Api.hs b/test/Api.hs new file mode 100644 --- /dev/null +++ b/test/Api.hs @@ -0,0 +1,134 @@ +{-# OPTIONS_GHC -fno-cse -fno-full-laziness #-} +-- A constant test tree compiled as a static closure would end in IND_STATIC +-- indirections, which the single-site traversal does not read. +module Main where + +import Control.DeepSeq (rnf) +import Control.Exception (evaluate) +import Control.Monad (unless) +import Control.Monad.Trans.Writer (runWriterT) +import Data.IORef (newIORef, readIORef) +import qualified Data.Map as M +import Data.Monoid (getProduct) +import Data.Word (Word64) +import LazyPPL (Meas(..), randomTree, runProb, sample, score, uniform) +import LazyPPL.Distributions (bernoulli) +import LazyPPL.SingleSite (mh1) +import qualified LazyPPL.SingleSite.Internal as S +import ReverseSupport (reverseSupport) +import System.Random (RandomGen(..), SplitGen(..), mkStdGen, setStdGen) +import System.Timeout (timeout) + +main :: IO () +main = do + supported <- reverseSupport + unless supported (fail "single-site rejects proposals without reverse support") + putStrLn "PASS single-site rejects proposals without reverse support" + constant <- constantMeasureRepeats + unless constant (fail "mh1 does not repeat the sample of a site-free measure") + putStrLn "PASS mh1 repeats the sample of a site-free measure" + rejected <- zeroWeightMoveRejected + unless rejected (fail "mh1 accepted a move from positive weight to weight 0") + putStrLn "PASS mh1 rejects a move from positive weight to weight 0" + accepted <- ratioOneMoveAccepted + unless accepted (fail "mh1 rejected a move whose acceptance ratio is 1") + putStrLn "PASS mh1 accepts a move whose acceptance ratio is 1" + setStdGen (mkStdGen 23) + coin <- map fst . take 60000 . drop 1000 <$> mh1 coinModel + posterior <- posteriorFraction "mh1 coin posterior" coin + unless posterior (fail "mh1 does not sample the coin posterior") + +-- A generator that returns one word forever. +newtype ConstWord = ConstWord Word64 + +instance RandomGen ConstWord where + genWord64 g@(ConstWord w) = (w, g) + +instance SplitGen ConstWord where + splitGen g = (g, g) + +-- A move from positive weight to weight 0 has acceptance ratio 0, so it is +-- never accepted. The word 2^64-1 makes the proposed label and the acceptance +-- uniform as small as the generator allows; the word 2^63 puts every label of +-- the initial tree near 1/2, where the weight is 1. +zeroWeightMoveRejected :: IO Bool +zeroWeightMoveRejected = do + -- Read through an IORef, so the compiler cannot build the tree as a static closure. + word <- newIORef (2 ^ (63 :: Int)) >>= readIORef + let model = do x <- sample uniform; score (if x > 0.1 then 1 else 0); pure x + Meas writer = model + treeSeed = ConstWord word + tree = randomTree treeSeed + (x0, w0) = runProb (runWriterT writer) tree + evaluate (rnf x0) + evaluate (rnf w0) + sites <- S.trunc tree + let (_, _, _, _, (x1, w1)) = S.mhStep model (treeSeed, ConstWord maxBound, M.empty, sites, (x0, w0)) + answered <- timeout 2000000 (evaluate (rnf x1)) + case answered of + Nothing -> putStrLn "zero-weight move: the step gave no answer within 2 s" >> pure False + Just () -> do + putStrLn $ "zero-weight move: from x=" ++ show x0 ++ " (weight " ++ show (getProduct w0) + ++ ") the step returned x=" ++ show x1 ++ " (weight " ++ show (getProduct w1) ++ ")" + pure (x1 == x0 && getProduct w1 == getProduct w0) + +-- A move between two samples of weight 1, each with one site, has acceptance +-- ratio 1, so it is always accepted. The word 0 makes the acceptance uniform as +-- large as the generator allows; the word 2^63 puts every label of the initial +-- tree near 1/2, and the proposed label is the largest label. +ratioOneMoveAccepted :: IO Bool +ratioOneMoveAccepted = do + -- Read through an IORef, so the compiler cannot build the tree as a static closure. + word <- newIORef (2 ^ (63 :: Int)) >>= readIORef + let model = do x <- sample uniform; score 1; pure x + Meas writer = model + treeSeed = ConstWord word + tree = randomTree treeSeed + (x0, w0) = runProb (runWriterT writer) tree + evaluate (rnf x0) + evaluate (rnf w0) + sites <- S.trunc tree + let (_, _, _, _, (x1, _)) = S.mhStep model (treeSeed, ConstWord 0, M.empty, sites, (x0, w0)) + answered <- timeout 2000000 (evaluate (rnf x1)) + case answered of + Nothing -> putStrLn "ratio-one move: the step gave no answer within 2 s" >> pure False + Just () -> do + putStrLn $ "ratio-one move: from x=" ++ show x0 ++ " the step returned x=" ++ show x1 + pure (x1 /= x0) + +-- A measure that reads no random site has a one-point sample space, and the +-- documented chain repeats its initial sample and weight, rather than failing +-- on an empty site list. `return` never splits the tree, while `sample (pure ())` +-- splits it once through the writer's fmap. +constantMeasureRepeats :: IO Bool +constantMeasureRepeats = do + pureMeasure <- take 5 <$> mh1 (return ()) + lifted <- take 5 <$> mh1 (sample (pure ())) + pure $ and [ map fst samples == replicate 5 () + && all (\(_, weight) -> getProduct weight == 1) samples + | samples <- [pureMeasure, lifted] ] + +-- A fair coin b. On true, a uniform u is read and scored 6u, which integrates +-- to 3; on false, the score is 1. So P(b) = 3/4 exactly, and the site u exists +-- only on the true branch. +coinModel :: Meas Bool +coinModel = do + b <- sample (bernoulli 0.5) + if b then sample uniform >>= \u -> score (6 * u) else score 1 + pure b + +-- The fraction of true draws is within five standard errors of 3/4. The +-- standard error comes from the means of 40 consecutive batches, and it must be +-- below 0.01 for the check to say anything. +posteriorFraction :: String -> [Bool] -> IO Bool +posteriorFraction name draws = do + let n = length draws + batches = 40 + size = n `div` batches + batchMean k = fromIntegral (length (filter id (take size (drop (k * size) draws)))) / fromIntegral size :: Double + means = map batchMean [0 .. batches - 1] + p = sum means / fromIntegral batches + se = sqrt (sum [(m - p) ^ (2 :: Int) | m <- means] / fromIntegral (batches - 1) / fromIntegral batches) + ok = se < 0.01 && abs (p - 0.75) < 5 * se + putStrLn $ (if ok then "PASS " else "FAIL ") ++ name ++ ": P(true) = " ++ show p ++ ", standard error " ++ show se ++ ", exact 3/4" + pure ok diff --git a/test/ReverseSupport.hs b/test/ReverseSupport.hs new file mode 100644 --- /dev/null +++ b/test/ReverseSupport.hs @@ -0,0 +1,52 @@ +{-# OPTIONS_GHC -fno-cse -fno-full-laziness #-} +module ReverseSupport (reverseSupport) where + +import Control.DeepSeq (NFData(..)) +import Control.Exception (evaluate) +import Control.Monad (forM) +import Data.IORef +import LazyPPL (Prob(..), Tree(..), sample) +import LazyPPL.SingleSite (mh1) +import System.IO.Unsafe (unsafePerformIO) +import System.Mem.StableName (makeStableName) +import System.Random (mkStdGen, setStdGen) + +-- We force the unused root label once, to make the root a site. A root +-- proposal leaves the new label lazy, so the proposed tree has no root +-- site (no reverse support). +data Observation = Observation Double Tree (IORef Bool) + +instance NFData Observation where + rnf (Observation value (Tree root _) first) = unsafePerformIO $ do + _ <- evaluate (rnf value) + initial <- atomicModifyIORef' first (\b -> (False, b)) + if initial then evaluate root >> pure () else pure () + +observe :: IORef Bool -> Tree -> Observation +observe first tree@(Tree _ (Tree value _ : _)) = Observation value tree first +observe _ _ = error "random tree has no first child" +{-# NOINLINE observe #-} + +valueOf :: Observation -> Double +valueOf (Observation value _ _) = value + +reverseSupport :: IO Bool +reverseSupport = do + results <- forM [1..16] $ \seed -> do + setStdGen (mkStdGen seed) + first <- newIORef True + chain <- mh1 (sample (Prob (observe first))) + snapshots <- forM (take 64 chain) $ \(draw, _) -> do + _ <- evaluate (rnf draw) + identity <- makeStableName =<< evaluate draw + pure (valueOf draw, identity) + let pairs = zip snapshots (drop 1 snapshots) + unsupported = length [() | ((x, a), (y, b)) <- pairs, x == y, a /= b] + moves = length [() | ((x, _), (y, _)) <- pairs, x /= y] + putStrLn $ "reverse support seed=" ++ show seed + ++ " accepted unused-root proposals=" ++ show unsupported + ++ " demanded-site moves=" ++ show moves + pure (unsupported, moves) + -- Rejecting every proposal would pass the unused-root check, so each chain + -- must also move. + pure (all ((== 0) . fst) results && all ((> 0) . snd) results) diff --git a/test/SiteDiscovery.hs b/test/SiteDiscovery.hs new file mode 100644 --- /dev/null +++ b/test/SiteDiscovery.hs @@ -0,0 +1,72 @@ +module Main where + +import Control.DeepSeq (force) +import Control.Exception (SomeException, evaluate, try) +import Control.Monad (forM_, unless) +import GHC.Exts.Heap (Closure, GenClosure(..), asBox, getBoxedClosureData) +import GHC.Stats (getRTSStats, gcs) +import LazyPPL (Tree(..), randomTree) +import qualified LazyPPL.SingleSite.Internal as S +import System.Mem (performGC) +import System.Random (mkStdGen) +import Test.QuickCheck +import Test.QuickCheck.Random (mkQCGen) + +at :: [Int] -> Tree -> Double +at [] (Tree x _) = x +at (i:is) (Tree _ ts) = at is (ts !! i) + +-- Compare the actual traversal with GHC's uncached decoder on the same +-- evaluated tree. Both traversals run without an intervening collection. +sameSites :: Property +sameSites = forAll arbitrary $ \seed -> + forAll (listOf1 (resize 4 (listOf (chooseInt (0,5))))) $ \paths -> ioProperty $ do + let tree = randomTree (mkStdGen seed) + forM_ paths (evaluate . (`at` tree)) + performGC + before <- getRTSStats + cached <- S.trunc tree >>= evaluate . force . S.getSites + plain <- S.helperT getBoxedClosureData (asBox tree) >>= evaluate . force . S.getSites + after <- getRTSStats + pure $ counterexample (show (seed,paths,cached,plain,gcs before,gcs after)) $ + cached === plain .&&. gcs before === gcs after + +-- A tree node evaluated to its constructor, with its label and its children +-- both unevaluated. The WebAssembly build meets such nodes in the clustering +-- demo; each field here is a call to a function that is not inlined, so it +-- stays a thunk. Nothing at or below the node was demanded, so the traversal +-- finds no site. +{-# NOINLINE opaqueLabel #-} +opaqueLabel :: Int -> Double +opaqueLabel k = fromIntegral k / 7 + +{-# NOINLINE opaqueChildren #-} +opaqueChildren :: Int -> [Tree] +opaqueChildren k = [lazyNode (k + 1)] + +{-# NOINLINE lazyNode #-} +lazyNode :: Int -> Tree +lazyNode k = Tree (opaqueLabel k) (opaqueChildren k) + +isThunk :: Closure -> Bool +isThunk ThunkClosure {} = True +isThunk _ = False + +unevaluatedNode :: IO Bool +unevaluatedNode = do + node@(Tree label children) <- evaluate (lazyNode 3) + shapes <- mapM getBoxedClosureData [asBox label, asBox children] + cached <- try (S.trunc node >>= evaluate . force . S.getSites) + plain <- try (S.helperT getBoxedClosureData (asBox node) >>= evaluate . force . S.getSites) + let shown = either (\e -> "an error: " ++ show (e :: SomeException)) show + ok = all isThunk shapes && either (const False) null cached && either (const False) null plain + putStrLn $ (if ok then "PASS" else "FAIL") ++ " a node with an unevaluated label and children has no site: fields " + ++ show (map isThunk shapes) ++ " thunks; trunc gives " ++ shown cached ++ "; helperT gives " ++ shown plain + pure ok + +main :: IO () +main = do + unevaluated <- unevaluatedNode + result <- quickCheckWithResult stdArgs {maxSuccess=200,replay=Just(mkQCGen 20261026,0)} sameSites + unless unevaluated (fail "A node with an unevaluated label and children was not traversed") + unless (isSuccess result) (fail "Cached and uncached traversal differ without intervening GC")