/* $NetBSD: subr_thmap.c,v 1.15 2023/10/17 11:57:20 riastradh Exp $ */

/* $NetBSD: subr_thmap.c,v 1.15 2023/10/17 11:57:20 riastradh Exp $ */

/*-
* Copyright (c) 2018 Mindaugas Rasiukevicius <rmind at noxt eu>
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions
* are met:
* 1. Redistributions of source code must retain the above copyright
* notice, this list of conditions and the following disclaimer.
* 2. Redistributions in binary form must reproduce the above copyright
* notice, this list of conditions and the following disclaimer in the
* documentation and/or other materials provided with the distribution.
*
* THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
* ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
* OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
* HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
* LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
* OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
* SUCH DAMAGE.
*
* Upstream: https://github.com/rmind/thmap/
*/

/*
* Concurrent trie-hash map.
*
* The data structure is conceptually a radix trie on hashed keys.
* Keys are hashed using a 32-bit function. The root level is a special
* case: it is managed using the compare-and-swap (CAS) atomic operation
* and has a fanout of 64. The subsequent levels are constructed using
* intermediate nodes with a fanout of 16 (using 4 bits). As more levels
* are created, more blocks of the 32-bit hash value might be generated
* by incrementing the seed parameter of the hash function.
*
* Concurrency
*
* - READERS: Descending is simply walking through the slot values of
* the intermediate nodes. It is lock-free as there is no intermediate
* state: the slot is either empty or has a pointer to the child node.
* The main assumptions here are the following:
*
* i) modifications must preserve consistency with the respect to the
* readers i.e. the readers can only see the valid node values;
*
* ii) any invalid view must "fail" the reads, e.g. by making them
* re-try from the root; this is a case for deletions and is achieved
* using the NODE_DELETED flag.
*
* iii) the node destruction must be synchronized with the readers,
* e.g. by using the Epoch-based reclamation or other techniques.
*
* - WRITERS AND LOCKING: Each intermediate node has a spin-lock (which
* is implemented using the NODE_LOCKED bit) -- it provides mutual
* exclusion amongst concurrent writers. The lock order for the nodes
* is "bottom-up" i.e. they are locked as we ascend the trie. A key
* constraint here is that parent pointer never changes.
*
* - DELETES: In addition to writer's locking, the deletion keeps the
* intermediate nodes in a valid state and sets the NODE_DELETED flag,
* to indicate that the readers must re-start the walk from the root.
* As the levels are collapsed, NODE_DELETED gets propagated up-tree.
* The leaf nodes just stay as-is until they are reclaimed.
*
* - ROOT LEVEL: The root level is a special case, as it is implemented
* as an array (rather than intermediate node). The root-level slot can
* only be set using CAS and it can only be set to a valid intermediate
* node. The root-level slot can only be cleared when the node it points
* at becomes empty, is locked and marked as NODE_DELETED (this causes
* the insert/delete operations to re-try until the slot is set to NULL).
*
* References:
*
* W. Litwin, 1981, Trie Hashing.
* Proceedings of the 1981 ACM SIGMOD, p. 19-29
* https://dl.acm.org/citation.cfm?id=582322
*
* P. L. Lehman and S. B. Yao.
* Efficient locking for concurrent operations on B-trees.
* ACM TODS, 6(4):650-670, 1981
* https://www.csd.uoc.gr/~hy460/pdf/p650-lehman.pdf
*/

#ifdef _KERNEL
#include <sys/cdefs.h>
#include <sys/param.h>
#include <sys/types.h>
#include <sys/thmap.h>
#include <sys/kmem.h>
#include <sys/lock.h>
#include <sys/atomic.h>
#include <sys/hash.h>
#include <sys/cprng.h>
#define THMAP_RCSID(a) __KERNEL_RCSID(0, a)
#else
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <stddef.h>
#include <inttypes.h>
#include <string.h>
#include <limits.h>
#define THMAP_RCSID(a) __RCSID(a)

#include "thmap.h"
#include "utils.h"
#endif

THMAP_RCSID("$NetBSD: subr_thmap.c,v 1.15 2023/10/17 11:57:20 riastradh Exp $");

#include <crypto/blake2/blake2s.h>

