/*-
 * Copyright (c) 2020 UPLEX Nils Goroll Systemoptimierung
 * All rights reserved
 *
 * Author: Geoffrey Simmons <geoffrey.simmons@uplex.de>
 *
 * 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 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.
 */

/*
 * Perfect hashing based on universal hashing. See:
 * - Fredman, Komlos & Szemeredi (1984), JACM 31 (3)
 * - Thorup (2020), arxiv.org 1504.06804
 * This is the same theory behind the NH algorithm applied as part of UMAC
 * message authentication, see RFC 4418 (esp section 5.2.2).
 */

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

#include "vdef.h"
#include "vas.h"
#include "miniobj.h"
#include "vbm.h"

#include "ph.h"
#include "rnd.h"

/*
 * Limit to MAXN strings in a set, so that we can round up the size of the
 * hash table to a power of 2, but can use UINT_MAX to signify "no match".
 *
 * MAX_BUCKETN is the largest N such that N^2 < MAXN.
 */
#define MAX_LG (31)
#define MAXN (1U << MAX_LG)
#define MAX_BUCKETN (46340)

struct hash {
	unsigned	magic;
#define HASH_MAGIC	0x11a887ce
	uint32_t	mask;
	uint64_t	*k;
	uint64_t	addend;
	uint32_t	*tbl;
	size_t		minlen;
	size_t		maxlen;
	size_t		l;
};

union tbl_t {
	struct hash	*h2;
	uint32_t	idx;
};

struct ph {
	unsigned	magic;
#define PH_MAGIC	0x00cd8c1d
	struct hash	*h1;
	union tbl_t	*tbl;
	struct vbitmap	*collision;
};

struct bucket {
	uint32_t	*idx;
	int		n;
};

/*
 * Multilinear-HM from Lemire & Kaser (2018), except that we add the tail
 * part and the length bounds. arxiv.org 1202.4961
 */
static inline uint32_t
hash(const struct hash * const restrict hash,
     const char * const restrict subject, size_t len)
{
	uint64_t h = hash->addend;
	const uint64_t *k = hash->k;
	const uint32_t *s = (const uint32_t *)subject, *e;
	uint32_t tail[2] = {0};
	size_t l = (len / 8) * 2;

	if (len < hash->minlen || len > hash->maxlen)
		return (UINT_MAX);
	assert((len >> 2) <= hash->l);

	for (e = s + l; s < e; k += 2, s += 2)
		h += (s[0] + k[0]) * (s[1] + k[1]);

	/* When the subject length is not an exact multiple of 8. */
	if ((len & 0x07) != 0) {
		memcpy(tail, e, (void *)(subject + len) - (void *)e);
		h += (tail[0] + k[0]) * (tail[1] + k[1]);
	}

	return ((h >> 32) & hash->mask);
}

/*
 * https://graphics.stanford.edu/~seander/bithacks.html#IntegerLogObvious
 */
static unsigned
lg(unsigned n)
{
	unsigned lg = 0;
	while (n >>= 1)
		lg++;
	return (lg);
}

static uint32_t
getsz(unsigned n)
{
	unsigned bits = lg(n);
	if (n != (unsigned)(1U << bits))
		bits++;
	assert(bits <= MAX_LG);
	return (1U << bits);
}

static uint64_t
rnd64(void)
{
	return (rnd_nxt() + ((uint64_t)rnd_nxt() << 32));
}

void
PH_Init(uint32_t seed[4])
{
	rnd_init(seed);
}

