Line data    Source code

       1             : //===--- Random.h - Utilities for random sampling -------------------------===//
       2             : //
       3             : //                     The LLVM Compiler Infrastructure
       4             : //
       5             : // This file is distributed under the University of Illinois Open Source
       6             : // License. See LICENSE.TXT for details.
       7             : //
       8             : //===----------------------------------------------------------------------===//
       9             : //
      10             : // Utilities for random sampling.
      11             : //
      12             : //===----------------------------------------------------------------------===//
      13             : 
      14             : #ifndef LLVM_FUZZMUTATE_RANDOM_H
      15             : #define LLVM_FUZZMUTATE_RANDOM_H
      16             : 
      17             : #include <random>
      18             : #include "llvm/Support/raw_ostream.h"
      19             : namespace llvm {
      20             : 
      21             : /// Return a uniformly distributed random value between \c Min and \c Max
      22             : template <typename T, typename GenT> T uniform(GenT &Gen, T Min, T Max) {
      23      150005 :   return std::uniform_int_distribution<T>(Min, Max)(Gen);
      24             : }
      25             : 
      26             : /// Return a uniformly distributed random value of type \c T
      27             : template <typename T, typename GenT> T uniform(GenT &Gen) {
      28             :   return uniform<T>(Gen, std::numeric_limits<T>::min(),
      29             :                     std::numeric_limits<T>::max());
      30             : }
      31             : 
      32             : /// Randomly selects an item by sampling into a set with an unknown number of
      33             : /// elements, which may each be weighted to be more likely choices.
      34           1 : template <typename T, typename GenT> class ReservoirSampler {
      35             :   GenT &RandGen;
      36             :   typename std::remove_const<T>::type Selection = {};
      37             :   uint64_t TotalWeight = 0;
      38             : 
      39             : public:
      40       15137 :   ReservoirSampler(GenT &RandGen) : RandGen(RandGen) {}
      41             : 
      42           0 :   uint64_t totalWeight() const { return TotalWeight; }
      43           1 :   bool isEmpty() const { return TotalWeight == 0; }
      44             : 
      45             :   const T &getSelection() const {
      46             :     assert(!isEmpty() && "Nothing selected");
      47             :     return Selection;
      48             :   }
      49             : 
      50          87 :   explicit operator bool() const { return !isEmpty();}
      51             :   const T &operator*() const { return getSelection(); }
      52             : 
      53             :   /// Sample each item in \c Items with unit weight
      54         171 :   template <typename RangeT> ReservoirSampler &sample(RangeT &&Items) {
      55      165596 :     for (auto &I : Items)
      56      150341 :       sample(I, 1);
      57         171 :     return *this;
      58             :   }
      59          88 : 
      60         132 :   /// Sample a single item with the given weight.
      61      150046 :   ReservoirSampler &sample(const T &Item, uint64_t Weight) {
      62      150090 :     if (!Weight)
      63             :       // If the weight is zero, do nothing.
      64          82 :       return *this;
      65      150103 :     TotalWeight += Weight;
      66          20 :     // Consider switching from the current element to this one.
      67      150083 :     if (uniform<uint64_t>(RandGen, 1, TotalWeight) <= Weight)
      68       44054 :       Selection = Item;
      69           1 :     return *this;
      70           2 :   }
      71         440 : };
      72         440 : 
      73             : template <typename GenT, typename RangeT,
      74             :           typename ElT = typename std::remove_reference<
      75         439 :               decltype(*std::begin(std::declval<RangeT>()))>::type>
      76          77 : ReservoirSampler<ElT, GenT> makeSampler(GenT &RandGen, RangeT &&Items) {
      77         516 :   ReservoirSampler<ElT, GenT> RS(RandGen);
      78         271 :   RS.sample(Items);
      79             :   return RS;
      80          77 : }
      81          31 : 
      82         108 : template <typename GenT, typename T>
      83          43 : ReservoirSampler<T, GenT> makeSampler(GenT &RandGen, const T &Item,
      84             :                                       uint64_t Weight) {
      85          31 :   ReservoirSampler<T, GenT> RS(RandGen);
      86          26 :   RS.sample(Item, Weight);
      87          55 :   return RS;
      88          19 : }
      89             : 
      90          24 : template <typename T, typename GenT>
      91         286 : ReservoirSampler<T, GenT> makeSampler(GenT &RandGen) {
      92         310 :   return ReservoirSampler<T, GenT>(RandGen);
      93             : }
      94             : 
      95         286 : } // End llvm namespace
      96           1 : 
      97         287 : #endif // LLVM_FUZZMUTATE_RANDOM_H