/*
* NetBSD kernel wrappers
*/
#ifdef _KERNEL
#define ASSERT KASSERT
#define atomic_thread_fence(x) membar_release() /* only used for release order */
#define atomic_compare_exchange_weak_explicit_32(p, e, n, m1, m2) \
(atomic_cas_32((p), *(e), (n)) == *(e))
#define atomic_compare_exchange_weak_explicit_ptr(p, e, n, m1, m2) \
(atomic_cas_ptr((p), *(void **)(e), (void *)(n)) == *(void **)(e))
#define atomic_exchange_explicit(o, n, m1) atomic_swap_ptr((o), (n))
#define murmurhash3 murmurhash2
#endif

/*
* The root level fanout is 64 (indexed by the last 6 bits of the hash
* value XORed with the length). Each subsequent level, represented by
* intermediate nodes, has a fanout of 16 (using 4 bits).
*
* The hash function produces 32-bit values.
*/

#define HASHVAL_SEEDLEN (16)
#define HASHVAL_BITS (32)
#define HASHVAL_MOD (HASHVAL_BITS - 1)
#define HASHVAL_SHIFT (5)

#define ROOT_BITS (6)
#define ROOT_SIZE (1 << ROOT_BITS)
#define ROOT_MASK (ROOT_SIZE - 1)
#define ROOT_MSBITS (HASHVAL_BITS - ROOT_BITS)

#define LEVEL_BITS (4)
#define LEVEL_SIZE (1 << LEVEL_BITS)
#define LEVEL_MASK (LEVEL_SIZE - 1)

/*
* Instead of raw pointers, we use offsets from the base address.
* This accommodates the use of this data structure in shared memory,
* where mappings can be in different address spaces.
*
* The pointers must be aligned, since pointer tagging is used to
* differentiate the intermediate nodes from leaves. We reserve the
* least significant bit.
*/
typedef uintptr_t thmap_ptr_t;
typedef uintptr_t atomic_thmap_ptr_t; // C11 _Atomic

#define THMAP_NULL ((thmap_ptr_t)0)

#define THMAP_LEAF_BIT (0x1)

#define THMAP_ALIGNED_P(p) (((uintptr_t)(p) & 3) == 0)
#define THMAP_ALIGN(p) ((uintptr_t)(p) & ~(uintptr_t)3)
#define THMAP_INODE_P(p) (((uintptr_t)(p) & THMAP_LEAF_BIT) == 0)

#define THMAP_GETPTR(th, p) ((void *)((th)->baseptr + (uintptr_t)(p)))
#define THMAP_GETOFF(th, p) ((thmap_ptr_t)((uintptr_t)(p) - (th)->baseptr))
#define THMAP_NODE(th, p) THMAP_GETPTR(th, THMAP_ALIGN(p))

/*
* State field.
*/

#define NODE_LOCKED (1U << 31) // lock (writers)
#define NODE_DELETED (1U << 30) // node deleted
#define NODE_COUNT(s) ((s) & 0x3fffffff) // slot count mask

/*
* There are two types of nodes:
* - Intermediate nodes -- arrays pointing to another level or a leaf;
* - Leaves, which store a key-value pair.
*/

typedef struct {
uint32_t state; // C11 _Atomic
thmap_ptr_t parent;
atomic_thmap_ptr_t slots[LEVEL_SIZE];
} thmap_inode_t;

#define THMAP_INODE_LEN sizeof(thmap_inode_t)

typedef struct {
thmap_ptr_t key;
size_t len;
void * val;
} thmap_leaf_t;

typedef struct {
const uint8_t * seed; // secret seed
unsigned rslot; // root-level slot index
unsigned level; // current level in the tree
unsigned hashidx; // current hash index (block of bits)
uint32_t hashval; // current hash value
} thmap_query_t;

union thmap_align {
void * p;
uint64_t v;
};

typedef struct thmap_gc thmap_gc_t;
struct thmap_gc {
size_t len;
thmap_gc_t * next;
char data[] __aligned(sizeof(union thmap_align));
};

#define THMAP_ROOT_LEN (sizeof(thmap_ptr_t) * ROOT_SIZE)

