path: root/tquery.c

#include <curses.h>

#include <fcntl.h>
#include <stdint.h>
#include <sys/wait.h>
#include <sys/poll.h>
#include <sys/fcntl.h>
#include <signal.h>
#include <dirent.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>

#define LOAD_BUFSIZ 16384
#define INPUT_BUFSIZ 2048

#define CTRL(x) ((x) & 0x1f)

#define MIN(a, b) ((a) > (b) ? (b) : (a))
#define MAX(a, b) ((a) > (b) ? (a) : (b))

struct nav {
	ssize_t wmin, wmax, wlen;
	ssize_t sel, min, max;
};

struct entry {
	size_t full;
	size_t visible;
};

static void sigint_handler(int sig);

static const struct sigaction sigint_action = {
	.sa_flags = SA_RESTART,
	.sa_handler = sigint_handler
};

static struct sigaction sigwinch_action = {
	.sa_flags = SA_RESTART,
	.sa_handler = SIG_IGN
};

static int child_fd = -1;
static pid_t child_pid = 0;
static const char **child_argv = NULL;

static pid_t hook_pid = -1;
static char hook_key_arg[6] = { 0 };
static char *hook_argv[4] = { NULL, hook_key_arg, NULL, NULL };
static bool hook_io = false;

static char *run_argv[3] = { NULL, NULL, NULL };
static bool run_cmd = false;
static bool run_io = false;

static char *loadbuf = NULL;

static char *output = NULL;
static size_t output_len = 0;
static size_t output_cap = 0;

static struct entry *entries = NULL;
static size_t entries_len = 0;
static size_t entries_cap = 0;

static struct nav output_nav;

static char *inputbuf = NULL;
static size_t inputlen = 0;

static bool split_args = false;
static bool oneshot = false;
static bool errlog = false;
static char delim = '\n';
static char vis_delim = '|';

static void
die(const char *fmt, ...)
{
	va_list ap;

	fputs("tquery: ", stderr);

	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);

	if (fmt[0] && fmt[strlen(fmt)-1] == ':') {
		fputc(' ', stderr);
		perror(NULL);
	} else {
		fputc('\n', stderr);
	}

	exit(1);
}

static void
nav_update_win(struct nav *nav)
{
	nav->wmin = MAX(nav->wmax - nav->wlen, nav->min);
	nav->wmax = MIN(nav->wmin + nav->wlen, nav->max);
	nav->sel = MAX(MIN(nav->sel, nav->wmax - 1), nav->wmin);
}

static void
nav_init(struct nav *nav)
{
	nav->sel = 0;
	nav->min = 0;
	nav->max = 0;
	nav->wlen = 0;
	nav->wmin = 0;
	nav->wmax = 0;
}

static void
nav_update_bounds(struct nav *nav, int min, int max)
{
	nav->min = min;
	nav->max = max;
	nav_update_win(nav);
}

static void
nav_update_wlen(struct nav *nav, ssize_t wlen)
{
	nav->wlen = wlen;
	nav_update_win(nav);
}

static void
nav_update_sel(struct nav *nav, ssize_t sel)
{
	nav->sel = MAX(MIN(sel, nav->max - 1), nav->min);
	if (nav->sel >= nav->wmax) {
		nav->wmax = nav->sel + 1;
		nav->wmin = MAX(nav->min, nav->wmax - nav->wlen);
	} else if (nav->sel < nav->wmin) {
		nav->wmin = nav->sel;
		nav->wmax = MIN(nav->wmin + nav->wlen, nav->max);
	}
}

static void
invoke(const char **argv, pid_t *pid, int *fd, bool in, bool out, bool err)
{
	int out_pipe[2];

	if (fd) {
		if (*fd >= 0) {
			close(*fd);
			*fd = -1;
		}
		if (pipe(out_pipe))
			die("pipe:");
	}

	*pid = fork();
	if (*pid < 0) die("fork:");
	if (!*pid) {
		if (child_fd) close(child_fd);
		if (!in || !out || !err) {
			int zfd = open("/dev/null", O_RDWR);
			if (zfd < 0) die("open /dev/null");
			if (!in) dup2(zfd, 0);
			if (!out) dup2(zfd, 1);
			if (!err) dup2(zfd, 2);
			close(zfd);
		}
		if (fd) {
			if (out) dup2(out_pipe[1], 1);
			close(out_pipe[0]);
			close(out_pipe[1]);
		}
		execvp(argv[0], (char *const *) argv);
		die("execv '%s':", argv[0]);
	} else {
		if (fd) {
			*fd = out_pipe[0];
			close(out_pipe[1]);
		}
	}
}