struct ph *
PH_Generate(char * const * const strings, unsigned n)
{
	size_t min = SIZE_MAX, max = 0, l;
	struct hash *h1 = NULL;
	union tbl_t *tbl = NULL;
	struct vbitmap *collision = NULL;
	struct bucket *buckets = NULL;
	uint32_t sz;
	struct ph *ph = NULL;

	AN(strings);
	if (n == 0 || n > MAXN) {
		errno = ERANGE;
		return (NULL);
	}

	sz = getsz(n);

	for (unsigned i = 0; i < n; i++) {
		size_t len = strlen(strings[i]);
		if (len < min)
			min = len;
		if (len > max)
			max = len;
	}
	l = ((max + 7) / 8) * 2;

	errno = 0;
	buckets = calloc(sz, sizeof(*buckets));
	if (buckets == NULL)
		return (NULL);
	tbl = malloc(sz * sizeof(*tbl));
	if (tbl == NULL)
		goto exit;
	memset(tbl, 0xff, sz * sizeof(*tbl));
	ALLOC_OBJ(h1, HASH_MAGIC);
	if (h1 == NULL)
		goto exit;
	collision = vbit_new(sz);

	h1->k = malloc(l * sizeof(*h1->k));
	if (h1->k == NULL)
		goto exit;
	h1->addend = rnd64();
	h1->minlen = min;
	h1->maxlen = max;
	h1->l = l;
	for (unsigned i = 0; i < l; i++)
		h1->k[i] = rnd64();
	h1->mask = sz - 1;

	for (unsigned i = 0; i < n; i++) {
		uint32_t h;

		h = hash(h1, strings[i], strlen(strings[i]));
		assert(h < sz);

		errno = 0;
		buckets[h].idx = realloc(buckets[h].idx,
					 (buckets[h].n + 1)
					 * sizeof(*buckets[h].idx));
		if (buckets[h].idx == NULL)
			goto exit;
		buckets[h].idx[buckets[h].n] = i;
		buckets[h].n++;
	}

	for (unsigned i = 0; i < sz; i++) {
		struct bucket *bucket;
		struct hash *bhsh;
		uint32_t bsz;

		bucket = &buckets[i];
		if (bucket->n == 0) {
			AZ(vbit_test(collision, i));
			assert(tbl[i].idx == UINT_MAX);
			continue;
		}
		AN(bucket->idx);
		if (bucket->n == 1) {
			AZ(vbit_test(collision, i));
			tbl[i].idx = *bucket->idx;
			continue;
		}
		/* It's infinitesimally improbable, just assert. */
		assert(bucket->n <= MAX_BUCKETN);

		vbit_set(collision, i);
		bsz = getsz(bucket->n * bucket->n);
		assert(bsz <= MAXN);

		errno = 0;
		ALLOC_OBJ(bhsh, HASH_MAGIC);
		if (bhsh == NULL)
			goto exit;
		tbl[i].h2 = bhsh;
		bhsh->tbl = malloc(bsz * sizeof(*bhsh->tbl));
		if (bhsh->tbl == NULL)
			goto exit;
		memset(bhsh->tbl, 0xff, bsz * sizeof(*bhsh->tbl));
		bhsh->mask = bsz - 1;

		bhsh->minlen = SIZE_MAX;
		bhsh->maxlen = 0;
		for (int j = 0; j < bucket->n; j++) {
			size_t len = strlen(strings[bucket->idx[j]]);
			if (len < bhsh->minlen)
				bhsh->minlen = len;
			if (len > bhsh->maxlen)
				bhsh->maxlen = len;
		}
		bhsh->l = ((bhsh->maxlen + 7) / 8) * 2;

		bhsh->k = malloc(bhsh->l * sizeof(*bhsh->k));
		if (bhsh->k == NULL)
			goto exit;
		bhsh->addend = rnd64();
		for (unsigned j = 0; j < bhsh->l; j++)
			bhsh->k[j] = rnd64();

		for (int j = 0; j < bucket->n; j++) {
			uint32_t h = hash(bhsh, strings[bucket->idx[j]],
					  strlen(strings[bucket->idx[j]]));
			assert(h < bsz);
			if (bhsh->tbl[h] != UINT_MAX) {
				j = -1;
				memset(bhsh->tbl, 0xff,
				       bsz * sizeof(*bhsh->tbl));
				bhsh->addend = rnd64();
				for (unsigned k = 0; k < bhsh->l; k++)
					bhsh->k[k] = rnd64();
				continue;
			}
			bhsh->tbl[h] = bucket->idx[j];
		}
	}

	errno = 0;
	ALLOC_OBJ(ph, PH_MAGIC);
	if (ph == NULL)
		goto exit;
	ph->h1 = h1;
	ph->tbl = tbl;
	ph->collision = collision;

 exit:
	AN(buckets);
	for (unsigned i = 0; i < sz; i++)
		if (buckets[i].idx != NULL)
			free(buckets[i].idx);
	free(buckets);
	if (ph == NULL) {
		if (h1 != NULL) {
			AZ(ph->h1->tbl);
			if (h1->k != NULL)
				free(h1->k);
			FREE_OBJ(h1);
		}
		if (tbl != NULL) {
			if (collision != NULL) {
				for (unsigned i = 0; i < sz; i++) {
					if (!vbit_test(collision, i))
						continue;
					CHECK_OBJ_NOTNULL(tbl[i].h2,
							  HASH_MAGIC);
					if (tbl[i].h2->tbl != NULL)
						free(tbl[i].h2->tbl);
					if (tbl[i].h2->k != NULL)
						free(tbl[i].h2->k);
					FREE_OBJ(tbl[i].h2);
				}
				vbit_destroy(collision);
			}
			free(tbl);
		}
	}

	return (ph);
}