struct thmap {
uintptr_t baseptr;
atomic_thmap_ptr_t * root;
unsigned flags;
const thmap_ops_t * ops;
thmap_gc_t * gc_list; // C11 _Atomic
uint8_t seed[HASHVAL_SEEDLEN];
};

static void stage_mem_gc(thmap_t *, uintptr_t, size_t);

/*
* A few low-level helper routines.
*/

static uintptr_t
alloc_wrapper(size_t len)
{
return (uintptr_t)kmem_intr_alloc(len, KM_NOSLEEP);
}

static void
free_wrapper(uintptr_t addr, size_t len)
{
kmem_intr_free((void *)addr, len);
}

static const thmap_ops_t thmap_default_ops = {
.alloc = alloc_wrapper,
.free = free_wrapper
};

static uintptr_t
gc_alloc(const thmap_t *thmap, size_t len)
{
const size_t alloclen = offsetof(struct thmap_gc, data[len]);
const uintptr_t gcaddr = thmap->ops->alloc(alloclen);

if (!gcaddr)
return 0;

thmap_gc_t *const gc = THMAP_GETPTR(thmap, gcaddr);
gc->len = len;
return THMAP_GETOFF(thmap, &gc->data[0]);
}

static void
gc_free(const thmap_t *thmap, uintptr_t addr, size_t len)
{
const size_t alloclen = offsetof(struct thmap_gc, data[len]);
char *const ptr = THMAP_GETPTR(thmap, addr);
thmap_gc_t *const gc = container_of(ptr, struct thmap_gc, data[0]);
const uintptr_t gcaddr = THMAP_GETOFF(thmap, gc);

KASSERTMSG(gc->len == len, "thmap=%p ops=%p addr=%p len=%zu"
" gc=%p gc->len=%zu",
thmap, thmap->ops, (void *)addr, len, gc, gc->len);
thmap->ops->free(gcaddr, alloclen);
}

/*
* NODE LOCKING.
*/

static inline bool __diagused
node_locked_p(thmap_inode_t *node)
{
return (atomic_load_relaxed(&node->state) & NODE_LOCKED) != 0;
}

static void
lock_node(thmap_inode_t *node)
{
unsigned bcount = SPINLOCK_BACKOFF_MIN;
uint32_t s;
again:
s = atomic_load_relaxed(&node->state);
if (s & NODE_LOCKED) {
SPINLOCK_BACKOFF(bcount);
goto again;
}
/* Acquire from prior release in unlock_node.() */
if (!atomic_compare_exchange_weak_explicit_32(&node->state,
&s, s | NODE_LOCKED, memory_order_acquire, memory_order_relaxed)) {
bcount = SPINLOCK_BACKOFF_MIN;
goto again;
}
}

static void
unlock_node(thmap_inode_t *node)
{
uint32_t s = atomic_load_relaxed(&node->state) & ~NODE_LOCKED;

ASSERT(node_locked_p(node));
/* Release to subsequent acquire in lock_node(). */
atomic_store_release(&node->state, s);
}

/*
* HASH VALUE AND KEY OPERATIONS.
*/

static inline uint32_t
hash(const uint8_t seed[static HASHVAL_SEEDLEN], const void *key, size_t len,
uint32_t level)
{
struct blake2s B;
uint32_t h;

if (level == 0)
return murmurhash3(key, len, 0);

/*
* Byte order is not significant here because this is
* intentionally secret and independent for each thmap.
*
* XXX We get 32 bytes of output at a time; we could march
* through them sequentially rather than throwing away 28 bytes
* and recomputing BLAKE2 each time. But the number of
* iterations ought to be geometric in the collision
* probability at each level which should be very small anyway.
*/
blake2s_init(&B, sizeof h, seed, HASHVAL_SEEDLEN);
blake2s_update(&B, &level, sizeof level);
blake2s_update(&B, key, len);
blake2s_final(&B, &h);

return h;
}

static inline void
hashval_init(thmap_query_t *query, const uint8_t seed[static HASHVAL_SEEDLEN],
const void * restrict key, size_t len)
{
const uint32_t hashval = hash(seed, key, len, 0);

query->seed = seed;
query->rslot = ((hashval >> ROOT_MSBITS) ^ len) & ROOT_MASK;
query->level = 0;
query->hashval = hashval;
query->hashidx = 0;
}

