teems/Enclave/sort.tcc

// set_key for each kind of key we sort on

template<>
inline void set_key<UidKey>(UidKey *key, const uint8_t *item, uint32_t index)
{
    key->uid_index = (uint64_t(*(const uint32_t *)item) << 32) + index;
}

template<>
inline void set_key<UidPriorityKey>(UidPriorityKey *key, const uint8_t *item, uint32_t index)
{
    key->uid_priority = (uint64_t(*(const uint32_t *)item) << 32) +
            (*(const uint32_t *)(item+4));
    key->idx = index;
}

template<>
inline void set_key<NidPriorityKey>(NidPriorityKey *key, const uint8_t *item, uint32_t index)
{
    constexpr uint32_t nid_mask = (~((1<<DEST_UID_BITS)-1));
    key->nid_priority = (uint64_t((*(const uint32_t *)item)&nid_mask) << 32) +
            (*(const uint32_t *)(item+4));
    key->idx = index;
}

// compare_keys for each kind of key we sort on.  Note that it must not
// be possible for any of these functions to return 0.  These functions
// must also be oblivious.  Return a positive (32-bit signed) int if *a
// is larger than *b, or a negative (32-bit signed) int otherwise.

template<>
int compare_keys<UidKey>(const void* a, const void* b)
{
    bool alarge = (((const UidKey*)a)->uid_index >
        ((const UidKey *)b)->uid_index);
    return oselect_uint32_t(-1, 1, alarge);
}

template<>
int compare_keys<UidPriorityKey>(const void* a, const void* b)
{
    uint64_t aup = ((const UidPriorityKey*)a)->uid_priority;
    uint64_t bup = ((const UidPriorityKey*)b)->uid_priority;
    uint32_t aidx = ((const UidPriorityKey*)a)->idx;
    uint32_t bidx = ((const UidPriorityKey*)b)->idx;
    bool auplarge = (aup > bup);
    bool aupeq = (aup == bup);
    bool aidxlarge = (aidx > bidx);
    bool alarge = auplarge | (aupeq & aidxlarge);
    return oselect_uint32_t(-1, 1, alarge);
}

template<>
int compare_keys<NidPriorityKey>(const void* a, const void* b)
{
    uint64_t anp = ((const NidPriorityKey*)a)->nid_priority;
    uint64_t bnp = ((const NidPriorityKey*)b)->nid_priority;
    uint32_t aidx = ((const NidPriorityKey*)a)->idx;
    uint32_t bidx = ((const NidPriorityKey*)b)->idx;
    bool anplarge = (anp > bnp);
    bool anpeq = (anp == bnp);
    bool aidxlarge = (aidx > bidx);
    bool alarge = anplarge | (anpeq & aidxlarge);
    return oselect_uint32_t(-1, 1, alarge);
}

// Perform the sort using up to nthreads threads.  The items to sort are
// byte arrays of size msg_size.  The key is the 10-bit storage server
// id concatenated with the 22-bit uid at the storage server.
template<typename T>
void sort_mtobliv(threadid_t nthreads, uint8_t* items, uint16_t msg_size,
    uint32_t Nr, uint32_t Na,
    // the arguments to the callback are items, the sorted indices, and
    // the number of non-padding items
    std::function<void(const uint8_t*, const T*, uint32_t Nr)> cb)
{
    // Shuffle the items
    uint32_t Nw = shuffle_mtobliv(nthreads, items, msg_size, Nr, Na);

    // Create the indices
    T *idx = new T[Nr];
    T *nextidx = idx;
    for (uint32_t i=0; i<Nw; ++i) {
        uint64_t padding = (*(uint32_t*)(items+msg_size*i));
        if (padding != uint32_t(-1)) {
            set_key<T>(nextidx, items+msg_size*i, i);
            ++nextidx;
        }
    }
    if (nextidx != idx + Nr) {
        printf("Found %u non-padding items, expected %u\n",
            nextidx-idx, Nr);
        assert(nextidx == idx + Nr);
    }
    // Sort the keys and indices
    T *backingidx = new T[Nr];
    bool whichbuf = mtmergesort<T>(idx, Nr, backingidx, nthreads);
    T *sortedidx = whichbuf ? backingidx : idx;
    cb(items, sortedidx, Nr);

    delete[] idx;
    delete[] backingidx;
}