static void *
addcap(void *alloc, size_t dsize, size_t min, size_t *cap)
{
	if (min > *cap) {
		*cap = *cap * 2;
		if (*cap < min) *cap = min;
		alloc = realloc(alloc, dsize * *cap);
		if (!alloc) die("realloc:");
	}

	return alloc;
}

static void
load(void)
{
	ssize_t nread;
	size_t i;
	char *c;

	nread = read(child_fd, loadbuf, LOAD_BUFSIZ);
	if (nread <= 0) {
		child_pid = 0;
		close(child_fd);
		child_fd = -1;
		return;
	}

	output = addcap(output, 1, output_len + (size_t) nread + 1, &output_cap);
	memcpy(output + output_len, loadbuf, (size_t) nread);

	if (!entries_cap)
		entries = addcap(entries, sizeof(struct entry), 1, &entries_cap);

	if (!entries_len) {
		entries[0].full = 0;
		entries[0].visible = 0;
		entries_len = 1;
	}

	c = output + output_len;
	for (i = output_len; i < output_len + (size_t) nread; i++, c++) {
		if (*c == vis_delim) {
			entries[entries_len-1].visible = i+1;
		} else if (*c == delim) {
			*c = '\0';
			entries = addcap(entries, sizeof(struct entry),
				entries_len + 1, &entries_cap);
			entries[entries_len].full = i + 1;
			entries[entries_len].visible = i + 1;
			entries_len++;
		}
	}

	output_len += (size_t) nread;
	output[output_len] = '\0';
}

static size_t
num_entries()
{
	if (!entries_len) return 0;
	if (output_len <= entries[entries_len-1].full)
		return entries_len-1;
	return entries_len;
}

static char *
full_entry(size_t i)
{
	return output + entries[i].full;
}

static char *
vis_entry(size_t i)
{
	return output + entries[i].visible;
}

static void
unload(void)
{
	int status;

	if (!child_pid) return;

	kill(child_pid, SIGKILL);
	do {
		waitpid(child_pid, &status, 0);
	} while (!WIFEXITED(status) && !WIFSIGNALED(status));

	close(child_fd);
	child_fd = -1;
	child_pid = 0;
}

static void
update(void)
{
	const ssize_t miny = 3;
	int width, height;
	size_t entry_cnt;
	ssize_t i, y;

	erase();

	width = getmaxx(stdscr);
	height = getmaxy(stdscr);

	entry_cnt = num_entries();
	mvprintw(1, 1, "%lu %li %lu %i", output_len,
		output_nav.sel, entry_cnt, child_pid);
	mvprintw(2, 1, "> %.*s", width - 1, inputbuf);

	nav_update_bounds(&output_nav, 0, (int) entry_cnt);
	nav_update_wlen(&output_nav, (int) (height - miny));

	if (output_len) {
		i = output_nav.wmin;
		for (y = miny; y < height && i < output_nav.wmax; y++, i++) {
			if (i == output_nav.sel)
				attr_on(A_REVERSE, NULL);

			mvprintw((int) y, 1, "%.*s", width - 1, vis_entry((size_t) i));

			if (i == output_nav.sel)
				attr_off(A_REVERSE, NULL);
		}
	}

	move(2, (int) (1 + 2 + inputlen));

	refresh();
}

static void
spawn(const char *query)
{
	static const char *argv[64];
	const char **arg;
	const char *tok, *ntok, *end;
	size_t argc;

	entries_len = 0;
	output_len = 0;

	argc = 0;
	for (arg = child_argv; *arg && argc < 64; arg++) {
		if (!strcmp(*arg, "{}")) {
			arg += 1;
			break;
		}
		argv[argc++] = *arg;
	}

	if (split_args) {
		ntok = tok = query;
		while (ntok && *ntok && argc < 64) {
			ntok = strchr(tok, ' ');
			end = ntok ? ntok : tok + strlen(tok);
			argv[argc++] = strndup(tok, (size_t) (end - tok));
			tok = ntok + 1;
		}
	} else if (argc < 64 && query) {
		argv[argc++] = query;
	}

	for (; *arg && argc < 64; arg++) {
		argv[argc++] = *arg;
	}

	if (argc < 64) {
		argv[argc++] = NULL;
	} else {
		die("Too many arguments");
	}

	if (child_pid) unload();
	invoke(argv, &child_pid, &child_fd, false, true, errlog);
}