/*
* hashval_getslot: given the key, compute the hash (if not already cached)
* and return the offset for the current level.
*/
static unsigned
hashval_getslot(thmap_query_t *query, const void * restrict key, size_t len)
{
const unsigned offset = query->level * LEVEL_BITS;
const unsigned shift = offset & HASHVAL_MOD;
const unsigned i = offset >> HASHVAL_SHIFT;

if (query->hashidx != i) {
/* Generate a hash value for a required range. */
query->hashval = hash(query->seed, key, len, i);
query->hashidx = i;
}
return (query->hashval >> shift) & LEVEL_MASK;
}

static unsigned
hashval_getleafslot(const thmap_t *thmap,
const thmap_leaf_t *leaf, unsigned level)
{
const void *key = THMAP_GETPTR(thmap, leaf->key);
const unsigned offset = level * LEVEL_BITS;
const unsigned shift = offset & HASHVAL_MOD;
const unsigned i = offset >> HASHVAL_SHIFT;

return (hash(thmap->seed, key, leaf->len, i) >> shift) & LEVEL_MASK;
}

static inline unsigned
hashval_getl0slot(const thmap_t *thmap, const thmap_query_t *query,
const thmap_leaf_t *leaf)
{
if (__predict_true(query->hashidx == 0)) {
return query->hashval & LEVEL_MASK;
}
return hashval_getleafslot(thmap, leaf, 0);
}

static bool
key_cmp_p(const thmap_t *thmap, const thmap_leaf_t *leaf,
const void * restrict key, size_t len)
{
const void *leafkey = THMAP_GETPTR(thmap, leaf->key);
return len == leaf->len && memcmp(key, leafkey, len) == 0;
}

/*
* INTER-NODE OPERATIONS.
*/

static thmap_inode_t *
node_create(thmap_t *thmap, thmap_inode_t *parent)
{
thmap_inode_t *node;
uintptr_t p;

p = gc_alloc(thmap, THMAP_INODE_LEN);
if (!p) {
return NULL;
}
node = THMAP_GETPTR(thmap, p);
ASSERT(THMAP_ALIGNED_P(node));

memset(node, 0, THMAP_INODE_LEN);
if (parent) {
/* Not yet published, no need for ordering. */
atomic_store_relaxed(&node->state, NODE_LOCKED);
node->parent = THMAP_GETOFF(thmap, parent);
}
return node;
}

static void
node_insert(thmap_inode_t *node, unsigned slot, thmap_ptr_t child)
{
ASSERT(node_locked_p(node) || node->parent == THMAP_NULL);
ASSERT((atomic_load_relaxed(&node->state) & NODE_DELETED) == 0);
ASSERT(atomic_load_relaxed(&node->slots[slot]) == THMAP_NULL);

ASSERT(NODE_COUNT(atomic_load_relaxed(&node->state)) < LEVEL_SIZE);

/*
* If node is public already, caller is responsible for issuing
* release fence; if node is not public, no ordering is needed.
* Hence relaxed ordering.
*/
atomic_store_relaxed(&node->slots[slot], child);
atomic_store_relaxed(&node->state,
atomic_load_relaxed(&node->state) + 1);
}

static void
node_remove(thmap_inode_t *node, unsigned slot)
{
ASSERT(node_locked_p(node));
ASSERT((atomic_load_relaxed(&node->state) & NODE_DELETED) == 0);
ASSERT(atomic_load_relaxed(&node->slots[slot]) != THMAP_NULL);

ASSERT(NODE_COUNT(atomic_load_relaxed(&node->state)) > 0);
ASSERT(NODE_COUNT(atomic_load_relaxed(&node->state)) <= LEVEL_SIZE);

/* Element will be GC-ed later; no need for ordering here. */
atomic_store_relaxed(&node->slots[slot], THMAP_NULL);
atomic_store_relaxed(&node->state,
atomic_load_relaxed(&node->state) - 1);
}

/*
* LEAF OPERATIONS.
*/

