/*

 * This file is part of ltrace.

 * Copyright (C) 2012, 2013 Petr Machata, Red Hat Inc.

 *

 * This program is free software; you can redistribute it and/or

 * modify it under the terms of the GNU General Public License as

 * published by the Free Software Foundation; either version 2 of the

 * License, or (at your option) any later version.

 *

 * This program is distributed in the hope that it will be useful, but

 * WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU

 * General Public License for more details.

 *

 * You should have received a copy of the GNU General Public License

 * along with this program; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA

 * 02110-1301 USA

 */

#include <string.h>

#include <stdlib.h>

#include <stdio.h>

#include "dict.h"

struct status_bits {

	unsigned char taken : 1;

	unsigned char erased : 1;

};

static struct status_bits *

bitp(struct dict *dict, size_t n)

{

	return VECT_ELEMENT(&dict->status, struct status_bits, n);

}

void

dict_init(struct dict *dict,

	  size_t key_size, size_t value_size,

	  size_t (*hash1)(const void *),

	  int (*eq)(const void *, const void *),

	  size_t (*hash2)(size_t))

{

	assert(hash1 != NULL);

	assert(eq != NULL);

	vect_init(&dict->keys, key_size);

	vect_init(&dict->values, value_size);

	VECT_INIT(&dict->status, struct status_bits);

	dict->size = 0;

	dict->hash1 = hash1;

	dict->hash2 = hash2;

	dict->eq = eq;

}

struct clone_data {

	struct dict *target;

	int (*clone_key)(void *tgt, const void *src, void *data);

	int (*clone_value)(void *tgt, const void *src, void *data);

	void (*dtor_key)(void *tgt, void *data);

	void (*dtor_value)(void *tgt, void *data);

	void *data;

};

static enum callback_status

clone_cb(void *key, void *value, void *data)

{

	struct clone_data *clone_data = data;

	char nkey[clone_data->target->keys.elt_size];

	if (clone_data->clone_key == NULL)

		memmove(nkey, key, sizeof(nkey));

	else if (clone_data->clone_key(&nkey, key, clone_data->data) < 0)

		return CBS_STOP;

	char nvalue[clone_data->target->values.elt_size];

	if (clone_data->clone_value == NULL) {

		memmove(nvalue, value, sizeof(nvalue));

	} else if (clone_data->clone_value(&nvalue, value,

					 clone_data->data) < 0) {

	fail:

		if (clone_data->clone_key != NULL)

			clone_data->dtor_key(&nkey, clone_data->data);

		return CBS_STOP;

	}

	if (dict_insert(clone_data->target, nkey, nvalue) < 0) {

		if (clone_data->clone_value != NULL)

			clone_data->dtor_value(&nvalue, clone_data->data);

		goto fail;

	}

	return CBS_CONT;

}

int

dict_clone(struct dict *target, const struct dict *source,

	   int (*clone_key)(void *tgt, const void *src, void *data),

	   void (*dtor_key)(void *tgt, void *data),

	   int (*clone_value)(void *tgt, const void *src, void *data),

	   void (*dtor_value)(void *tgt, void *data),

	   void *data)

{

	assert((clone_key != NULL) == (dtor_key != NULL));

	assert((clone_value != NULL) == (dtor_value != NULL));

	dict_init(target, source->keys.elt_size, source->values.elt_size,

		  source->hash1, source->eq, source->hash2);

	struct clone_data clone_data = {

		target, clone_key, clone_value, dtor_key, dtor_value, data

	};

	if (dict_each((struct dict *)source, NULL,

		      clone_cb, &clone_data) != NULL) {

		dict_destroy(target, dtor_key, dtor_value, data);

		return -1;

	}

	return 0;

}

size_t

dict_size(const struct dict *dict)

{

	return dict->size;

}

int

dict_empty(const struct dict *dict)

{

	return dict->size == 0;

}

struct destroy_data {

	void (*dtor_key)(void *tgt, void *data);

	void (*dtor_value)(void *tgt, void *data);

	void *data;

};

static enum callback_status

destroy_cb(void *key, void *value, void *data)

{

	struct destroy_data *destroy_data = data;

	if (destroy_data->dtor_key)

		destroy_data->dtor_key(key, destroy_data->data);

	if (destroy_data->dtor_value)

		destroy_data->dtor_value(value, destroy_data->data);

	return CBS_CONT;

}

void