static void
run(char *input)
{
	if (run_io) {
		def_prog_mode();
		endwin();
	}

	int status, pid;
	run_argv[1] = input;
	invoke((const char **)run_argv, &pid, NULL, run_io, run_io, errlog);
	do {
		int rc = waitpid(pid, &status, 0);
		if (rc != pid) die("waitpid:");
	} while (!WIFEXITED(status) && !WIFSIGNALED(status));

	if (oneshot) exit(0);

	if (run_io) {
		reset_prog_mode();
	}

	if (WIFEXITED(status) && WEXITSTATUS(status) == 3)
		spawn(inputbuf);
	update();
}

static void
hook(uint8_t key, const char *line)
{
	int status, rc;

	if (!hook_argv[0]) return;

	snprintf(hook_argv[1], 6, "%u", key);
	free(hook_argv[2]);
	hook_argv[2] = strdup(line);

	if (hook_io) {
		def_prog_mode();
		endwin();
	}

	invoke((const char **)hook_argv, &hook_pid, NULL, hook_io, hook_io, errlog);
	do {
		rc = waitpid(hook_pid, &status, 0);
		if (rc != hook_pid) die("waitpid:");
	} while (!WIFEXITED(status) && !WIFSIGNALED(status));

	if (hook_io) {
		reset_prog_mode();
	}

	if (WIFEXITED(status) && WEXITSTATUS(status) == 3)
		spawn(inputbuf);
	update();
}

static void
input(void)
{
	bool dirty;
	bool reload;
	int c;

	dirty = false;
	reload = false;
	while ((c = getch()) != ERR) {
		switch (c) {
		case KEY_BACKSPACE:
			if (inputlen) inputlen--;
			dirty = true;
			reload = true;
			break;
		case KEY_PPAGE:
			nav_update_sel(&output_nav,
				output_nav.sel - output_nav.wlen / 2);
			dirty = true;
			break;
		case KEY_NPAGE:
			nav_update_sel(&output_nav,
				output_nav.sel + output_nav.wlen / 2);
			dirty = true;
			break;
		case KEY_DOWN:
			nav_update_sel(&output_nav, output_nav.sel + 1);
			dirty = true;
			break;
		case KEY_UP:
			nav_update_sel(&output_nav, output_nav.sel - 1);
			dirty = true;
			break;
		case '\x1b':
			/* TODO: word-wise seek */
			switch (getch()) {
			case 'c':
				break;
			case 'd':
				break;
			}
			break;
		case CTRL('w'):
			inputlen = 0;
			dirty = true;
			reload = true;
			break;
		case CTRL('l'):
			clear();
			dirty = true;
			break;
		case CTRL('b'):
			nav_update_sel(&output_nav, output_nav.min);
			dirty = true;
			break;
		case CTRL('g'):
			nav_update_sel(&output_nav, output_nav.max-1);
			dirty = true;
			break;
		case CTRL('x'):
			if (output_nav.max > 0 && output_nav.sel >= 0) {
				timeout(-1);
				hook((uint8_t)getch(), full_entry((size_t)output_nav.sel));
				timeout(0);
			}
			break;
		case '\n':
			if (output_nav.max > 0 && output_nav.sel >= 0) {
				if (run_cmd) {
					run(full_entry((size_t) output_nav.sel));
				} else {
					puts(full_entry((size_t) output_nav.sel));
					if (oneshot) exit(0);
				}
			}
			break;
		default:
			if (c < 32 || c >= 127)
				continue;
			if (inputlen < INPUT_BUFSIZ - 1)
				inputbuf[inputlen++] = (char) c;
			reload = true;
			dirty = true;
			break;
		}
		inputbuf[inputlen] = '\0';
	}

	if (dirty) {
		if (reload) spawn(inputbuf);
		update();
	}
}

static void
sigint_handler(int sig)
{
	if (child_pid) {
		unload();
		update();
	} else {
		exit(0);
	}
}