static thmap_leaf_t *
leaf_create(const thmap_t *thmap, const void *key, size_t len, void *val)
{
thmap_leaf_t *leaf;
uintptr_t leaf_off, key_off;

leaf_off = gc_alloc(thmap, sizeof(thmap_leaf_t));
if (!leaf_off) {
return NULL;
}
leaf = THMAP_GETPTR(thmap, leaf_off);
ASSERT(THMAP_ALIGNED_P(leaf));

if ((thmap->flags & THMAP_NOCOPY) == 0) {
/*
* Copy the key.
*/
key_off = gc_alloc(thmap, len);
if (!key_off) {
gc_free(thmap, leaf_off, sizeof(thmap_leaf_t));
return NULL;
}
memcpy(THMAP_GETPTR(thmap, key_off), key, len);
leaf->key = key_off;
} else {
/* Otherwise, we use a reference. */
leaf->key = (uintptr_t)key;
}
leaf->len = len;
leaf->val = val;
return leaf;
}

static void
leaf_free(const thmap_t *thmap, thmap_leaf_t *leaf)
{
if ((thmap->flags & THMAP_NOCOPY) == 0) {
gc_free(thmap, leaf->key, leaf->len);
}
gc_free(thmap, THMAP_GETOFF(thmap, leaf), sizeof(thmap_leaf_t));
}

static thmap_leaf_t *
get_leaf(const thmap_t *thmap, thmap_inode_t *parent, unsigned slot)
{
thmap_ptr_t node;

/* Consume from prior release in thmap_put(). */
node = atomic_load_consume(&parent->slots[slot]);
if (THMAP_INODE_P(node)) {
return NULL;
}
return THMAP_NODE(thmap, node);
}

/*
* ROOT OPERATIONS.
*/

/*
* root_try_put: Try to set a root pointer at query->rslot.
*
* => Implies release operation on success.
* => Implies no ordering on failure.
*/
static inline int
root_try_put(thmap_t *thmap, const thmap_query_t *query, thmap_leaf_t *leaf)
{
thmap_ptr_t expected;
const unsigned i = query->rslot;
thmap_inode_t *node;
thmap_ptr_t nptr;
unsigned slot;

/*
* Must pre-check first. No ordering required because we will
* check again before taking any actions, and start over if
* this changes from null.
*/
if (atomic_load_relaxed(&thmap->root[i])) {
return EEXIST;
}

/*
* Create an intermediate node. Since there is no parent set,
* it will be created unlocked and the CAS operation will
* release it to readers.
*/
node = node_create(thmap, NULL);
if (__predict_false(node == NULL)) {
return ENOMEM;
}
slot = hashval_getl0slot(thmap, query, leaf);
node_insert(node, slot, THMAP_GETOFF(thmap, leaf) | THMAP_LEAF_BIT);
nptr = THMAP_GETOFF(thmap, node);
again:
if (atomic_load_relaxed(&thmap->root[i])) {
gc_free(thmap, nptr, THMAP_INODE_LEN);
return EEXIST;
}
/* Release to subsequent consume in find_edge_node(). */
expected = THMAP_NULL;
if (!atomic_compare_exchange_weak_explicit_ptr(&thmap->root[i], &expected,
nptr, memory_order_release, memory_order_relaxed)) {
goto again;
}
return 0;
}

/*
* find_edge_node: given the hash, traverse the tree to find the edge node.
*
* => Returns an aligned (clean) pointer to the parent node.
* => Returns the slot number and sets current level.
*/
static thmap_inode_t *
find_edge_node(const thmap_t *thmap, thmap_query_t *query,
const void * restrict key, size_t len, unsigned *slot)
{
thmap_ptr_t root_slot;
thmap_inode_t *parent;
thmap_ptr_t node;
unsigned off;

ASSERT(query->level == 0);

/* Consume from prior release in root_try_put(). */
root_slot = atomic_load_consume(&thmap->root[query->rslot]);
parent = THMAP_NODE(thmap, root_slot);
if (!parent) {
return NULL;
}
descend:
off = hashval_getslot(query, key, len);
/* Consume from prior release in thmap_put(). */
node = atomic_load_consume(&parent->slots[off]);

/* Descend the tree until we find a leaf or empty slot. */
if (node && THMAP_INODE_P(node)) {
parent = THMAP_NODE(thmap, node);
query->level++;
goto descend;
}
/*
* NODE_DELETED does not become stale until GC runs, which
* cannot happen while we are in the middle of an operation,
* hence relaxed ordering.
*/
if (atomic_load_relaxed(&parent->state) & NODE_DELETED) {
return NULL;
}
*slot = off;
return parent;
}