dict_destroy(struct dict *dict,

	     void (*dtor_key)(void *tgt, void *data),

	     void (*dtor_value)(void *tgt, void *data),

	     void *data)

{

	/* Some slots are unused (the corresponding keys and values

	 * are uninitialized), so we can't call dtors for them.

	 * Iterate DICT instead.  */

	if (dtor_key != NULL || dtor_value != NULL) {

		struct destroy_data destroy_data = {

			dtor_key, dtor_value, data

		};

		dict_each(dict, NULL, destroy_cb, &destroy_data);

	}

	vect_destroy(&dict->keys, NULL, NULL);

	vect_destroy(&dict->values, NULL, NULL);

	vect_destroy(&dict->status, NULL, NULL);

}

static size_t

default_secondary_hash(size_t pos)

{

	return pos % 97 + 1;

}

static size_t

small_secondary_hash(size_t pos)

{

	return 1;

}

static inline size_t

n(struct dict *dict)

{

	return vect_size(&dict->keys);

}

static inline size_t (*

hash2(struct dict *dict))(size_t)

{

	if (dict->hash2 != NULL)

		return dict->hash2;

	else if (n(dict) < 100)

		return small_secondary_hash;

	else

		return default_secondary_hash;

}

static void *

getkey(struct dict *dict, size_t pos)

{

	return ((unsigned char *)dict->keys.data)

		+ dict->keys.elt_size * pos;

}

static void *

getvalue(struct dict *dict, size_t pos)

{

	return ((unsigned char *)dict->values.data)

		+ dict->values.elt_size * pos;

}

static size_t

find_slot(struct dict *dict, const void *key,

	  int *foundp, int *should_rehash, size_t *pi)

{

	assert(n(dict) > 0);

	size_t pos = dict->hash1(key) % n(dict);

	size_t pos0 = -1;

	size_t d = hash2(dict)(pos);

	size_t i = 0;

	*foundp = 0;

	/* We skip over any taken or erased slots.  But we remember

	 * the first erased that we find, and if we don't find the key

	 * later, we return that position.  */

	for (; bitp(dict, pos)->taken || bitp(dict, pos)->erased;

	     pos = (pos + d) % n(dict)) {

		if (pos0 == (size_t)-1 && bitp(dict, pos)->erased)

			pos0 = pos;

		/* If there is a loop, but we've seen an erased

		 * element, take that one.  Otherwise give up.  */

		if (++i > dict->size) {

			if (pos0 != (size_t)-1)

				break;

			return (size_t)-1;

		}

		if (bitp(dict, pos)->taken

		    && dict->eq(getkey(dict, pos), key)) {

			*foundp = 1;

			break;

		}

	}

	if (!*foundp && pos0 != (size_t)-1)

		pos = pos0;

	/* If the hash table degraded into a linked list, request a

	 * rehash.  */

	if (should_rehash != NULL)

		*should_rehash = i > 10 && i > n(dict) / 10;

	if (pi != NULL)

		*pi = i;

	return pos;

}

static enum callback_status

rehash_move(void *key, void *value, void *data)

{

	if (dict_insert(data, key, value) < 0)

		return CBS_STOP;

	else

		return CBS_CONT;

}

static int

rehash(struct dict *dict, size_t nn)

{

	assert(nn != n(dict));

	int ret = -1;

	struct dict tmp;

	dict_init(&tmp, dict->keys.elt_size, dict->values.elt_size,

		  dict->hash1, dict->eq, dict->hash2);

	/* To honor all invariants (so that we can safely call

	 * dict_destroy), we first make a request to _reserve_ enough

	 * room in all vectors.  This has no observable effect on

	 * contents of vectors.  */

	if (vect_reserve(&tmp.keys, nn) < 0

	    || vect_reserve(&tmp.values, nn) < 0

	    || vect_reserve(&tmp.status, nn) < 0)

		goto done;

	/* Now that we know that there is enough size in vectors, we

	 * simply bump the size.  */

	tmp.keys.size = nn;

	tmp.values.size = nn;

	size_t old_size = tmp.status.size;

	tmp.status.size = nn;

	memset(VECT_ELEMENT(&tmp.status, struct status_bits, old_size),

	       0, (tmp.status.size - old_size) * tmp.status.elt_size);

	/* At this point, TMP is once more an empty dictionary with NN

	 * slots.  Now move stuff from DICT to TMP.  */

	if (dict_each(dict, NULL, rehash_move, &tmp) != NULL)

		goto done;

	/* And now swap contents of DICT and TMP, and we are done.  */

	{

		struct dict tmp2 = *dict;

		*dict = tmp;

		tmp = tmp2;

	}

	ret = 0;

done:

	/* We only want to release the containers, not the actual data

	 * that they hold, so it's fine if we don't pass any dtor.  */

	dict_destroy(&tmp, NULL, NULL, NULL);

	return ret;

}

