/*-
 * 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.
 */

/*
 * Inspired by Varnish hash_critbit.c, tarsnap's Patricia implementation,
 * and Tony Finch's qp tries.
 */

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

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

#include "qp.h"
#include "popcnt_compat.h"

#define ANIB(x) AZ((x) & ~0x0f)

/*
 * A trie is comprised by a tree of qp_y nodes and an array of strings,
 * both of which are owned by a VMOD object. Both are passed into the QP_*
 * functions.
 *
 * idx is an index to the strings array, and the region of the string
 * represented at a node is strings[idx][off] to strings[idx}[off+len],
 * inclusive.
 *
 * Terminating nulls form part of the target string. So "foo\0" and
 * "foobar\0" differ at the high nibble of the fourth byte.
 *
 * term is true if this node is a terminating node. Note that terminating
 * nodes may be non-leaves, if the set of strings contains common
 * prefixes.
 *
 * hinib is true iff the most significant nibble at off+len+1 is the
 * critical nibble.
 *
 * bitmap represents values of the critical nibble for which this node
 * branches -- 0 at the LSB, 15 at the MSB.
 *
 * branches is a sparse array of length popcnt(bitmap) for the child
 * nodes. The branch for the least significat set bit in bitmap is at
 * branch[0], for the next most significant set bit at branch[1], and so
 * on.
 */
struct qp_y {
	unsigned	magic;
#define QP_Y_MAGIC	0x6dfde24a
	unsigned	idx;
	struct qp_y	**branch;
	unsigned short	off;
	unsigned short	len;
	uint16_t	bitmap;
	unsigned int	hinib:1;
	unsigned int	term:1;
};

static struct qp_y *
y_alloc(unsigned idx, unsigned short off, size_t len)
{
	struct qp_y *y;

	if (len > USHRT_MAX) {
		errno = ERANGE;
		return (NULL);
	}

	errno = 0;
	ALLOC_OBJ(y, QP_Y_MAGIC);
	if (y == NULL)
		return (NULL);

	y->idx = idx;
	y->off = off;
	y->len = (unsigned short) len;
	AZ(y->branch);
	AZ(y->bitmap);
	AZ(y->hinib);
	AZ(y->term);
	return (y);
}

static inline struct qp_y *
y_leaf_alloc(unsigned idx, unsigned char *c, unsigned char *b)
{
	struct qp_y *y = y_alloc(idx, (unsigned short)(uintptr_t)(c - b),
				 strlen((char *)c));
	y->term = 1;
	return (y);
}

static int
y_realloc_branch(struct qp_y * const y, uint16_t bitmap)
{
	int len;
	uint16_t prev, lobitmap;
	int16_t higher;
	uint8_t loidx;

	assert(popcount(bitmap) == 1);
	AN(y->bitmap & bitmap);

	len = popcount(y->bitmap);
	assert(len <= 16);
	errno = 0;
	y->branch = realloc(y->branch, len * sizeof(*y->branch));
	if (y->branch == NULL)
		return (-1);

	/*
	 * If there was a bit in the previous bitmap such that the new bit
	 * is lower, move up the array entries from that index.
	 */
	prev = y->bitmap & ~bitmap;
	higher = prev & ~((bitmap << 1) - 1);
	if (higher == 0)
		return (0);
	assert(bitmap != 0x8000);
	lobitmap = higher & -higher;
	loidx = popcount(prev & (lobitmap - 1));
	memmove(&y->branch[loidx + 1], &y->branch[loidx],
		(len - loidx - 1) * sizeof(*y->branch));

	return (0);
}

static struct qp_y *
y_dup(struct qp_y *y0, unsigned short len)
{
	struct qp_y *y;

	assert(len < y0->len);

	y = y_alloc(y0->idx, y0->off + len, y0->len - len);
	if (y == NULL)
		return (NULL);

	y->bitmap = y0->bitmap;
	y->hinib = y0->hinib;
	y->term = y0->term;
	y->branch = y0->branch;
	return (y);
}

static inline uint16_t
getbits(const struct qp_y * const restrict y, unsigned char c)
{
	unsigned shift = y->hinib << 2;
	unsigned mask = 0x0f << shift;
	return (1 << ((c & mask) >> shift));
}