/*
* find_edge_node_locked: traverse the tree, like find_edge_node(),
* but attempt to lock the edge node.
*
* => Returns NULL if the deleted node is found. This indicates that
* the caller must re-try from the root, as the root slot might have
* changed too.
*/
static thmap_inode_t *
find_edge_node_locked(const thmap_t *thmap, thmap_query_t *query,
const void * restrict key, size_t len, unsigned *slot)
{
thmap_inode_t *node;
thmap_ptr_t target;
retry:
/*
* Find the edge node and lock it! Re-check the state since
* the tree might change by the time we acquire the lock.
*/
node = find_edge_node(thmap, query, key, len, slot);
if (!node) {
/* The root slot is empty -- let the caller decide. */
query->level = 0;
return NULL;
}
lock_node(node);
if (__predict_false(atomic_load_relaxed(&node->state) & NODE_DELETED)) {
/*
* The node has been deleted. The tree might have a new
* shape now, therefore we must re-start from the root.
*/
unlock_node(node);
query->level = 0;
return NULL;
}
target = atomic_load_relaxed(&node->slots[*slot]);
if (__predict_false(target && THMAP_INODE_P(target))) {
/*
* The target slot has been changed and it is now an
* intermediate node. Re-start from the top internode.
*/
unlock_node(node);
query->level = 0;
goto retry;
}
return node;
}

/*
* thmap_get: lookup a value given the key.
*/
void *
thmap_get(thmap_t *thmap, const void *key, size_t len)
{
thmap_query_t query;
thmap_inode_t *parent;
thmap_leaf_t *leaf;
unsigned slot;

hashval_init(&query, thmap->seed, key, len);
parent = find_edge_node(thmap, &query, key, len, &slot);
if (!parent) {
return NULL;
}
leaf = get_leaf(thmap, parent, slot);
if (!leaf) {
return NULL;
}
if (!key_cmp_p(thmap, leaf, key, len)) {
return NULL;
}
return leaf->val;
}

/*
* thmap_put: insert a value given the key.
*
* => If the key is already present, return the associated value.
* => Otherwise, on successful insert, return the given value.
*/
void *
thmap_put(thmap_t *thmap, const void *key, size_t len, void *val)
{
thmap_query_t query;
thmap_leaf_t *leaf, *other;
thmap_inode_t *parent, *child;
unsigned slot, other_slot;
thmap_ptr_t target;

/*
* First, pre-allocate and initialize the leaf node.
*/
leaf = leaf_create(thmap, key, len, val);
if (__predict_false(!leaf)) {
return NULL;
}
hashval_init(&query, thmap->seed, key, len);
retry:
/*
* Try to insert into the root first, if its slot is empty.
*/
switch (root_try_put(thmap, &query, leaf)) {
case 0:
/* Success: the leaf was inserted; no locking involved. */
return val;
case EEXIST:
break;
case ENOMEM:
return NULL;
default:
__unreachable();
}

/*
* Release node via store in node_insert (*) to subsequent
* consume in get_leaf() or find_edge_node().
*/
atomic_thread_fence(memory_order_release);

/*
* Find the edge node and the target slot.
*/
parent = find_edge_node_locked(thmap, &query, key, len, &slot);
if (!parent) {
goto retry;
}
target = atomic_load_relaxed(&parent->slots[slot]); // tagged offset
if (THMAP_INODE_P(target)) {
/*
* Empty slot: simply insert the new leaf. The release
* fence is already issued for us.
*/
target = THMAP_GETOFF(thmap, leaf) | THMAP_LEAF_BIT;
node_insert(parent, slot, target); /* (*) */
goto out;
}

/*
* Collision or duplicate.
*/
other = THMAP_NODE(thmap, target);
if (key_cmp_p(thmap, other, key, len)) {
/*
* Duplicate. Free the pre-allocated leaf and
* return the present value.
*/
leaf_free(thmap, leaf);
val = other->val;
goto out;
}
descend:
/*
* Collision -- expand the tree. Create an intermediate node
* which will be locked (NODE_LOCKED) for us. At this point,
* we advance to the next level.
*/
child = node_create(thmap, parent);
if (__predict_false(!child)) {
leaf_free(thmap, leaf);
val = NULL;
goto out;
}
query.level++;

/*
* Insert the other (colliding) leaf first. The new child is
* not yet published, so memory order is relaxed.
*/
other_slot = hashval_getleafslot(thmap, other, query.level);
target = THMAP_GETOFF(thmap, other) | THMAP_LEAF_BIT;
node_insert(child, other_slot, target);

/*
* Insert the intermediate node into the parent node.
* It becomes the new parent for the our new leaf.
*
* Ensure that stores to the child (and leaf) reach global
* visibility before it gets inserted to the parent, as
* consumed by get_leaf() or find_edge_node().
*/
atomic_store_release(&parent->slots[slot], THMAP_GETOFF(thmap, child));

unlock_node(parent);
ASSERT(node_locked_p(child));
parent = child;

/*
* Get the new slot and check for another collision
* at the next level.
*/
slot = hashval_getslot(&query, key, len);
if (slot == other_slot) {
/* Another collision -- descend and expand again. */
goto descend;
}

/*
* Insert our new leaf once we expanded enough. The release
* fence is already issued for us.
*/
target = THMAP_GETOFF(thmap, leaf) | THMAP_LEAF_BIT;
node_insert(parent, slot, target); /* (*) */
out:
unlock_node(parent);
return val;
}

