#include <err.h>
#include <errno.h>
#include <inttypes.h>
#include <limits.h>
#include <math.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <time.h>
#include <sys/mman.h>

#define	ARRSZ	(16*1024*1024)

struct snprintf_arg {
        char    *str;
        size_t  remain;
};

static int
snprintf_func(const char *p, int len, void *arg)
{
        struct snprintf_arg *const info = arg;
	int copy;

	copy = len;
	if (info->remain - copy < 2)
		copy = info->remain - 2;

	strncpy(info->str, p, copy);
	info->str += copy;
	info->remain -= copy;

        if (info->remain >= 2)
		return (0);
	errno = ENOSPC;
	return (-1);
}

static int
my_sprintf(char *p, const char *fmt, ...)
{
        struct snprintf_arg info;
	va_list ap;
	int rc;

	/* Nothing to do on flush. */
	if (p == NULL)
		return (0);

	va_start(ap, fmt);
	info.str = p;
        info.remain = INT_MAX;
        rc = vcb_printf(snprintf_func, &info, fmt, ap);
        if (info.remain >= 1)
                *info.str++ = '\0';
	va_end(ap);
        return (rc);
}

struct __siov {
	void	*iov_base;
	size_t	iov_len;
};
struct __suio {
	struct	__siov *uio_iov;
	int	uio_iovcnt;
	int	uio_resid;
};
#define NIOV 8
struct  io_state {
	FILE *fp;
	struct __suio uio;	/* output information: summary */
	struct __siov iov[NIOV];/* ... and individual io vectors */
};
extern int __sfvwrite(FILE *fp, struct __suio *uio);

/*
 * Flush out all the vectors defined by the given uio,
 * then reset it so that it can be reused.
 */
static int
__sprint(FILE *fp, struct __suio *uio)
{
	int err;

	if (uio->uio_resid == 0) {
		uio->uio_iovcnt = 0;
		return (0);
	}
	err = __sfvwrite(fp, uio);
	uio->uio_resid = 0;
	uio->uio_iovcnt = 0;
	if (err != 0)
		printf("returning %d\n", err);
	return (err);
}

static inline void
io_init(struct io_state *iop, FILE *fp)
{

	iop->uio.uio_iov = iop->iov;
	iop->uio.uio_resid = 0;
	iop->uio.uio_iovcnt = 0;
	iop->fp = fp;
}

static int
fprintf_func(const char *ptr, int len, void *arg)
{
	struct io_state *iop = arg;

	if (ptr == NULL)
		return (__sprint(iop->fp, &iop->uio));

	iop->iov[iop->uio.uio_iovcnt].iov_base = (char *)ptr;
	iop->iov[iop->uio.uio_iovcnt].iov_len = len;
	iop->uio.uio_resid += len;
	if (++iop->uio.uio_iovcnt >= NIOV)
		return (__sprint(iop->fp, &iop->uio));
	else
		return (0);
}

static int
my_fprintf(FILE *fp, const char *fmt, ...)
{
	struct io_state io;
	va_list ap;
	int rc;

	io_init(&io, fp);

	va_start(ap, fmt);
        rc = vcb_printf(fprintf_func, &io, fmt, ap);
	va_end(ap);
        return (rc);
}


static int
snprintf2_func(int ch, void *arg)
{
        struct snprintf_arg *const info = arg;

	if (info->remain >= 2) {
		*info->str++ = ch;
		info->remain--;
		return (0);
	}
	errno = ENOSPC;
	return (-1);
}

static int
my_sprintf2(char *p, const char *fmt, ...)
{
        struct snprintf_arg info;
	va_list ap;
	int rc;

	va_start(ap, fmt);
	info.str = p;
        info.remain = INT_MAX;
        rc = vcb_printf2(snprintf2_func, &info, fmt, ap);
        if (info.remain >= 1)
                *info.str++ = '\0';
	va_end(ap);
        return (rc);
}


static int
fprintf2_func(int ch, void *arg)
{
        FILE *fp = arg;
	int rc;

	rc = putc_unlocked(ch, fp);
	if (rc != ch) {
		errno = EINVAL;
		return (-1);
	}
	return (0);
}

static int
my_fprintf2(FILE *fp, const char *fmt, ...)
{
	va_list ap;
	int rc;

	va_start(ap, fmt);
        rc = vcb_printf2(fprintf2_func, fp, fmt, ap);
	va_end(ap);
        return (rc);
}

/*------------------------------------------------------------------*/