static inline uint8_t
getidx(const struct qp_y * const restrict y, uint16_t bitmap)
{
	return (popcount(y->bitmap & (bitmap - 1)));
}

int
QP_Insert(struct qp_y * * restrict root, unsigned idx,
	  char * const restrict * const restrict strings,
	  unsigned allow_overlaps)
{
	struct qp_y *y;
	unsigned char *c, *b;

	AN(root);
	CHECK_OBJ_ORNULL(*root, QP_Y_MAGIC);
	AN(strings);
	AN(strings[idx]);

	if (*root == NULL) {
		*root = y_alloc(idx, 0, strlen(strings[idx]));
		if (*root == NULL)
			return (-1);
		(*root)->term = 1;
		return (0);
	}

	y = *root;
	b = (unsigned char *)strings[idx];
	c = b;
	errno = 0;

	for (;;) {
		unsigned short i;
		unsigned char *s;
		uint16_t bitmap;
		uint8_t n;
		struct qp_y *y_new, *y_old;

		CHECK_OBJ(y, QP_Y_MAGIC);

		s = (unsigned char *)(strings[y->idx] + y->off);
		for (i = 0; *c != '\0' && i < y->len && s[i] == *c; i++)
			c++;

		if (s[i] == '\0' && *c == '\0') {
			/*
			 * The string to be inserted is already in the
			 * trie.
			 */
			assert(i == y->len);
			errno = EINVAL;
			return (-1);
		}

		if (!allow_overlaps && i == y->len && y->term) {
			/*
			 * The current node is terminating and has a
			 * prefix in common with the current string,
			 * reject if overlaps are not permitted.
			 */
			errno = EPERM;
			return (-1);
		}

		if (i == y->len && y->branch != NULL) {
			/*
			 * The string to be inserted has a prefix that is
			 * already in the trie.
			 */
			AN(y->bitmap);
			bitmap = getbits(y, *c);
			if ((y->bitmap & bitmap) != 0) {
				/*
				 * Other strings in the trie have the same
				 * prefix, follow the branch.
				 */
				n = getidx(y, bitmap);
				ANIB(n);
				AN(y->branch[n]);
				y = y->branch[n];
				continue;
			}
		}

		assert(s[i] != *c);
		AN(*c);

		if (i == y->len && y->branch != NULL) {
			AN(y->bitmap);
			bitmap = getbits(y, *c);
			y_new = y_leaf_alloc(idx, c, b);
			if (y_new == NULL)
				return (-1);
			y->bitmap |= bitmap;
			if (y_realloc_branch(y, bitmap) != 0) {
				FREE_OBJ(y_new);
				return (-1);
			}
			n = getidx(y, bitmap);
			ANIB(n);
			y->branch[n] = y_new;
			return (0);
		}

		if (y->branch == NULL) {
			/*
			 * Current node is a leaf.
			 */
			AZ(y->bitmap);

			y_new = y_leaf_alloc(idx, c, b);
			if (y_new == NULL)
				return (-1);
			y_old = NULL;
			if (s[i] != '\0') {
				y_old = y_dup(y, i);
				if (y_old == NULL) {
					FREE_OBJ(y_new);
					return (-1);
				}
			}

			y->hinib = ((s[i] ^ *c) & 0xf0) != 0;

			bitmap = getbits(y, *c);
			y->bitmap = bitmap;
			errno = 0;
			y->branch = malloc(sizeof(*y->branch));
			if (y->branch == NULL) {
				FREE_OBJ(y_new);
				return (-1);
			}
			AZ(getidx(y, bitmap));
			y->branch[0] = y_new;

			if (s[i] == '\0') {
				/*
				 * The current node is a proper prefix, so
				 * we're done after adding the leaf.
				 */
				assert(i == y->len);
				return (0);
			}

			/*
			 * Move the current node down as a branch.
			 */
			if (y->branch != NULL)
				y->term = 0;
			bitmap = getbits(y, s[i]);
			y->bitmap |= bitmap;
			if (y_realloc_branch(y, bitmap) != 0) {
				FREE_OBJ(y_old);
				FREE_OBJ(y_new);
				return (-1);
			}
			n = getidx(y, bitmap);
			ANIB(n);
			y->branch[n] = y_old;

			y->len = i;
			return (0);
		}

		AN(y->bitmap);
		assert(i < y->len);
		AN(s[i]);
		y_new = y_leaf_alloc(idx, c, b);
		if (y_new == NULL)
			return (-1);
		y_old = y_dup(y, i);
		if (y_old == NULL) {
			FREE_OBJ(y_new);
			return (-1);
		}

		y->hinib = ((s[i] ^ *c) & 0xf0) != 0;
		y->term = 0;

		bitmap = getbits(y, *c);
		y->bitmap = bitmap;
		errno = 0;
		y->branch = malloc(sizeof(*y->branch));
		if (y->branch == NULL) {
			FREE_OBJ(y_new);
			return (-1);
		}
		AZ(getidx(y, bitmap));
		y->branch[0] = y_new;

		bitmap = getbits(y, s[i]);
		y->bitmap |= bitmap;
		if (y_realloc_branch(y, bitmap) != 0) {
			FREE_OBJ(y_old);
			FREE_OBJ(y_new);
			return (-1);
		}
		n = getidx(y, bitmap);
		ANIB(n);
		y->branch[n] = y_old;

		y->len = i;
		return (0);
	}
}

