/*	$OpenBSD$ */

/*
 * Copyright (c) 2014 David Gwynne <dlg@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/time.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

#include <sys/select.h>
#include <poll.h>

int     uselect(int, fd_set * __restrict, fd_set * __restrict,
	    fd_set * __restrict, struct timeval * __restrict);
int     upselect(int, fd_set * __restrict, fd_set * __restrict,
	    fd_set * __restrict, const struct timespec * __restrict,
	    const sigset_t * __restrict);

static int
_fd_isset(int fd, fd_set *fds)
{
	if (fds == NULL)
		return (0);

	return (FD_ISSET(fd, fds));
}

static int
_fd_set(fd_set *fds, struct pollfd *pfd, short mask)
{
	if (fds == NULL || (pfd->events & mask) == 0 ||
	    (pfd->revents & mask) == 0)
		return (0);

	FD_SET(pfd->fd, fds);
	return (1);
}

int
uselect(int nfds, fd_set *rfds, fd_set *wfds, fd_set *efds, struct timeval *tv)
{
	struct timespec *tsp = NULL, ts;

	if (tv != NULL) {
		tsp = &ts;
		TIMEVAL_TO_TIMESPEC(tv, tsp);
	}

	return (upselect(nfds, rfds, wfds, efds, tsp, NULL));
}

int
upselect(int nfds, fd_set *rfds, fd_set *wfds, fd_set *efds,
    const struct timespec *ts, const sigset_t *mask)
{
	struct pollfd *pfds = NULL, *pfd;
	int pfdlen = 0, npfds = 0;
	int i;
	int serrno;
	int rv = -1;
	short events;

	if (nfds < 0 || nfds > FD_SETSIZE) {
		errno = EINVAL;
		return (-1);
	}

	for (i = 0; i < nfds; i++) {
		events = 0;
		if (_fd_isset(i, rfds))
			events |= POLLIN;
		if (_fd_isset(i, wfds))
			events |= POLLOUT;
		if (_fd_isset(i, efds)) /* XXX */
			events |= POLLPRI;

		if (events == 0)
			continue;

		if (npfds >= pfdlen) {
			pfdlen += 4; /* greater than one, two, or many */
			pfd = reallocarray(pfds, pfdlen, sizeof(*pfd));
			if (pfd == NULL)
				goto done;

			pfds = pfd;
		}

		pfd = &pfds[npfds++];
		pfd->fd = i;
		pfd->events = events;
		pfd->revents = 0;
	}

	rv = ppoll(pfds, npfds, ts, mask);
	if (rv == -1)
		goto done;

	if (rv > npfds)
		abort();

	i = howmany(nfds, NFDBITS) * sizeof(fd_mask);
	if (rfds != NULL)
		memset(rfds, 0, i);
	if (wfds != NULL)
		memset(wfds, 0, i);
	if (efds != NULL)
		memset(efds, 0, i);

	npfds = rv;
	rv = 0;

	pfd = pfds;
	for (i = 0; i < npfds; i++) {
		/* hopefully revents are consistent with poll()s rv */
		while (pfd->revents == 0)
			pfd++;

		rv += _fd_set(rfds, pfd, POLLIN);
		rv += _fd_set(wfds, pfd, POLLOUT);
		rv += _fd_set(efds, pfd, POLLPRI); /* XXX */
	}

done:
	serrno = errno;
	free(pfds);
	errno = serrno;
	return (rv);
}