static void
usage(int rc, bool full)
{
	fprintf(stderr, "Usage: tquery [OPT..] -- CMD [ARG..]\n");
	if (full) {
		fprintf(stderr, "\n");
		fprintf(stderr, "  -h, --help        Show this message\n");
		fprintf(stderr, "  -d, --delim       Set input entry delim\n");
		fprintf(stderr, "  -D, --vis-delim   Set visible entry delim\n");
		fprintf(stderr, "  -o, --oneshot     Exit after oneshot selection\n");
		fprintf(stderr, "  -e, --stderr      Dont close child stderr\n");
		fprintf(stderr, "  -s, --split       Split the query into args\n");
		fprintf(stderr, "  -r, --run         Program to run on Enter\n");
		fprintf(stderr, "  -R, --run-io      Program to run interactively on Enter\n");
		fprintf(stderr, "  -x, --hook        Program to run on Ctrl-X\n");
		fprintf(stderr, "  -X, --hook-io     Program to run interactively on Ctrl-X\n");
		fprintf(stderr, "\n");
	}
	exit(rc);
}

static void
parse(int argc, char **argv)
{
	char **arg;

	for (arg = &argv[1]; *arg; arg++) {
		if (!strcmp(*arg, "-h") || !strcmp(*arg, "--help")) {
			usage(0, true);
		} else if (!strcmp(*arg, "-o") || !strcmp(*arg, "--oneshot")) {
			oneshot = true;
		} else if (!strcmp(*arg, "-e") || !strcmp(*arg, "--stderr")) {
			errlog = true;
		} else if (!strcmp(*arg, "-s") || !strcmp(*arg, "--split")) {
			split_args = true;
		} else if (!strcmp(*arg, "-d") || !strcmp(*arg, "--delim")) {
			delim = **++arg;
		} else if (!strcmp(*arg, "-D") || !strcmp(*arg, "--vis-delim")) {
			vis_delim = **++arg;
		} else if (!strcmp(*arg, "-r") || !strcmp(*arg, "--run")) {
			run_argv[0] = *++arg;
			run_cmd = true;
		} else if (!strcmp(*arg, "-R") || !strcmp(*arg, "--run-io")) {
			run_argv[0] = *++arg;
			run_cmd = true;
			run_io = true;
		} else if (!strcmp(*arg, "-x") || !strcmp(*arg, "--hook")) {
			hook_argv[0] = *++arg;
		} else if (!strcmp(*arg, "-X") || !strcmp(*arg, "--hook-io")) {
			hook_argv[0] = *++arg;
			hook_io = true;
		} else if (!strcmp(*arg, "--")) {
			child_argv = (const char **)arg + 1;
			break;
		} else {
			die("unknown option '%s'", arg[0]);
		}
	}

	if (!child_argv || !*child_argv)
		usage(0, false);
}

int
main(int argc, char **argv)
{
	struct pollfd fds[2];
	int rc;

	parse(argc, argv);

	nav_init(&output_nav);

	atexit((void (*)()) endwin);
	if (!newterm(NULL, stderr, stdin))
		die("newterm stdin -> stderr");
	cbreak();
	noecho();
	keypad(stdscr, true);

	sigaction(SIGINT, &sigint_action, NULL);
	sigaction(SIGWINCH, NULL, &sigwinch_action);
	sigwinch_action.sa_flags |= SA_RESTART;
	sigaction(SIGWINCH, &sigwinch_action, NULL);

	timeout(0);

	spawn("");

	inputbuf = malloc(INPUT_BUFSIZ);
	if (!inputbuf) die("malloc inputbuf:");

	loadbuf = malloc(LOAD_BUFSIZ);
	if (!loadbuf) die("malloc loadbuf:");

	inputlen = 0;
	inputbuf[inputlen] = '\0';

	update();

	while (1) {
		fds[0].fd = 0;
		fds[0].events = POLLIN;
		fds[1].fd = child_fd;
		fds[1].events = POLLIN;

		rc = poll(fds, child_pid ? 2 : 1, -1);
		if (rc < 0 && errno == EINTR) {
			update();
			continue;
		} else if (rc < 0) {
			die("select:");
		}

		if (fds[0].revents & POLLIN) {
			input();
		} else if (fds[1].revents & POLLIN) {
			load();
			update();
		} else if (fds[1].revents & POLLHUP) {
			unload();
			update();
		}
	}
}