path: root/stdlib/source/lux/math/random.lux

(.module: {#.doc "Pseudo-random number generation (PRNG) algorithms."}
  [lux #- list i64 nat int rev char]
  (lux (control [functor #+ Functor]
                [apply #+ Apply]
                [monad #+ do Monad]
                hash)
       (data [bit]
             [text "text/" Monoid<Text>]
             (text [unicode #+ Char Segment])
             [product]
             [maybe]
             [number #+ hex]
             (number ["r" ratio]
                     ["c" complex])
             (coll [list "list/" Fold<List>]
                   [array]
                   (dictionary ["dict" unordered #+ Dict])
                   [queue #+ Queue]
                   (set ["set" unordered #+ Set])
                   [stack #+ Stack]
                   [row #+ Row]
                   (tree [finger #+ Tree])))
       (type [refinement #+ Refiner Refined])
       ))

(type: #export #rec PRNG
  {#.doc "An abstract way to represent any PRNG."}
  (-> Any [PRNG I64]))

(type: #export (Random a)
  {#.doc "A producer of random values based on a PRNG."}
  (-> PRNG [PRNG a]))

(structure: #export _ (Functor Random)
  (def: (map f fa)
    (function (_ state)
      (let [[state' a] (fa state)]
        [state' (f a)]))))

(structure: #export _ (Apply Random)
  (def: functor Functor<Random>)

  (def: (apply ff fa)
    (function (_ state)
      (let [[state' f] (ff state)
            [state'' a] (fa state')]
        [state'' (f a)]))))

(structure: #export _ (Monad Random)
  (def: functor Functor<Random>)

  (def: (wrap a)
    (function (_ state)
      [state a]))

  (def: (join ffa)
    (function (_ state)
      (let [[state' fa] (ffa state)]
        (fa state')))))

(def: #export (filter pred gen)
  {#.doc "Retries the generator until the output satisfies a predicate."}
  (All [a] (-> (-> a Bool) (Random a) (Random a)))
  (do Monad<Random>
    [sample gen]
    (if (pred sample)
      (wrap sample)
      (filter pred gen))))

(def: #export (refine refiner gen)
  {#.doc "Retries the generator until the output can be refined."}
  (All [t r] (-> (Refiner t r) (Random t) (Random (Refined t r))))
  (do Monad<Random>
    [sample gen]
    (case (refiner sample)
      (#.Some refined)
      (wrap refined)

      #.None
      (refine refiner gen))))

(def: #export bool
  (Random Bool)
  (function (_ prng)
    (let [[prng output] (prng [])]
      [prng (|> output (bit.and +1) (n/= +1))])))

(def: #export i64
  (Random I64)
  (function (_ prng)
    (let [[prng left] (prng [])
          [prng right] (prng [])]
      [prng (|> left
                (bit.left-shift +32)
                ("lux i64 +" right))])))

(def: #export nat
  (Random Nat)
  (:: Monad<Random> map .nat ..i64))

(def: #export int
  (Random Int)
  (:: Monad<Random> map .int ..i64))

(def: #export rev
  (Random Rev)
  (:: Monad<Random> map .rev ..i64))

(def: #export frac
  (Random Frac)
  (:: Monad<Random> map number.bits-to-frac nat))

(def: #export (char set)
  (-> unicode.Set (Random Char))
  (let [summary (finger.tag set)
        start (unicode.start summary)
        size (unicode.size summary)
        in-range (: (-> Char Char)
                    (|>> (n/% size) (n/+ start)))]
    (|> nat
        (:: Monad<Random> map in-range)
        (..filter (function (_ char)
                    (finger.found? (function (_ segment)
                                     (unicode.within? segment char))
                                   set))))))

(def: #export (text char-gen size)
  (-> (Random Char) Nat (Random Text))
  (if (n/= +0 size)
    (:: Monad<Random> wrap "")
    (do Monad<Random>
      [x char-gen
       xs (text char-gen (dec size))]
      (wrap (text/compose (text.from-code x) xs)))))

(def: #export unicode
  (-> Nat (Random Text))
  (text (char unicode.full)))

(do-template [<name> <type> <ctor> <gen>]
  [(def: #export <name>
     (Random <type>)
     (do Monad<Random>
       [left <gen>
        right <gen>]
       (wrap (<ctor> left right))))]

  [ratio   r.Ratio   r.ratio   nat]
  [complex c.Complex c.complex frac]
  )

(def: #export (seq left right)
  {#.doc "Sequencing combinator."}
  (All [a b] (-> (Random a) (Random b) (Random [a b])))
  (do Monad<Random>
    [=left left
     =right right]
    (wrap [=left =right])))

(def: #export (alt left right)
  {#.doc "Heterogeneous alternative combinator."}
  (All [a b] (-> (Random a) (Random b) (Random (| a b))))
  (do Monad<Random>
    [? bool]
    (if ?
      (do @
        [=left left]
        (wrap (+0 =left)))
      (do @
        [=right right]
        (wrap (+1 =right))))))

(def: #export (either left right)
  {#.doc "Homogeneous alternative combinator."}
  (All [a] (-> (Random a) (Random a) (Random a)))
  (do Monad<Random>
    [? bool]
    (if ?
      left
      right)))

(def: #export (rec gen)
  {#.doc "A combinator for producing recursive random generators."}
  (All [a] (-> (-> (Random a) (Random a)) (Random a)))
  (function (_ state)
    (let [gen' (gen (rec gen))]
      (gen' state))))

(def: #export (maybe value-gen)
  (All [a] (-> (Random a) (Random (Maybe a))))
  (do Monad<Random>
    [some? bool]
    (if some?
      (do @
        [value value-gen]
        (wrap (#.Some value)))
      (wrap #.None))))

(do-template [<name> <type> <zero> <plus>]
  [(def: #export (<name> size value-gen)
     (All [a] (-> Nat (Random a) (Random (<type> a))))
     (if (n/> +0 size)
       (do Monad<Random>
         [x value-gen
          xs (<name> (dec size) value-gen)]
         (wrap (<plus> x xs)))
       (:: Monad<Random> wrap <zero>)))]

  [list     List   (.list)      #.Cons]
  [row Row row.empty row.add]
  )

(do-template [<name> <type> <ctor>]
  [(def: #export (<name> size value-gen)
     (All [a] (-> Nat (Random a) (Random (<type> a))))
     (do Monad<Random>
       [values (list size value-gen)]
       (wrap (|> values <ctor>))))]

  [array Array array.from-list]
  [queue Queue queue.from-list]
  [stack Stack (list/fold stack.push stack.empty)]
  )

(def: #export (set Hash<a> size value-gen)
  (All [a] (-> (Hash a) Nat (Random a) (Random (Set a))))
  (if (n/> +0 size)
    (do Monad<Random>
      [xs (set Hash<a> (dec size) value-gen)]
      (loop [_ []]
        (do @
          [x value-gen
           #let [xs+ (set.add x xs)]]
          (if (n/= size (set.size xs+))
            (wrap xs+)
            (recur [])))))
    (:: Monad<Random> wrap (set.new Hash<a>))))

(def: #export (dict Hash<a> size key-gen value-gen)
  (All [k v] (-> (Hash k) Nat (Random k) (Random v) (Random (Dict k v))))
  (if (n/> +0 size)
    (do Monad<Random>
      [kv (dict Hash<a> (dec size) key-gen value-gen)]
      (loop [_ []]
        (do @
          [k key-gen
           v value-gen
           #let [kv+ (dict.put k v kv)]]
          (if (n/= size (dict.size kv+))
            (wrap kv+)
            (recur [])))))
    (:: Monad<Random> wrap (dict.new Hash<a>))))

(def: #export (run prng calc)
  (All [a] (-> PRNG (Random a) [PRNG a]))
  (calc prng))

(def: pcg-32-magic-mult Nat +6364136223846793005)

(def: #export (pcg-32 [inc seed])
  {#.doc "An implementation of the PCG32 algorithm.

          For more information, please see: http://www.pcg-random.org/"}
  (-> [(I64 Any) (I64 Any)] PRNG)
  (function (_ _)
    [(|> seed .nat (n/* pcg-32-magic-mult) ("lux i64 +" inc) [inc] pcg-32)
     (let [rot (|> seed .i64 (bit.logical-right-shift +59))]
       (|> seed
           (bit.logical-right-shift +18)
           (bit.xor seed)
           (bit.logical-right-shift +27)
           (bit.rotate-right rot)
           .i64))]))

(def: #export (xoroshiro-128+ [s0 s1])
  {#.doc "An implementation of the Xoroshiro128+ algorithm.

          For more information, please see: http://xoroshiro.di.unimi.it/"}
  (-> [(I64 Any) (I64 Any)] PRNG)
  (function (_ _)
    [(let [s01 (bit.xor s0 s1)]
       (xoroshiro-128+ [(|> s0
                            (bit.rotate-left +55)
                            (bit.xor s01)
                            (bit.xor (bit.left-shift +14 s01)))
                        (bit.rotate-left +36 s01)]))
     ("lux i64 +" s0 s1)]))

(def: (swap from to vec)
  (All [a] (-> Nat Nat (Row a) (Row a)))
  (|> vec
      (row.put to (maybe.assume (row.nth from vec)))
      (row.put from (maybe.assume (row.nth to vec)))))

(def: #export (shuffle seed row)
  {#.doc "Shuffle a row randomly based on a seed value."}
  (All [a] (-> Nat (Row a) (Row a)))
  (let [_size (row.size row)
        _shuffle (monad.fold Monad<Random>
                             (function (_ idx vec)
                               (do Monad<Random>
                                 [rand nat]
                                 (wrap (swap idx (n/% _size rand) vec))))
                             row
                             (list.n/range +0 (dec _size)))]
    (|> _shuffle
        (run (pcg-32 [+123 seed]))
        product.right)))