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

(.module: {#.doc "Pseudo-random number generation (PRNG) algorithms."}
  [lux #- list]
  (lux (control [functor #+ Functor]
                [apply #+ Apply]
                [monad #+ do Monad]
                hash)
       (data [bit]
             [text "text/" Monoid<Text>]
             [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]
                   [sequence #+ Sequence]))
       ))

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

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

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

(struct: #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)]))))

(struct: #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 nat
  (Random Nat)
  (function (_ prng)
    (let [[prng left] (prng [])
          [prng right] (prng [])]
      [prng (n/+ (bit.left-shift +32 left)
                 right)])))

(def: #export int
  (Random Int)
  (:: Monad<Random> map nat-to-int nat))

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

(def: (bits n)
  (-> Nat (Random Nat))
  (function (_ prng)
    (let [[prng output] (prng [])]
      [prng (bit.logical-right-shift (n/- n +64) output)])))

(def: #export frac
  (Random Frac)
  (do Monad<Random>
    [left (bits +26)
     right (bits +27)]
    (wrap (|> right
              (n/+ (bit.left-shift +27 left))
              nat-to-int
              int-to-frac
              (f// (|> +1 (bit.left-shift +53) nat-to-int int-to-frac))))))

(def: #export deg
  (Random Deg)
  (:: Monad<Random> map (|>> (:! Deg)) nat))

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

(type: Char-Range [Nat Nat])

(do-template [<name> <from> <to>]
  [(def: <name> Char-Range [(hex <from>) (hex <to>)])]

  [Thaana               "+0780" "+07BF"]
  [Khmer-Symbols        "+19E0" "+19FF"]
  [Phonetic-Extensions "+1D00" "+1D7F"]
  [Hangul-Syllables     "+AC00" "+D7AF"]

  [Cypriot-Syllabary                       "+10800" "+1083F"]
  [Tai-Xuan-Jing-Symbols                   "+1D300" "+1D35F"]
  [Mathematical-Alphanumeric-Symbols       "+1D400" "+1D7FF"]
  [CJK-Unified-Ideographs-Extension-B      "+20000" "+2A6DF"]
  [CJK-Compatibility-Ideographs-Supplement "+2F800" "+2FA1F"]
  )

(def: (within? [from to] char)
  (-> Char-Range Nat Bool)
  (and (n/>= from char) (n/<= to char)))

(def: unicode-ceiling (n/inc (product.right CJK-Compatibility-Ideographs-Supplement)))

(def: #export unicode
  (Random Nat)
  (|> ..nat
      (:: Monad<Random> map (n/% unicode-ceiling))
      (..filter (function (_ raw)
                  ## From "Basic Latin" to "Syriac"
                  (or (n/<= (hex "+074F") raw)
                      (within? Thaana raw)
                      ## From "Devanagari" to "Ethiopic"
                      (and (n/>= (hex "+0900") raw)
                           (n/<= (hex "+137F") raw))
                      ## From "Cherokee" to "Mongolian"
                      (and (n/>= (hex "+13A0") raw)
                           (n/<= (hex "+18AF") raw))
                      ## From "Limbu" to "Tai Le"
                      (and (n/>= (hex "+1900") raw)
                           (n/<= (hex "+197F") raw))
                      (within? Khmer-Symbols raw)
                      (within? Phonetic-Extensions raw)
                      ## From "Latin Extended Additional" to "Miscellaneous Symbols and Arrows"
                      (and (n/>= (hex "+1E00") raw)
                           (n/<= (hex "+2BFF") raw))
                      ## From "CJK Radicals Supplement" to "Kangxi Radicals"
                      (and (n/>= (hex "+2E80") raw)
                           (n/<= (hex "+2FDF") raw))
                      ## From "Ideographic Description Characters" to "Bopomofo Extended"
                      (and (n/>= (hex "+2FF0") raw)
                           (n/<= (hex "+31BF") raw))
                      ## From "Katakana Phonetic Extensions" to "CJK Unified Ideographs"
                      (and (n/>= (hex "+31F0") raw)
                           (n/<= (hex "+9FAF") raw))
                      ## From "Yi Syllables" to "Yi Radicals"
                      (and (n/>= (hex "+A000") raw)
                           (n/<= (hex "+A4CF") raw))
                      (within? Hangul-Syllables raw)
                      ## From "CJK Compatibility Ideographs" to "Arabic Presentation Forms-A"
                      (and (n/>= (hex "+F900") raw)
                           (n/<= (hex "+FDFF") raw))
                      ## From "Combining Half Marks" to "Halfwidth and Fullwidth Forms"
                      (and (n/>= (hex "+FE20") raw)
                           (n/<= (hex "+FFEF") raw))
                      ## From "Linear B Syllabary" to "Aegean Numbers"
                      (and (n/>= (hex "+10000") raw)
                           (n/<= (hex "+1013F") raw))
                      ## From "Old Italic" to "Osmanya"
                      (and (n/>= (hex "+10300") raw)
                           (n/<= (hex "+104AF") raw))
                      (within? Cypriot-Syllabary raw)
                      ## From "Byzantine Musical Symbols" to "Musical Symbols"
                      (and (n/>= (hex "+1D000") raw)
                           (n/<= (hex "+1D1FF") raw))
                      (within? Tai-Xuan-Jing-Symbols raw)
                      (within? Mathematical-Alphanumeric-Symbols raw)
                      (within? CJK-Unified-Ideographs-Extension-B raw)
                      (within? CJK-Compatibility-Ideographs-Supplement raw)
                      )))))

(def: #export (text size)
  (-> Nat (Random Text))
  (text' unicode size))

(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> (n/dec size) value-gen)]
         (wrap (<plus> x xs)))
       (:: Monad<Random> wrap <zero>)))]

  [list     List   (.list)      #.Cons]
  [sequence Sequence sequence.empty sequence.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> (n/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> (n/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/"}
  (-> [Nat Nat] PRNG)
  (function (_ _)
    (let [seed' (|> seed (n/* pcg-32-magic-mult) (n/+ inc))
          xor-shifted (|> seed (bit.logical-right-shift +18) (bit.xor seed) (bit.logical-right-shift +27))
          rot (|> seed (bit.logical-right-shift +59))]
      [(pcg-32 [inc seed']) (bit.rotate-right rot xor-shifted)]
      )))

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

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

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

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