1 /* 2 * Copyright 2021 Google LLC 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * https://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #pragma once 18 19 #include <cstdint> 20 #include <optional> 21 #include <random> 22 23 namespace dist_proc { 24 namespace aggregation { 25 26 class RandomGenerator { 27 public: ~RandomGenerator()28 virtual ~RandomGenerator(){}; 29 30 // Returns a uniformly distributed pseudorandom integer in [0, n) 31 // where n >= 0. 32 virtual uint64_t UnbiasedUniform(uint64_t n) = 0; 33 }; 34 35 class MTRandomGenerator : public RandomGenerator { 36 public: 37 MTRandomGenerator(std::optional<uint64_t> seed = std::nullopt) { 38 if (seed.has_value()) { 39 bit_gen_ = std::mt19937(seed.value()); 40 } else { 41 std::random_device rd; 42 bit_gen_ = std::mt19937(rd()); 43 } 44 } UnbiasedUniform(uint64_t n)45 uint64_t UnbiasedUniform(uint64_t n) override { 46 std::uniform_int_distribution<uint64_t> distrib(0, n - 1); 47 return distrib(bit_gen_); 48 } 49 50 private: 51 std::mt19937 bit_gen_; 52 }; 53 54 } // namespace aggregation 55 } // namespace dist_proc 56