/*
* thmap_del: remove the entry given the key.
*/
void *
thmap_del(thmap_t *thmap, const void *key, size_t len)
{
thmap_query_t query;
thmap_leaf_t *leaf;
thmap_inode_t *parent;
unsigned slot;
void *val;

hashval_init(&query, thmap->seed, key, len);
parent = find_edge_node_locked(thmap, &query, key, len, &slot);
if (!parent) {
/* Root slot empty: not found. */
return NULL;
}
leaf = get_leaf(thmap, parent, slot);
if (!leaf || !key_cmp_p(thmap, leaf, key, len)) {
/* Not found. */
unlock_node(parent);
return NULL;
}

/* Remove the leaf. */
ASSERT(THMAP_NODE(thmap, atomic_load_relaxed(&parent->slots[slot]))
== leaf);
node_remove(parent, slot);

/*
* Collapse the levels if removing the last item.
*/
while (query.level &&
NODE_COUNT(atomic_load_relaxed(&parent->state)) == 0) {
thmap_inode_t *node = parent;

ASSERT(atomic_load_relaxed(&node->state) == NODE_LOCKED);

/*
* Ascend one level up.
* => Mark our current parent as deleted.
* => Lock the parent one level up.
*/
query.level--;
slot = hashval_getslot(&query, key, len);
parent = THMAP_NODE(thmap, node->parent);
ASSERT(parent != NULL);

lock_node(parent);
ASSERT((atomic_load_relaxed(&parent->state) & NODE_DELETED)
== 0);

/*
* Lock is exclusive, so nobody else can be writing at
* the same time, and no need for atomic R/M/W, but
* readers may read without the lock and so need atomic
* load/store. No ordering here needed because the
* entry itself stays valid until GC.
*/
atomic_store_relaxed(&node->state,
atomic_load_relaxed(&node->state) | NODE_DELETED);
unlock_node(node); // memory_order_release

ASSERT(THMAP_NODE(thmap,
atomic_load_relaxed(&parent->slots[slot])) == node);
node_remove(parent, slot);

/* Stage the removed node for G/C. */
stage_mem_gc(thmap, THMAP_GETOFF(thmap, node), THMAP_INODE_LEN);
}

/*
* If the top node is empty, then we need to remove it from the
* root level. Mark the node as deleted and clear the slot.
*
* Note: acquiring the lock on the top node effectively prevents
* the root slot from changing.
*/
if (NODE_COUNT(atomic_load_relaxed(&parent->state)) == 0) {
const unsigned rslot = query.rslot;
const thmap_ptr_t nptr =
atomic_load_relaxed(&thmap->root[rslot]);

ASSERT(query.level == 0);
ASSERT(parent->parent == THMAP_NULL);
ASSERT(THMAP_GETOFF(thmap, parent) == nptr);

/* Mark as deleted and remove from the root-level slot. */
atomic_store_relaxed(&parent->state,
atomic_load_relaxed(&parent->state) | NODE_DELETED);
atomic_store_relaxed(&thmap->root[rslot], THMAP_NULL);

stage_mem_gc(thmap, nptr, THMAP_INODE_LEN);
}
unlock_node(parent);

/*
* Save the value and stage the leaf for G/C.
*/
val = leaf->val;
if ((thmap->flags & THMAP_NOCOPY) == 0) {
stage_mem_gc(thmap, leaf->key, leaf->len);
}
stage_mem_gc(thmap, THMAP_GETOFF(thmap, leaf), sizeof(thmap_leaf_t));
return val;
}