unsigned
QP_Lookup(const struct qp_y * const restrict root,
	  char * const restrict * const restrict strings,
	  const char * const restrict subject)
{
	const struct qp_y *y;
	size_t len;

	AN(strings);
	AN(subject);
	if (root == NULL)
		return UINT_MAX;
	len = strlen(subject);

	for (y = root;;) {
		size_t l;
		uint16_t bitmap;
		uint8_t idx;

		CHECK_OBJ(y, QP_Y_MAGIC);
		l = y->off + y->len;
		if (l > len)
			return UINT_MAX;
		if (y->branch == NULL)
			break;
		bitmap = getbits(y, subject[l]);
		if ((y->bitmap & bitmap) == 0)
			break;
		idx = getidx(y, bitmap);
		ANIB(idx);
		y = y->branch[idx];
		AN(y);
	}

	if (strcmp(subject, strings[y->idx]) == 0)
		return y->idx;
	return (UINT_MAX);
}

static inline int
update_match(struct match_data * const match, unsigned idx, size_t len,
	     size_t l)
{
	if (match->n == match->limit)
		return (-1);
	match->indices[match->n] = idx;
	match->n++;
	if (idx < match->min)
		match->min = idx;
	if (idx > match->max)
		match->max = idx;
	if (l == len)
		match->exact = idx;
	return (0);
}

int
QP_Prefixes(const struct qp_y * const restrict root,
	    char * const restrict * const restrict strings,
	    const char * const restrict subject,
	    struct match_data * const restrict match)
{
	size_t len;

	CHECK_OBJ_NOTNULL(match, MATCH_DATA_MAGIC);
	AN(match->indices);
	AN(match->limit);
	AN(strings);
	AN(subject);

	match->n = 0;
	if (root == NULL)
		return (0);

	match->min = UINT_MAX;
	match->max = 0;
	match->exact = UINT_MAX;

	len = strlen(subject);
	for (const struct qp_y *y = root;;) {
		size_t l;
		uint16_t bitmap;
		int idx = -1;

		CHECK_OBJ(y, QP_Y_MAGIC);
		l = y->off + y->len;
		if (l > len)
			return (0);
		if (y->term) {
			if (strncmp(subject, strings[y->idx], l) != 0)
				return (0);
			if (update_match(match, y->idx, len, l) != 0)
				return (-1);
			if (l == len)
				return (0);
		}
		if (y->branch == NULL)
			return (0);
		bitmap = getbits(y, subject[l]);
		if ((y->bitmap & bitmap) != 0) {
			idx = getidx(y, bitmap);
			ANIB(idx);
		}

		if (idx == -1)
			return (0);
		y = y->branch[idx];
		AN(y);
	}
}

void
QP_Free(struct qp_y *y)
{
	if (y == NULL)
		return;
	CHECK_OBJ(y, QP_Y_MAGIC);
	if (y->branch != NULL) {
		AN(y->bitmap);
		for (int i = 0; i < popcount(y->bitmap); i++) {
			AN(y->branch[i]);
			QP_Free(y->branch[i]);
		}
		free(y->branch);
	}
	FREE_OBJ(y);
}

