fix frequency

Author: safareli

.travis.yml

@@ -14,6 +14,7 @@ install:
   - bower install
 script:
   - npm run -s build
+  - npm run -s test
 after_success:
 - >-
   test $TRAVIS_TAG &&

bower.json

@@ -21,5 +21,12 @@
     "purescript-tuples": "^4.0.0",
     "purescript-unfoldable": "^3.0.0",
     "purescript-integers": "^3.0.0"
+  },
+  "devDependencies": {
+    "purescript-console": "^3.0.0",
+    "purescript-assert": "^3.0.0",
+    "purescript-arrays": "^4.2.1",
+    "purescript-transformers": "^3.4.0",
+    "purescript-math": "^2.1.0"
   }
 }

package.json

@@ -2,6 +2,7 @@
   "private": true,
   "scripts": {
     "clean": "rimraf output && rimraf .pulp-cache",
+    "test": "pulp test",
     "build": "pulp build -- --censor-lib --strict"
   },
   "devDependencies": {

src/Control/Monad/Gen.purs

@@ -14,7 +14,6 @@ import Control.Monad.Gen.Class (class MonadGen, Size, chooseBool, chooseFloat, c
 import Control.Monad.Rec.Class (class MonadRec, Step(..), tailRecM)
 
 import Data.Foldable (class Foldable, length, foldl, foldMap)
-import Data.Int as Int
 import Data.Maybe (Maybe(..), fromMaybe)
 import Data.Monoid.Additive (Additive(..))
 import Data.Newtype (alaF)
@@ -46,15 +45,26 @@ frequency
   => Foldable f
   => NonEmpty f (Tuple Number (m a))
   -> m a
-frequency (x :| xs) =
+frequency (x :| xs) = 
   let
     first = fst x
     total = first + alaF Additive foldMap fst xs
-  in do
-    pos <- chooseFloat 0.0 total
-    let n = Int.round (pos / total * length xs)
-    snd $ if n == 0 then x else fromIndex (n - 1) x xs
-
+  in 
+    chooseFloat 0.0 total >>= pick
+  where
+  pick pos = 
+    let
+      initial = go (Tuple 0.0 $ snd x) x
+    in 
+      snd $ foldl go initial xs
+    where
+    go (Tuple weight val) (Tuple currWeight currVal) =
+      let
+        nextWeight = weight + currWeight 
+      in 
+        if weight <= pos && pos <= nextWeight
+          then Tuple nextWeight currVal
+          else Tuple nextWeight val
 -- | Creates a generator that outputs a value chosen from a selection with
 -- | uniform probability.
 elements :: forall m f a. MonadGen m => Foldable f => NonEmpty f a -> m a

test/Test/Main.purs

@@ -0,0 +1,69 @@
+module Test.Main where
+
+import Prelude
+
+import Test.Assert (ASSERT, assert)
+import Partial.Unsafe (unsafeCrashWith)
+import Control.Monad.Eff (Eff)
+import Control.Monad.Eff.Console (CONSOLE, log)
+import Control.Monad.Gen (class MonadGen, frequency)
+import Data.NonEmpty ((:|), NonEmpty(..))
+import Data.Traversable (sequence)
+import Data.Array (replicate, group', length)
+import Data.Tuple (Tuple(..))
+import Data.Newtype (unwrap)
+import Math (remainder)
+import Control.Monad.State (State, class MonadState, get, put, evalStateT)
+
+type TestEffects = (console :: CONSOLE, assert :: ASSERT)
+type Tests = Eff TestEffects Unit
+
+newtype TestGen a = TestGen (State Number a)
+derive newtype instance testGenFunctor :: Functor TestGen
+derive newtype instance testGenApply :: Apply TestGen
+derive newtype instance testGenBind :: Bind TestGen
+derive newtype instance testGenApplicative :: Applicative TestGen
+derive newtype instance testGenMonad :: Monad TestGen
+derive newtype instance testGenMonadState :: MonadState Number TestGen
+
+instance testGenMonadGen :: MonadGen TestGen where
+  sized _ = unsafeCrashWith "sized should not be called"
+  resize _ _ = unsafeCrashWith "resize should not be called"
+  chooseBool = pure unit >>= \_ -> unsafeCrashWith "chooseBool should not be called"
+  chooseFloat s e = do 
+    c <- get
+    put (c + 1.0)
+    pure ((s + c) `remainder` e)
+  chooseInt _ _ = unsafeCrashWith "chooseFloat should not be called"
+
+runTestGen :: TestGen ~> State Number
+runTestGen (TestGen x) = x
+
+main :: Tests
+main = do
+  log "check frequency"
+  let
+    abcGen :: TestGen String
+    abcGen = 
+      frequency $ 
+        ( Tuple 10.0 $ pure "A" )  :| 
+        [ Tuple 20.0 $ pure "B" 
+        , Tuple 0.0 $ pure "Z" 
+        , Tuple 30.0 $ pure "C" 
+        , Tuple 40.0 $ pure "D" 
+        , Tuple 50.0 $ pure "E" 
+        , Tuple 50.0 $ pure "F" 
+        ]
+    abcArrGen = sequence $ replicate 200 abcGen
+    abcArr = runTestGen abcArrGen `evalStateT` 0.0 # unwrap
+    actual = group' abcArr
+      <#> \(NonEmpty x xs)  -> Tuple (length xs + 1) x
+    expected = 
+      [ (Tuple 10 "A")
+      , (Tuple 20 "B")
+      , (Tuple 30 "C")
+      , (Tuple 40 "D")
+      , (Tuple 50 "E")
+      , (Tuple 50 "F")
+      ]
+  assert (expected == actual)