/*
* G/C routines.
*/

static void
stage_mem_gc(thmap_t *thmap, uintptr_t addr, size_t len)
{
char *const ptr = THMAP_GETPTR(thmap, addr);
thmap_gc_t *head, *gc;

gc = container_of(ptr, struct thmap_gc, data[0]);
KASSERTMSG(gc->len == len,
"thmap=%p ops=%p ptr=%p len=%zu gc=%p gc->len=%zu",
thmap, thmap->ops, (char *)addr, len, gc, gc->len);
retry:
head = atomic_load_relaxed(&thmap->gc_list);
gc->next = head; // not yet published

/* Release to subsequent acquire in thmap_stage_gc(). */
if (!atomic_compare_exchange_weak_explicit_ptr(&thmap->gc_list, &head, gc,
memory_order_release, memory_order_relaxed)) {
goto retry;
}
}

void *
thmap_stage_gc(thmap_t *thmap)
{
/* Acquire from prior release in stage_mem_gc(). */
return atomic_exchange_explicit(&thmap->gc_list, NULL,
memory_order_acquire);
}

void
thmap_gc(thmap_t *thmap, void *ref)
{
thmap_gc_t *gc = ref;

while (gc) {
thmap_gc_t *next = gc->next;
gc_free(thmap, THMAP_GETOFF(thmap, &gc->data[0]), gc->len);
gc = next;
}
}

/*
* thmap_create: construct a new trie-hash map object.
*/
thmap_t *
thmap_create(uintptr_t baseptr, const thmap_ops_t *ops, unsigned flags)
{
thmap_t *thmap;
uintptr_t root;

/*
* Setup the map object.
*/
if (!THMAP_ALIGNED_P(baseptr)) {
return NULL;
}
thmap = kmem_zalloc(sizeof(thmap_t), KM_SLEEP);
thmap->baseptr = baseptr;
thmap->ops = ops ? ops : &thmap_default_ops;
thmap->flags = flags;

if ((thmap->flags & THMAP_SETROOT) == 0) {
/* Allocate the root level. */
root = gc_alloc(thmap, THMAP_ROOT_LEN);
if (!root) {
kmem_free(thmap, sizeof(thmap_t));
return NULL;
}
thmap->root = THMAP_GETPTR(thmap, root);
memset(thmap->root, 0, THMAP_ROOT_LEN);
}

cprng_strong(kern_cprng, thmap->seed, sizeof thmap->seed, 0);

return thmap;
}

int
thmap_setroot(thmap_t *thmap, uintptr_t root_off)
{
if (thmap->root) {
return -1;
}
thmap->root = THMAP_GETPTR(thmap, root_off);
return 0;
}

uintptr_t
thmap_getroot(const thmap_t *thmap)
{
return THMAP_GETOFF(thmap, thmap->root);
}

void
thmap_destroy(thmap_t *thmap)
{
uintptr_t root = THMAP_GETOFF(thmap, thmap->root);
void *ref;

ref = thmap_stage_gc(thmap);
thmap_gc(thmap, ref);

if ((thmap->flags & THMAP_SETROOT) == 0) {
gc_free(thmap, root, THMAP_ROOT_LEN);
}
kmem_free(thmap, sizeof(thmap_t));
}