static const size_t primes[] = {

	13, 31, 61, 127, 251, 509, 1021, 2039, 4093,

	8191, 16381, 32749, 65521, 130981, 0

};

static size_t

larger_size(size_t current)

{

	if (current == 0)

		return primes[0];

	if (current < primes[sizeof(primes)/sizeof(*primes) - 2]) {

		size_t i;

		for (i = 0; primes[i] != 0; ++i)

			if (primes[i] > current)

				return primes[i];

		abort();

	}

	/* We ran out of primes, so invent a new one.  The following

	 * gives primes until about 17M elements (and then some more

	 * later).  */

	return 2 * current + 6585;

}

static size_t

smaller_size(size_t current)

{

	if (current <= primes[0])

		return primes[0];

	if (current <= primes[sizeof(primes)/sizeof(*primes) - 2]) {

		size_t i;

		size_t prev = 0;

		for (i = 0; primes[i] != 0; ++i) {

			if (primes[i] >= current)

				return prev;

			prev = primes[i];

		}

		abort();

	}

	return (current - 6585) / 2;

}

int

dict_insert(struct dict *dict, void *key, void *value)

{

	if (n(dict) == 0 || dict->size > 0.7 * n(dict))

	rehash:

		if (rehash(dict, larger_size(n(dict))) < 0)

			return -1;

	int found;

	int should_rehash;

	size_t slot_n = find_slot(dict, key, &found, &should_rehash, NULL);

	if (slot_n == (size_t)-1)

		return -1;

	if (found)

		return 1;

	assert(!bitp(dict, slot_n)->taken);

	/* If rehash was requested, do that, and retry.  But just live

	 * with it for apparently sparse tables.  No resizing can fix

	 * a rubbish hash.  */

	if (should_rehash && dict->size > 0.3 * n(dict))

		goto rehash;

	memmove(getkey(dict, slot_n), key, dict->keys.elt_size);

	memmove(getvalue(dict, slot_n), value, dict->values.elt_size);

	bitp(dict, slot_n)->taken = 1;

	bitp(dict, slot_n)->erased = 0;

	++dict->size;

	return 0;

}

void *

dict_find(struct dict *dict, const void *key)

{

	if (dict->size == 0)

		return NULL;

	assert(n(dict) > 0);

	int found;

	size_t slot_n = find_slot(dict, key, &found, NULL, NULL);

	if (found)

		return getvalue(dict, slot_n);

	else

		return NULL;

}

int

dict_erase(struct dict *dict, const void *key,

	   void (*dtor_key)(void *tgt, void *data),

	   void (*dtor_value)(void *tgt, void *data),

	   void *data)

{

	int found;

	size_t i;

	size_t slot_n = find_slot(dict, key, &found, NULL, &i);

	if (!found)

		return -1;

	if (dtor_key != NULL)

		dtor_key(getkey(dict, slot_n), data);

	if (dtor_value != NULL)

		dtor_value(getvalue(dict, slot_n), data);

	bitp(dict, slot_n)->taken = 0;

	bitp(dict, slot_n)->erased = 1;

	--dict->size;

	if (dict->size < 0.3 * n(dict)) {

		size_t smaller = smaller_size(n(dict));

		if (smaller != n(dict))

			/* Don't mind if it fails when shrinking.  */

			rehash(dict, smaller);

	}

	return 0;

}

void *

dict_each(struct dict *dict, void *start_after,

	  enum callback_status (*cb)(void *, void *, void *), void *data)

{

	size_t i;

	if (start_after != NULL)

		i = ((start_after - dict->keys.data) / dict->keys.elt_size) + 1;

	else

		i = 0;

	for (; i < dict->keys.size; ++i)

		if (bitp(dict, i)->taken && !bitp(dict, i)->erased) {

			void *key = getkey(dict, i);

			if (cb(key, getvalue(dict, i), data) != CBS_CONT)

				return key;

		}

	return NULL;

}

size_t

dict_hash_int(const int *key)

{

	return (size_t)(*key * 2654435761U);

}

int

dict_eq_int(const int *key1, const int *key2)

