bktree.cc

#include <stdint.h>
#include <stdio.h>

#include <map>
#include <vector>
#include <memory>
#include <type_traits>
#include <utility>
#include <iostream>
#include <chrono> 
#include <cstdlib>
#include <shared_mutex>


static const uint8_t bitsinbyte[256] = {0,1,1,2,1,2,2,3,1,2,2,3,2,3,3,4,1,2,2,3,2,3,3,4,2,3,3,4,3,4,4,5,1,2,2,3,2,3,3,4,2,3,3,4,3,4,4,5,2,3,3,4,3,4,4,5,3,4,4,5,4,5,5,6,1,2,2,3,2,3,3,4,2,3,3,4,3,4,4,5,2,3,3,4,3,4,4,5,3,4,4,5,4,5,5,6,2,3,3,4,3,4,4,5,3,4,4,5,4,5,5,6,3,4,4,5,4,5,5,6,4,5,5,6,5,6,6,7,1,2,2,3,2,3,3,4,2,3,3,4,3,4,4,5,2,3,3,4,3,4,4,5,3,4,4,5,4,5,5,6,2,3,3,4,3,4,4,5,3,4,4,5,4,5,5,6,3,4,4,5,4,5,5,6,4,5,5,6,5,6,6,7,2,3,3,4,3,4,4,5,3,4,4,5,4,5,5,6,3,4,4,5,4,5,5,6,4,5,5,6,5,6,6,7,3,4,4,5,4,5,5,6,4,5,5,6,5,6,6,7,4,5,5,6,5,6,6,7,5,6,6,7,6,7,7,8};

int popcount(uint64_t v) {
    int s = 0;

    uint8_t p; 
    for (int i = 0; i < sizeof(i); i++) {
        p = ((uint8_t *)(&v))[i];
        s += bitsinbyte[p];
    }

    return s;
}

int hdistance(uint64_t x, uint64_t y) {
//    return __builtin_popcountll(x ^ y);
    return popcount(x ^ y);
}

class HammingDistance {
public:
    explicit HammingDistance() {}
    int operator()(const uint64_t s, const uint64_t t) const {
        return hdistance(s, t);
    }
};

template <typename T, typename Unit, typename Metric> class bktree;

template<typename T, typename Unit, typename Metric>
class bktree_node {
    friend class bktree<T, Unit, Metric>;
    typedef bktree_node<T, Unit, Metric> node_type;
    typedef std::unique_ptr<node_type> node_ptr_type;

    T value_;
    std::map<Unit, node_ptr_type> children_;
    
    bktree_node(const T &value) : value_(value) {}

    int insert(const T& value, const Metric& distance) {
        printf("%ld\n", value);
        int inserted = 0;
        Unit dist = distance(value, this->value_);
        if (dist > 0) {
            auto it = children_.find(dist);
            if (it == children_.end()) {
                children_.insert(std::make_pair(dist, node_ptr_type(new node_type(value))));
                inserted = 1;
            } else {
                inserted = it->second->insert(value, distance);
            }
        }
        return inserted;
    }

    template <class OutputIt>
    OutputIt find(const T &value, const Unit& limit, const Metric& distance, OutputIt it) const {
        Unit dist = distance(value, this->value_);
        if (dist <= limit) {
            *it++ = std::make_pair(this->value_, dist);
        }

        for (auto iter = children_.begin(); iter != children_.end(); ++iter) {
            if (abs(dist - limit) <= iter->first && abs(dist + limit) >= iter->first) {
                iter->second->find(value, limit, distance, it);
            }
        }

        return it;
    }
};

template<typename T, typename Unit, typename Metric>
class bktree {
private:
    typedef typename bktree_node<T, Unit, Metric>::node_type node_type;
    typedef typename bktree_node<T, Unit, Metric>::node_ptr_type node_ptr_type;

    node_ptr_type head_;
    const Metric distance_;
    size_t size_; 

    mutable std::shared_mutex mutex_;
public:
    bktree(const Metric& distance = Metric()) : head_(nullptr), distance_(distance), size_(0L) {
        static_assert(std::is_integral<Unit>::value, "Integral unit type required.");
    }

    int insert(const T& value) {
        std::unique_lock lock(mutex_);

        int inserted = 0;
        if (head_ == nullptr) {
            head_ = node_ptr_type(new node_type(value));
            size_ = 1;
            inserted = 1;
        } else if (head_->insert(value, distance_)) {
            ++size_;
            inserted = 1;
        }
        
        return inserted;
    }

    template<class OutputIt>
    OutputIt find(const T& value, const Unit& limit, OutputIt it) const {
        std::shared_lock lock(mutex_);
        return head_->find(value, limit, distance_, it);
    }

    size_t size() const {
        return size_;
    }
};

void test_insert(bktree<uint64_t, int, HammingDistance>& tree, int range) {
    auto start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < range; i++) {
        tree.insert(i);
    }
    auto stop = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(stop - start); 

    std::cout << "Time taken by test_insert function: "
              << duration.count() << " milliseconds" << std::endl;
}

void test_find(bktree<uint64_t, int, HammingDistance>& tree, std::vector<std::pair<uint64_t, int>>& results) {
    auto start = std::chrono::high_resolution_clock::now();
    tree.find(10000001, 3, std::back_inserter(results));

    auto stop = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(stop - start); 

    std::cout << "Time taken by test_find function: "
              << duration.count() << " milliseconds" << std::endl;
}

int main(int argc, char *argv[]) {
    bktree<uint64_t, int, HammingDistance> tree;
    std::vector<std::pair<uint64_t, int>> results;

    test_insert(tree, 10000000);
    test_find(tree, results);

//    for (auto it = results.begin(); it != results.end(); ++it) {
//        std::cout << it->first << " (distance " << it->second << ")\n";
//    }

    for (int i = 0; i < 100000000; i++) {
    }
    return 0;
}