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

(.module: {#.doc "Pseudo-random number generation (PRNG) algorithms."}
  [lux (#- or and list i64 nat int rev char)
   [abstract
    [functor (#+ Functor)]
    [apply (#+ Apply)]
    ["." monad (#+ do Monad)]
    hash]
   [data
    ["." product]
    ["." maybe]
    [number (#+ hex)
     ["." i64]
     ["n" nat]
     ["i" int]
     ["r" ratio]
     ["c" complex]
     ["f" frac]]
    ["." text (#+ Char) ("#@." monoid)
     ["." unicode]]
    [collection
     ["." list ("#@." fold)]
     ["." array (#+ Array)]
     ["." dictionary (#+ Dictionary)]
     ["." queue (#+ Queue)]
     ["." set (#+ Set)]
     ["." stack (#+ Stack)]
     ["." row (#+ Row)]
     [tree
      ["." finger (#+ Tree)]]]]
   [time
    ["." instant (#+ Instant)]
    ["." date (#+ Date)]
    ["." duration (#+ Duration)]
    ["." month (#+ Month)]
    ["." day (#+ Day)]]
   [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
  (Functor Random)
  
  (def: (map f fa)
    (function (_ state)
      (let [[state' a] (fa state)]
        [state' (f a)]))))

(structure: #export apply
  (Apply Random)
  
  (def: &functor ..functor)

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

(structure: #export monad
  (Monad Random)
  
  (def: &functor ..functor)

  (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 Bit) (Random a) (Random a)))
  (do ..monad
    [sample gen]
    (if (pred sample)
      (wrap sample)
      (filter pred gen))))

(def: #export (one check random)
  (All [a b]
    (-> (-> a (Maybe b)) (Random a) (Random b)))
  (do ..monad
    [sample random]
    (case (check sample)
      (#.Some output)
      (wrap output)

      #.None
      (one check random))))

(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
    [sample gen]
    (case (refiner sample)
      (#.Some refined)
      (wrap refined)

      #.None
      (refine refiner gen))))

(def: #export bit
  (Random Bit)
  (function (_ prng)
    (let [[prng output] (prng [])]
      [prng (|> output (i64.and 1) (n.= 1))])))

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

(template [<name> <type> <cast>]
  [(def: #export <name>
     (Random <type>)
     (:: ..monad map <cast> ..i64))]

  [nat Nat .nat]
  [int Int .int]
  [rev Rev .rev]
  )

(def: #export frac
  (Random Frac)
  (:: ..monad map (|>> .i64 f.from-bits) ..nat))

(def: #export safe-frac
  (Random Frac)
  (let [mantissa-range (.int (i64.left-shift 53 1))
        mantissa-max (i.frac (dec mantissa-range))]
    (:: ..monad map
        (|>> (i.% mantissa-range)
             i.frac
             (f./ mantissa-max))
        ..int)))

(def: #export (char set)
  (-> unicode.Set (Random Char))
  (let [[start end] (unicode.range set)
        size (n.- start end)
        in-range (: (-> Char Char)
                    (|>> (n.% size) (n.+ start)))]
    (|> ..nat
        (:: ..monad map in-range)
        (..filter (unicode.member? set)))))

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

(template [<name> <set>]
  [(def: #export <name>
     (-> Nat (Random Text))
     (text (char <set>)))]

  [unicode           unicode.full]
  [ascii             unicode.ascii]
  [ascii/alpha       unicode.ascii/alpha]
  [ascii/alpha-num   unicode.ascii/alpha-num]
  [ascii/upper-alpha unicode.ascii/upper-alpha]
  [ascii/lower-alpha unicode.ascii/lower-alpha]
  )

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

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

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

(def: #export (or left right)
  {#.doc "Heterogeneous alternative combinator."}
  (All [a b] (-> (Random a) (Random b) (Random (| a b))))
  (do {! ..monad}
    [? bit]
    (if ?
      (do !
        [=left left]
        (wrap (0 #0 =left)))
      (do !
        [=right right]
        (wrap (0 #1 =right))))))

(def: #export (either left right)
  {#.doc "Homogeneous alternative combinator."}
  (All [a] (-> (Random a) (Random a) (Random a)))
  (do ..monad
    [? bit]
    (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}
    [some? bit]
    (if some?
      (do !
        [value value-gen]
        (wrap (#.Some value)))
      (wrap #.None))))

(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
         [x value-gen
          xs (<name> (dec size) value-gen)]
         (wrap (<plus> x xs)))
       (:: ..monad wrap <zero>)))]

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

(template [<name> <type> <ctor>]
  [(def: #export (<name> size value-gen)
     (All [a] (-> Nat (Random a) (Random (<type> a))))
     (do ..monad
       [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}
      [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 wrap (set.new Hash<a>))))

(def: #export (dictionary Hash<a> size key-gen value-gen)
  (All [k v] (-> (Hash k) Nat (Random k) (Random v) (Random (Dictionary k v))))
  (if (n.> 0 size)
    (do {! ..monad}
      [kv (dictionary Hash<a> (dec size) key-gen value-gen)]
      (loop [_ []]
        (do !
          [k key-gen
           v value-gen
           #let [kv+ (dictionary.put k v kv)]]
          (if (n.= size (dictionary.size kv+))
            (wrap kv+)
            (recur [])))))
    (:: ..monad wrap (dictionary.new Hash<a>))))

(def: #export instant
  (Random Instant)
  (:: ..monad map instant.from-millis ..int))

(def: #export date
  (Random Date)
  (:: ..monad map instant.date ..instant))

(def: #export duration
  (Random Duration)
  (:: ..monad map duration.from-millis ..int))

(def: #export month
  (Random Month)
  (let [(^open "/@.") ..monad]
    (..either (..either (..either (/@wrap #month.January)
                                  (..either (/@wrap #month.February)
                                            (/@wrap #month.March)))
                        (..either (/@wrap #month.April)
                                  (..either (/@wrap #month.May)
                                            (/@wrap #month.June))))
              (..either (..either (/@wrap #month.July)
                                  (..either (/@wrap #month.August)
                                            (/@wrap #month.September)))
                        (..either (/@wrap #month.October)
                                  (..either (/@wrap #month.November)
                                            (/@wrap #month.December)))))))

(def: #export day
  (Random Day)
  (let [(^open "/@.") ..monad]
    (..either (..either (/@wrap #day.Sunday)
                        (..either (/@wrap #day.Monday)
                                  (/@wrap #day.Tuesday)))
              (..either (..either (/@wrap #day.Wednesday)
                                  (/@wrap #day.Thursday))
                        (..either (/@wrap #day.Friday)
                                  (/@wrap #day.Saturday))))))

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

(def: pcg-32-magic Nat 6364136223846793005)

(def: #export (pcg-32 [increase seed])
  {#.doc (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) ("lux i64 +" increase) [increase] pcg-32)
     (let [rot (|> seed .i64 (i64.logic-right-shift 59))]
       (|> seed
           (i64.logic-right-shift 18)
           (i64.xor seed)
           (i64.logic-right-shift 27)
           (i64.rotate-right rot)
           .i64))]))

(def: #export (xoroshiro-128+ [s0 s1])
  {#.doc (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 (i64.xor s0 s1)]
       (xoroshiro-128+ [(|> s0
                            (i64.rotate-left 55)
                            (i64.xor s01)
                            (i64.xor (i64.left-shift 14 s01)))
                        (i64.rotate-left 36 s01)]))
     ("lux i64 +" s0 s1)]))