{

	return *key1 == *key2;

}

size_t

dict_hash_string(const char **key)

{

	size_t h = 5381;

	const char *str = *key;

	while (*str != 0)

		h = h * 33 ^ *str++;

	return h;

}

int

dict_eq_string(const char **key1, const char **key2)

{

	return strcmp(*key1, *key2) == 0;

}

void

dict_dtor_string(const char **key, void *data)

{

	free((char *)*key);

}

int

dict_clone_string(const char **tgt, const char **src, void *data)

{

	*tgt = strdup(*src);

	return *tgt != NULL ? 0 : -1;

}

#ifdef TEST

static enum callback_status

dump(int *key, int *value, void *data)

{

	char *seen = data;

	assert(seen[*key] == 0);

	seen[*key] = 1;

	assert(*value == *key * 2 + 1);

	return CBS_STOP;

}

static size_t

dict_hash_int_silly(const int *key)

{

	return *key % 10;

}

static void

verify(struct dict *di, size_t len, char *seen)

{

	size_t ct = 0;

	int *it;

	for (it = NULL; (it = DICT_EACH(di, int, int, it, dump, seen)) != NULL;)

		ct++;

	assert(ct == len);

	memset(seen, 0, len);

}

static enum callback_status

fill_keys(int *key, int *value, void *data)

{

	int *array = data;

	array[++array[0]] = *key;

	return CBS_CONT;

}

static void

test1(void)

{

	struct dict di;

	DICT_INIT(&di, int, int, dict_hash_int, dict_eq_int, NULL);

	char seen[100000] = {};

	size_t i;

	for (i = 0; i < sizeof(seen); ++i) {

		int key = i;

		int value = 2 * i + 1;

		DICT_INSERT(&di, &key, &value);

		int *valp = DICT_FIND_REF(&di, &key, int);

		assert(valp != NULL);

		assert(*valp == value);

		assert(dict_size(&di) == i + 1);

	}

	verify(&di, sizeof(seen), seen);

	struct dict d2;

	DICT_CLONE(&d2, &di, int, int, NULL, NULL, NULL, NULL, NULL);

	DICT_DESTROY(&di, int, int, NULL, NULL, NULL);

	verify(&d2, sizeof(seen), seen);

	/* Now we try to gradually erase all elements.  We can't erase

	 * inside a DICT_EACH call, so copy first keys to a separate

	 * memory area first.  */

	int keys[d2.size + 1];

	size_t ct = 0;

	keys[0] = 0;

	DICT_EACH(&d2, int, int, NULL, fill_keys, keys);

	for (i = 0; i < (size_t)keys[0]; ++i) {

		assert(DICT_ERASE(&d2, &keys[i + 1], int,

				  NULL, NULL, NULL) == 0);

		++ct;

	}

	assert(ct == sizeof(seen));

	DICT_DESTROY(&d2, int, int, NULL, NULL, NULL);

}

static void

test_erase(void)

{

	int i;

	/* To test erase, we need a relatively bad hash function, so

	 * that there are some overlapping chains in the table.  */

	struct dict d2;

	DICT_INIT(&d2, int, int, dict_hash_int_silly, dict_eq_int, NULL);

	const int limit = 500;

	for (i = 0; i < limit; ++i) {

		int key = 2 * i + 1;

		int value = 2 * key + 1;

		DICT_INSERT(&d2, &key, &value);

	}

	/* Now we try to delete each of the keys, and verify that none

	 * of the chains was broken.  */

	for (i = 0; i < limit; ++i) {

		struct dict copy;

		DICT_CLONE(&copy, &d2, int, int, NULL, NULL, NULL, NULL, NULL);

		int key = 2 * i + 1;

		DICT_ERASE(&copy, &key, int, NULL, NULL, NULL);

		assert(dict_size(&copy) == dict_size(&d2) - 1);

		int j;

		for (j = 0; j < limit; ++j) {

			key = 2 * j + 1;

			int *valp = DICT_FIND_REF(&copy, &key, int);

			if (i != j) {

				assert(valp != NULL);

				assert(*valp == 2 * key + 1);

			} else {

				assert(valp == NULL);

			}

		}

		DICT_DESTROY(&copy, int, int, NULL, NULL, NULL);

	}

	DICT_DESTROY(&d2, int, int, NULL, NULL, NULL);

}

int main(int argc, char *argv[])

{

	test1();

	test_erase();

	return 0;

}

#endif