unsigned
PH_Lookup(const struct ph * const restrict ph,
	  char * const restrict * const restrict strings,
	  const char * const restrict subject)
{
	size_t len;
	uint32_t h;
	unsigned idx = UINT_MAX;

	if (ph == NULL)
		return (UINT_MAX);

	CHECK_OBJ(ph, PH_MAGIC);
	CHECK_OBJ_NOTNULL(ph->h1, HASH_MAGIC);
	AN(ph->tbl);
	AN(strings);
	AN(subject);

	len = strlen(subject);
	h = hash(ph->h1, subject, len);
	if (h == UINT_MAX)
		return (UINT_MAX);
	idx = ph->tbl[h].idx;
	if (vbit_test(ph->collision, h)) {
		struct hash *h2 = ph->tbl[h].h2;
		CHECK_OBJ_NOTNULL(h2, HASH_MAGIC);
		AN(h2->tbl);
		h = hash(h2, subject, len);
		if (h == UINT_MAX)
			return (UINT_MAX);
		idx = h2->tbl[h];
	}

	if (idx == UINT_MAX || strcmp(subject, strings[idx]) != 0)
		return (UINT_MAX);
	return (idx);
}

struct vsb *
PH_Dump(struct ph *ph, char **strings)
{
	struct vsb *sb = VSB_new_auto();

	if (ph == NULL) {
		VSB_finish(sb);
		return (sb);
	}

	CHECK_OBJ(ph, PH_MAGIC);
	CHECK_OBJ_NOTNULL(ph->h1, HASH_MAGIC);
	AN(strings);
	VSB_printf(sb, "minlen = %zu\n", ph->h1->minlen);
	VSB_printf(sb, "maxlen = %zu\n", ph->h1->maxlen);
	VSB_printf(sb, "l = %zu\n", ph->h1->l);

	VSB_printf(sb, "h1->mask = 0x%0x\n", ph->h1->mask);
	VSB_printf(sb, "h1->addend = 0x%0lx\n", ph->h1->addend);
	for (unsigned i = 0; i < ph->h1->l; i++)
		VSB_printf(sb, "h1->k[%u] = 0x%0lx\n", i, ph->h1->k[i]);
	for (unsigned i = 0; i <= ph->h1->mask; i++) {
		VSB_printf(sb, "\n");
		if (!vbit_test(ph->collision, i)) {
			VSB_printf(sb, "tbl[%u].idx = %u\n", i, ph->tbl[i].idx);
			if (ph->tbl[i].idx != UINT_MAX)
				VSB_printf(sb, "\tstrings[%u] = %s\n",
					   ph->tbl[i].idx,
					   strings[ph->tbl[i].idx]);
			continue;
		}

		struct hash *h2 = ph->tbl[i].h2;
		VSB_printf(sb, "tbl[%u].h2 = %p\n", i, h2);
		if (h2 == NULL)
			continue;
		CHECK_OBJ(h2, HASH_MAGIC);
		VSB_printf(sb, "tbl[%u].h2->l = %zu\n", i, h2->l);
		VSB_printf(sb, "tbl[%u].h2->minlen = %zu\n", i, h2->minlen);
		VSB_printf(sb, "tbl[%u].h2->maxlen = %zu\n", i, h2->maxlen);
		VSB_printf(sb, "tbl[%u].h2->mask = 0x%0x\n", i, h2->mask);
		VSB_printf(sb, "tbl[%u].h2->addend = 0x%0lx\n", i, h2->addend);
		for (unsigned j = 0; j < h2->l; j++)
			VSB_printf(sb, "tbl[%u].h2->k[%u] = 0x%0lx\n", i, j,
				   h2->k[j]);
		for (unsigned j = 0; j <= h2->mask; j++) {
			VSB_printf(sb, "tbl[%u].h2->tbl[%u] = %u\n", i, j,
				   h2->tbl[j]);
			if (h2->tbl[j] != UINT_MAX)
				VSB_printf(sb, "\tstrings[%u] = %s\n",
					   h2->tbl[j], strings[h2->tbl[j]]);
		}
	}

	VSB_finish(sb);
	return (sb);
}

