Commits

Aleksey Khudyakov committed c2cf05e

Fix looping in tableFromProbabilities

It's possible to construct probability array which will
send `correctWeights` into infinite loop by adding element which
have probability of 1:

> tableFromProbabilities $ U.fromList [(0,1),(1::Int,0)]

Fixes #17

Comments (0)

Files changed (1)

System/Random/MWC/CondensedTable.hs

        => v (a, Double) -> CondensedTable v a
 {-# INLINE tableFromProbabilities #-}
 tableFromProbabilities v
-  | G.null v  = pkgError "tableFromProbabilities" "empty vector of outcomes"
-  | otherwise = tableFromIntWeights $ G.map (second $ round . (* mlt)) v
+  | G.null tbl = pkgError "tableFromProbabilities" "empty vector of outcomes"
+  | otherwise  = tableFromIntWeights $ G.map (second $ round . (* mlt)) tbl
   where
-    mlt = 4.294967296e9 -- 2^32
+    mlt = 4.294967296e9            -- 2^32
+    tbl = G.filter ((> 0) . snd) v -- Drop non-positive probabilities
 
 -- | Same as 'tableFromProbabilities' but treats number as weights not
 -- probilities. Non-positive weights are discarded, and those
 -- | Generate a condensed lookup table from integer weights. Weights
 -- should sum to @2^32@. If they don't, the algorithm will alter the
 -- weights so that they do. This approach should work reasonably well
--- for rounding error.
+-- for rounding errors.
 tableFromIntWeights :: (Vector v (a,Word32), Vector v a, Vector v Word32)
                     => v (a, Word32)
                     -> CondensedTable v a