static void
qp_print_tree(struct qp_y *y, struct vsb *sb, char **strings)
{
	CHECK_OBJ_NOTNULL(y, QP_Y_MAGIC);
	CHECK_OBJ_NOTNULL(sb, VSB_MAGIC);

	VSB_printf(sb, "node = %p\n", y);
	VSB_printf(sb, "idx = %u\n", y->idx);
	VSB_printf(sb, "off = %u\n", y->off);
	VSB_printf(sb, "len = %u\n", y->len);
	AN(strings[y->idx]);
	VSB_printf(sb, "strings[idx] = %s\n", strings[y->idx]);
	VSB_printf(sb, "strings[idx][0]..[off] = %.*s\n", y->off,
		   strings[y->idx]);
	VSB_printf(sb, "strings[idx][off]..[off+len] = %.*s\n", y->len,
		   strings[y->idx] + y->off);
	VSB_printf(sb, "bitmap = 0x%04x\n", y->bitmap);
	VSB_printf(sb, "hinib = %d\n", y->hinib);
	VSB_printf(sb, "term = %d\n", y->term);
	VSB_printf(sb, "branch = %p\n", y->branch);
	VSB_printf(sb, "branches = %d\n", popcount(y->bitmap));
	if (y->bitmap != 0) {
		VSB_printf(sb, "next nibbles = ");
		for (int i = 0; i < 16; i++)
			if (y->bitmap & (1 << i))
				VSB_printf(sb, "%x ", i);
		VSB_printf(sb, "\n");
		AN(y->branch);
		for (int i = 0; i < popcount(y->bitmap); i++)
			VSB_printf(sb, "branch[%d] = %p\n", i, y->branch[i]);
	}
	VSB_printf(sb, "\n");

	if (y->bitmap != 0)
		for (int i = 0; i < popcount(y->bitmap); i++)
			qp_print_tree(y->branch[i], sb, strings);
}

struct vsb *
QP_Dump(struct qp_y *root, char **strings)
{
	struct vsb *sb = VSB_new_auto();


	VSB_printf(sb, "root = %p\n\n", root);
	if (root != NULL) {
		AN(strings);
		qp_print_tree(root, sb, strings);
	}
	VSB_finish(sb);
	return (sb);
}

void
qp_stats(const struct qp_y * const restrict y,
	 char * const restrict * const restrict strings,
	 struct qp_stats * const restrict stats, unsigned depth)
{
	uint8_t fanout;

	if (y == NULL)
		return;

	CHECK_OBJ(y, QP_Y_MAGIC);
	depth++;
	stats->nodes++;

	if (strings[y->idx][y->off + y->len] == '\0') {
		if (depth < stats->dmin)
			stats->dmin = depth;
		if (depth > stats->dmax)
			stats->dmax = depth;
		stats->davg += (depth - stats->davg) / (stats->terms + 1.);
		stats->terms++;
	}

	if (y->bitmap == 0) {
		AZ(y->branch);
		stats->leaves++;
		return;
	}

	AN(y->branch);
	fanout = popcount(y->bitmap);
	assert(fanout <= 16);
	if (fanout < stats->fmin)
		stats->fmin = fanout;
	if (fanout > stats->fmax)
		stats->fmax = fanout;
	stats->favg +=
		(fanout - stats->favg) / ((stats->nodes - stats->leaves));

	for (int i = 0; i < popcount(y->bitmap); i++)
		qp_stats(y->branch[i], strings, stats, depth);
}

void
QP_Stats(const struct qp_y * const restrict root,
	 char * const restrict * const restrict strings,
	 struct qp_stats * const restrict stats)
{
	CHECK_OBJ_NOTNULL(stats, QP_STATS_MAGIC);

	stats->nodes = 0;
	stats->leaves = 0;
	stats->terms = 0;
	stats->dmin = UINT64_MAX;
	stats->dmax = 0;
	stats->davg = 0.;
	stats->fmin = UINT64_MAX;
	stats->fmax = 0;
	stats->favg = 0.;
	stats->nodesz = sizeof(*root);

	qp_stats(root, strings, stats, 0);
	if (stats->fmin == UINT64_MAX)
		stats->fmin = 0;
}