static uintmax_t
trial(int (*spf)(char *, const char *, ...), char *p)
{
	struct timespec start, end;
	int i, rc, tot;
	uintmax_t nsec;

	rc = clock_gettime(CLOCK_REALTIME_FAST, &start);
	if (rc != 0)
		err(1, "gettime - start");

	for (tot = 0, i = 0; i < ARRSZ / 64; i++) {
		rc = spf(p + tot, "%-64s", "abc123");
		if (rc != 64)
			errx(1, "Got rc %d, not 64", rc);
		tot += rc;
	}

	rc = clock_gettime(CLOCK_REALTIME_FAST, &end);
	if (rc != 0)
		err(1, "gettime - end");

	nsec = (end.tv_sec - start.tv_sec) * 1000000000;
	nsec += (end.tv_nsec - start.tv_nsec);
	return (nsec);
}

static uintmax_t
trial2(int (*spf)(FILE *, const char *, ...), FILE *fp)
{
	struct timespec start, end;
	int i, rc, tot;
	uintmax_t nsec;

	rc = clock_gettime(CLOCK_REALTIME_FAST, &start);
	if (rc != 0)
		err(1, "gettime - start");

	for (tot = 0, i = 0; i < ARRSZ / 64; i++) {
		rc = spf(fp, "%-64s", "abc123");
		if (rc != 64)
			errx(1, "Got rc %d, not 64", rc);
		tot += rc;
	}

	rc = clock_gettime(CLOCK_REALTIME_FAST, &end);
	if (rc != 0)
		err(1, "gettime - end");

	nsec = (end.tv_sec - start.tv_sec) * 1000000000;
	nsec += (end.tv_nsec - start.tv_nsec);
	return (nsec);
}

/*------------------------------------------------------------------*/

#define	NTRIALS	5

static uintmax_t
do_avg(uintmax_t nsec[NTRIALS])
{
	uintmax_t sum = 0;
	int i;

	for (i = 0; i < NTRIALS; i++)
		sum += nsec[i];
	return (sum / NTRIALS);
}

static double
do_std(uintmax_t avg, uintmax_t nsec[NTRIALS])
{
	double sum2 = 0.0;
	int i;

	/* Work in seconds, not ns */
	for (i = 0; i < NTRIALS; i++) {
		double diff = (avg - nsec[i]) / 1e9;

		sum2 += diff * diff;
	}
	return (sqrt(sum2 / NTRIALS));
}

static void
do_report(char *name, uintmax_t nsec[NTRIALS])
{
	uintmax_t avg = do_avg(nsec);
	double stdd = do_std(avg, nsec);

	printf("%s: avg %jd.%09jd sec, stdd %6.4g\n",
	    name, avg / 1000000000, avg % 1000000000, stdd);
}


int
main(void)
{
	char *p;
	int i, rc;
	FILE *fp;
	uintmax_t nsec0[NTRIALS];
	uintmax_t nsec1[NTRIALS];
	uintmax_t nsec2[NTRIALS];

	rc = mlockall(MCL_CURRENT | MCL_FUTURE);
	if (rc != 0)
		err(1, "mlockall");

	p = malloc(ARRSZ + 100);
	if (p == NULL)
		err(1, "malloc");

	/* Make sure all pages are memory resident. */
	bzero(p, ARRSZ);

	for (i = 0; i < NTRIALS; i++)
		nsec0[i] = trial(sprintf, p);
	for (i = 0; i < NTRIALS; i++)
		nsec1[i] = trial(my_sprintf, p);
	for (i = 0; i < NTRIALS; i++)
		nsec2[i] = trial(my_sprintf2, p);

	do_report("sprintf    ", nsec0);
	do_report("my_sprintf ", nsec1);
	do_report("my_sprintf2", nsec2);

	/* Writing to /dev/null. */
	fp = fopen("/dev/null", "w+");
	for (i = 0; i < NTRIALS; i++)
		nsec0[i] = trial2(fprintf, fp);
	for (i = 0; i < NTRIALS; i++)
		nsec1[i] = trial2(my_fprintf, fp);
	for (i = 0; i < NTRIALS; i++)
		nsec2[i] = trial2(my_fprintf2, fp);

	do_report("fprintf     /dev/null", nsec0);
	do_report("my_fprintf  /dev/null", nsec1);
	do_report("my_fprintf2 /dev/null", nsec2);

	/* Writing to /tmp. */
	fp = fopen("/tmp/temp_file", "w");
	for (i = 0; i < NTRIALS; i++)
		nsec0[i] = trial2(fprintf, fp);
	fclose(fp);
	fp = fopen("/tmp/temp_file", "w");
	for (i = 0; i < NTRIALS; i++)
		nsec1[i] = trial2(my_fprintf, fp);
	fclose(fp);
	fp = fopen("/tmp/temp_file", "w");
	for (i = 0; i < NTRIALS; i++)
		nsec2[i] = trial2(my_fprintf2, fp);
	fclose(fp);

	do_report("fprintf     /tmp", nsec0);
	do_report("my_fprintf  /tmp", nsec1);
	do_report("my_fprintf2 /tmp", nsec2);

	return (0);
}