#include <cuco/static_multiset.cuh>

#include <thrust/device_vector.h>

#include <thrust/functional.h>

#include <thrust/logical.h>

#include <thrust/sequence.h>

#include <iostream>

#include <limits>

int main(void)

{

using key_type = int;

// Empty slots are represented by reserved "sentinel" values. These values should be selected such

// that they never occur in your input data.

key_type constexpr empty_key_sentinel = -1;

// Number of keys to be inserted

std::size_t constexpr num_keys = 50'000;

// Compute capacity based on a 50% load factor

auto constexpr load_factor = 0.5;

std::size_t const capacity = std::ceil(num_keys / load_factor);

// Constructs a set with at least `capacity` slots using -1 as the empty keys sentinel.

cuco::static_multiset<key_type> multiset{capacity, cuco::empty_key{empty_key_sentinel}};

// Create a sequence of keys {0, 1, 2, .., i}

// We're going to insert each key twice so we only need 'num_keys / 2' distinct keys.

thrust::device_vector<key_type> keys(num_keys / 2);

thrust::sequence(keys.begin(), keys.end(), 0);

// Inserts all keys into the hash set

multiset.insert(keys.begin(), keys.end());

// Insert the same set of keys again, so each distinct key should occur twice in the multiset

multiset.insert(keys.begin(), keys.end());

// Counts the occurrences of matching keys contained in the multiset.

std::size_t const counted_output_size = multiset.count(keys.begin(), keys.end());

// Storage for result

thrust::device_vector<key_type> output_probes(counted_output_size);

thrust::device_vector<key_type> output_matches(counted_output_size);

// Retrieve all matching keys

auto const [output_probes_end, _] =

multiset.retrieve(keys.begin(), keys.end(), output_probes.begin(), output_matches.begin());

std::size_t const retrieved_output_size = output_probes_end - output_probes.begin();

if ((retrieved_output_size == counted_output_size) and (retrieved_output_size == num_keys)) {

std::cout << "Success! Found all keys.\n";

} else {

std::cout << "Fail! Something went wrong.\n";

}

return 0;

}