/* libmisc/map.c - A map/dict data structure
 *
 * Copyright (C) 2024-2025  Luke T. Shumaker <lukeshu@lukeshu.com>
 * SPDX-License-Identifier: AGPL-3.0-or-later
 */

#include <stdlib.h>
#include <string.h>

#include <libmisc/hash.h>
#include <libmisc/assert.h>
#include <libmisc/map.h>

#define FLAG_ITER (UINT8_C(1)<<0)
#define FLAG_DEL  (UINT8_C(1)<<1)

/* Internal utilities *********************************************************/

#define cast(n) lm_dll_node_cast(struct _map_kv, n)

static inline void *_map_kv_keyp(struct _map *m, struct _map_kv *kv) {
	assert(m);
	assert(kv);
	return ((void*)kv)+m->offsetof_k;
}
static inline void *_map_kv_valp(struct _map *m, struct _map_kv *kv) {
	assert(m);
	assert(kv);
	return ((void*)kv)+m->offsetof_v;
}

static inline void _map_lookup(struct _map *m, void *keyp,
                 hash_t *ret_hash,
                 lm_dll_root **ret_bucket,
                 struct _map_kv **ret_kv) {
	assert(m);
	assert(keyp);
	assert(ret_hash);
	assert(ret_bucket);
	assert(ret_kv);
	*ret_hash = hash(keyp, m->sizeof_k);
	if (m->nbuckets == 0) {
		*ret_bucket = NULL;
		*ret_kv = NULL;
		return;
	}
	*ret_bucket = &m->buckets[*ret_hash % m->nbuckets];
	for (struct _map_kv *kv = cast((*ret_bucket)->front); kv; kv = cast(kv->rear)) {
		if (!(kv->flags & FLAG_DEL) &&
		    memcmp(_map_kv_keyp(m, kv), keyp, m->sizeof_k) == 0) {
			*ret_kv = kv;
			return;
		}
	}
	*ret_kv = NULL;
}

static inline void _map_resize(struct _map *m, size_t new_nbuckets) {
	assert(m);
	assert(new_nbuckets);
	lm_dll_root *new_buckets = calloc(new_nbuckets, sizeof(lm_dll_root));
	for (size_t i = 0; i < m->nbuckets; i++) {
		while (m->buckets[i].front) {
			struct _map_kv *kv = cast(m->buckets[i].front);
			lm_dll_pop_from_front(&m->buckets[i]);
			hash_t h = hash(_map_kv_keyp(m, kv), m->sizeof_k);
			lm_dll_push_to_rear(&new_buckets[h % new_nbuckets], kv);
		}
	}
	m->nbuckets = new_nbuckets;
	free(m->buckets);
	m->buckets = new_buckets;
}

static bool _map_autoresize(struct _map *m) {
	assert(m);
	if (m->len > (m->nbuckets * 8/10)) {
		size_t nbuckets = 1;
		while (m->len > (nbuckets * 8/10))
			nbuckets <<= 1;
		_map_resize(m, nbuckets);
		return true;
	}
	return false;
}

/* Methods ********************************************************************/

void *_map_load(struct _map *m, void *keyp) {
	assert(m);
	assert(keyp);

	hash_t h;
	lm_dll_root *bucket;
	struct _map_kv *kv;
	_map_lookup(m, keyp, &h, &bucket, &kv);

	if (!kv)
		return NULL;
	return _map_kv_valp(m, kv);
}

bool _map_del(struct _map *m, void *keyp) {
	assert(m);
	assert(keyp);

	hash_t h;
	lm_dll_root *bucket;
	struct _map_kv *kv;
	_map_lookup(m, keyp, &h, &bucket, &kv);

	if (!kv)
		return false;
	if (kv->flags & FLAG_ITER) {
		kv->flags |= FLAG_DEL;
	} else {
		lm_dll_remove(bucket, kv);
		free(kv);
	}
	m->len--;
	return true;
}

void *_map_store(struct _map *m, void *keyp, void *valp) {
	assert(m);
	assert(keyp);
	assert(valp);

	hash_t h;
	lm_dll_root *bucket;
	struct _map_kv *old;
	_map_lookup(m, keyp, &h, &bucket, &old);

	if (old) {
		lm_dll_remove(bucket, old);
		free(old);
		m->len--;
	}
	m->len++;
	if (!m->iterating && _map_autoresize(m)) {
		h = hash(keyp, m->sizeof_k);
		bucket = &m->buckets[h % m->nbuckets];
	}
	struct _map_kv *kv = calloc(1, m->sizeof_kv);
	memcpy(_map_kv_keyp(m, kv), keyp, m->sizeof_k);
	memcpy(_map_kv_valp(m, kv), valp, m->sizeof_v);
	lm_dll_push_to_rear(bucket, kv);
	return _map_kv_valp(m, kv);
}

void _map_free(struct _map *m) {
	assert(m);

	for (size_t i = 0; i < m->nbuckets; i++) {
		while (m->buckets[i].front) {
			struct _map_kv *kv = cast(m->buckets[i].front);
			lm_dll_pop_from_front(&m->buckets[i]);
			free(kv);
		}
	}
	free(m->buckets);
	m->len = 0;
	m->nbuckets = 0;
	m->buckets = NULL;
}

/* Iteration ******************************************************************/

struct _map_iter _map_iter_before(struct _map *m, void *keyp, void **valpp) {
	assert(m);
	assert(keyp);
	assert(valpp);

	struct _map_iter state = {
		.m = m,
		.keyp = keyp,
		.valpp = valpp,
	};
	m->iterating++;
	return state;
}

void _map_iter_after(struct _map_iter *state) {
	assert(state);
	assert(state->m);

	state->m->iterating--;
	if (!state->m->iterating)
		_map_autoresize(state->m);
}

bool _map_iter_next(struct _map_iter *state) {
	assert(state);
	assert(state->m);
	assert(state->valpp);

	if (!state->kv) {
		if (!state->m->len)
			return false;
		while (!state->m->buckets[state->i].front)
			state->i++;
		state->kv = cast(state->m->buckets[state->i].front);
	} else {
		struct _map_kv *old_kv = state->kv;
		state->kv = cast(old_kv->rear);

		old_kv->flags &= ~FLAG_ITER;
		if (old_kv->flags & FLAG_DEL) {
			lm_dll_remove(&state->m->buckets[state->i], old_kv);
			free(old_kv);
		}

		while (!state->kv) {
			state->i++;
			if (state->i == state->m->nbuckets)
				return false;
			state->kv = cast(state->m->buckets[state->i].front);
		}
	}
	state->kv->flags |= FLAG_ITER;
	memcpy(state->keyp, _map_kv_keyp(state->m, state->kv), state->m->sizeof_k);
	*(state->valpp) = _map_kv_valp(state->m, state->kv);
	return true;
}