void
PH_Stats(const struct ph * const restrict ph,
	 char * const restrict * const restrict strings,
	 struct ph_stats * const restrict stats)
{
	CHECK_OBJ_NOTNULL(stats, PH_STATS_MAGIC);

	memset(stats, 0, sizeof(*stats));
	if (ph == NULL)
		return;
	CHECK_OBJ(ph, PH_MAGIC);
	CHECK_OBJ_NOTNULL(ph->h1, HASH_MAGIC);
	AN(ph->tbl);
	AN(ph->collision);
	AN(strings);

	stats->buckets = ph->h1->mask + 1;
	stats->klen = ph->h1->l;
	stats->minlen = ph->h1->minlen;
	stats->maxlen = ph->h1->maxlen;

	stats->h2buckets_min = UINT64_MAX;
	stats->h2strings_min = UINT64_MAX;
	stats->h2klen_min = SIZE_MAX;
	for (unsigned i = 0; i <= ph->h1->mask; i++)
		if (vbit_test(ph->collision, i)) {
			struct hash *h2;
			uint64_t sz, nstrings = 0;

			h2 = ph->tbl[i].h2;
			CHECK_OBJ_NOTNULL(h2, HASH_MAGIC);
			sz = h2->mask + 1;

			stats->collisions++;
			if (sz < stats->h2buckets_min)
				stats->h2buckets_min = sz;
			if (sz > stats->h2buckets_max)
				stats->h2buckets_max = sz;
			stats->h2buckets_avg +=
				((double)sz - stats->h2buckets_avg)
				/ (double)stats->collisions;

			if (h2->l < stats->h2klen_min)
				stats->h2klen_min = h2->l;
			if (h2->l > stats->h2klen_max)
				stats->h2klen_max = h2->l;
			stats->h2klen_avg +=
				((double)h2->l - stats->h2klen_avg)
				/ (double)stats->collisions;

			for (unsigned j = 0; j < sz; j++) {
				if (h2->tbl[j] == UINT_MAX)
					continue;
				nstrings++;
			}
			if (nstrings < stats->h2strings_min)
				stats->h2strings_min = nstrings;
			if (nstrings > stats->h2strings_max)
				stats->h2strings_max = nstrings;
			stats->h2strings_avg +=
				((double)nstrings - stats->h2strings_avg)
				/ (double)stats->collisions;
		}
	if (stats->h2buckets_min == UINT64_MAX)
		stats->h2buckets_min = 0;
	if (stats->h2strings_min == UINT64_MAX)
		stats->h2strings_min = 0;
	if (stats->h2klen_min == SIZE_MAX)
		stats->h2klen_min = 0;
}

void
PH_Free(struct ph *ph)
{
	if (ph == NULL)
		return;
	CHECK_OBJ(ph, PH_MAGIC);
	if (ph->tbl != NULL) {
		if (ph->collision != NULL) {
			for (unsigned i = 0; i <= ph->h1->mask; i++) {
				if (!vbit_test(ph->collision, i))
					continue;
				CHECK_OBJ_NOTNULL(ph->tbl[i].h2, HASH_MAGIC);
				if (ph->tbl[i].h2->tbl != NULL)
					free(ph->tbl[i].h2->tbl);
				if (ph->tbl[i].h2->k != NULL)
					free(ph->tbl[i].h2->k);
				FREE_OBJ(ph->tbl[i].h2);
			}
			vbit_destroy(ph->collision);
		}
		free(ph->tbl);
	}
	if (ph->h1 != NULL) {
		AZ(ph->h1->tbl);
		free(ph->h1->k);
		FREE_OBJ(ph->h1);
	}
	FREE_OBJ(ph);
}
