virgo.c

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <shellapi.h>

#define sb_free(a) ((a) ? HeapFree(GetProcessHeap(), 0, stb__sbraw(a)), 0 : 0)
#define sb_push(a, v) (stb__sbmaybegrow(a, 1), (a)[stb__sbn(a)++] = (v))
#define sb_count(a) ((a) ? stb__sbn(a) : 0)

#define stb__sbraw(a) ((int *)(a)-2)
#define stb__sbm(a) stb__sbraw(a)[0]
#define stb__sbn(a) stb__sbraw(a)[1]

#define stb__sbneedgrow(a, n) ((a) == 0 || stb__sbn(a) + (n) >= stb__sbm(a))
#define stb__sbmaybegrow(a, n) (stb__sbneedgrow(a, (n)) ? stb__sbgrow(a, n) : 0)
#define stb__sbgrow(a, n) ((a) = stb__sbgrowf((a), (n), sizeof(*(a))))

#ifndef MOD_NOREPEAT
#define MOD_NOREPEAT 0x4000
#endif

#define NUM_DESKTOPS 9
#define NUM_PINNED 9

typedef struct {
	HWND *windows;
	unsigned count;
} Windows;

typedef struct {
	unsigned current;
	unsigned handle_hotkeys;
	Windows desktops[NUM_DESKTOPS];
	HWND pinned[NUM_PINNED];
	HWND focused[NUM_DESKTOPS];
} Virgo;

static void *stb__sbgrowf(void *arr, unsigned increment, unsigned itemsize)
{
	unsigned dbl_cur = arr ? 2 * stb__sbm(arr) : 0;
	unsigned min_needed = sb_count(arr) + increment;
	unsigned m = dbl_cur > min_needed ? dbl_cur : min_needed;
	unsigned *p;
	if (arr) {
		p = HeapReAlloc(GetProcessHeap(), 0, stb__sbraw(arr),
						itemsize * m + sizeof(unsigned) * 2);
	} else {
		p = HeapAlloc(GetProcessHeap(), 0, itemsize * m + sizeof(unsigned) * 2);
	}
	if (p) {
		if (!arr) {
			p[1] = 0;
		}
		p[0] = m;
		return p + 2;
	} else {
		ExitProcess(1);
		return (void *)(2 * sizeof(unsigned));
	}
}

static void windows_mod(Windows *wins, unsigned state)
{
	unsigned i;
	for (i = 0; i < wins->count; i++) {
		ShowWindow(wins->windows[i], state);
	}
}

static void windows_show(Windows *wins) { windows_mod(wins, SW_SHOW); }

static void windows_hide(Windows *wins) { windows_mod(wins, SW_HIDE); }

static void windows_add(Windows *wins, HWND hwnd)
{
	if (wins->count >= sb_count(wins->windows)) {
		sb_push(wins->windows, hwnd);
	} else {
		wins->windows[wins->count] = hwnd;
	}
	wins->count++;
}

static void windows_del(Windows *wins, HWND hwnd)
{
	unsigned i, e;
	for (i = 0; i < wins->count; i++) {
		if (wins->windows[i] != hwnd) {
			continue;
		}
		if (i != wins->count - 1) {
			for (e = i; e < wins->count - 1; e++) {
				wins->windows[e] = wins->windows[e + 1];
			}
		}
		wins->count--;
		break;
	}
}

static unsigned is_valid_window(HWND hwnd)
{
	WINDOWINFO wi;
	wi.cbSize = sizeof(wi);
	GetWindowInfo(hwnd, &wi);
	return (wi.dwStyle & WS_VISIBLE) && !(wi.dwExStyle & WS_EX_TOOLWINDOW);
}

static void register_hotkey(unsigned id, unsigned mod, unsigned vk)
{
	if (!RegisterHotKey(NULL, id, mod, vk)) {
		MessageBox(NULL, "could not register hotkey", "error",
				   MB_ICONEXCLAMATION);
		ExitProcess(1);
	}
}

static unsigned virgo_is_pinned(Virgo *v, HWND hwnd) {
	unsigned i;
	for (i = 0; i < NUM_PINNED; ++i) {
		if (v->pinned[i] == hwnd) {
			return 1;
		}
	}
	return 0;
}

static void virgo_toggle_pin(Virgo *v, HWND hwnd) {
	unsigned i;
	unsigned empty = NUM_PINNED;
	for (i = 0; i < NUM_PINNED; ++i) {
		if (v->pinned[i] == hwnd) {
			v->pinned[i] = 0;
			windows_add(&v->desktops[v->current], hwnd);
			return;
		}
		if (!v->pinned[i]) {
			empty = i;
		}
	}
	if (empty == NUM_PINNED) {
		MessageBox(NULL, "reached pinned windows limit", "error",
				   MB_ICONEXCLAMATION);
		return;
	}
	v->pinned[empty] = hwnd;
	windows_del(&v->desktops[v->current], hwnd);
}

static BOOL enum_func(HWND hwnd, LPARAM lParam)
{
	unsigned i, e;
	Virgo *v;
	Windows *desk;
	v = (Virgo *)lParam;
	if (!is_valid_window(hwnd)) {
		return 1;
	}
	for (i = 0; i < NUM_DESKTOPS; i++) {
		desk = &(v->desktops[i]);
		for (e = 0; e < desk->count; e++) {
			if (desk->windows[e] == hwnd) {
				return 1;
			}
		}
	}
	if (virgo_is_pinned(v, hwnd)) {
		return 1;
	}
	windows_add(&(v->desktops[v->current]), hwnd);
	return 1;
}

static void virgo_update(Virgo *v)
{
	unsigned i, e;
	Windows *desk;
	HWND hwnd;
	for (i = 0; i < NUM_DESKTOPS; i++) {
		desk = &(v->desktops[i]);
		for (e = 0; e < desk->count; e++) {
			hwnd = desk->windows[e];
			if (!GetWindowThreadProcessId(desk->windows[e], NULL)) {
				windows_del(desk, hwnd);
			}
		}
	}
	for (i = 0; i < NUM_PINNED; i++) {
		hwnd = v->pinned[i];
		if (!GetWindowThreadProcessId(hwnd, NULL)) {
			v->pinned[i] = 0;
		}
	}
	desk = &v->desktops[v->current];
	for (i = 0; i < desk->count; i++) {
		hwnd = desk->windows[i];
		if (!IsWindowVisible(hwnd)) {
			windows_del(desk, hwnd);
		}
	}
	EnumWindows((WNDENUMPROC)&enum_func, (LPARAM)v);
}

static void virgo_toggle_hotkeys(Virgo *v)
{
	unsigned i;
	v->handle_hotkeys = !v->handle_hotkeys;
	if (v->handle_hotkeys) {
		for (i = 0; i < NUM_DESKTOPS; i++) {
			register_hotkey(i * 2, MOD_ALT | MOD_CONTROL | MOD_NOREPEAT, i + 1 + '0');
			register_hotkey(i * 2 + 1, MOD_ALT | MOD_SHIFT | MOD_NOREPEAT, i + 1 + '0');

		}
	} else {
		for (i = 0; i < NUM_DESKTOPS; i++) {
			UnregisterHotKey(NULL, i * 2);
			UnregisterHotKey(NULL, i * 2 + 1);
		}
	}
}

static void virgo_init(Virgo *v)
{
	unsigned i;
	v->handle_hotkeys = 1;
	for (i = 0; i < NUM_DESKTOPS; i++) {
		register_hotkey(i * 2, MOD_ALT | MOD_CONTROL | MOD_NOREPEAT, i + 1 + '0');
		register_hotkey(i * 2 + 1, MOD_ALT | MOD_SHIFT | MOD_NOREPEAT, i + 1 + '0');
	}
	register_hotkey(i * 2, MOD_ALT | MOD_CONTROL | MOD_SHIFT | MOD_NOREPEAT,
					'Q');
	register_hotkey(i * 2 + 1, MOD_ALT | MOD_CONTROL | MOD_SHIFT | MOD_NOREPEAT,
					'S');
	register_hotkey(i * 2 + 2, MOD_ALT | MOD_CONTROL | MOD_SHIFT | MOD_NOREPEAT,
					'P');
}

static void virgo_deinit(Virgo *v)
{
	unsigned i;
	for (i = 0; i < NUM_DESKTOPS; i++) {
		windows_show(&v->desktops[i]);
		sb_free(v->desktops[i].windows);
	}
}

static void virgo_move_to_desk(Virgo *v, unsigned desk)
{
	HWND hwnd;
	if (v->current == desk) {
		return;
	}
	virgo_update(v);
	hwnd = GetForegroundWindow();
	if (virgo_is_pinned(v, hwnd)) {
		return;
	}
	if (!hwnd || !is_valid_window(hwnd)) {
		return;
	}
	windows_del(&v->desktops[v->current], hwnd);
	windows_add(&v->desktops[desk], hwnd);
	ShowWindow(hwnd, SW_HIDE);
}

static void virgo_go_to_desk(Virgo *v, unsigned desk)
{
	if (v->current == desk) {
		return;
	}
	virgo_update(v);
	v->focused[v->current] = GetForegroundWindow();
	windows_hide(&v->desktops[v->current]);
	windows_show(&v->desktops[desk]);
	SetForegroundWindow(v->focused[desk]);
	v->current = desk;
}

void __main(void) __asm__("__main");
void __main(void)
{
	Virgo v = {0};
	MSG msg;
	virgo_init(&v);
	while (GetMessage(&msg, NULL, 0, 0)) {
		if (msg.message != WM_HOTKEY) {
			continue;
		}
		if (msg.wParam == NUM_DESKTOPS * 2) {
			break;
		}
		if (msg.wParam == NUM_DESKTOPS * 2 + 1) {
			virgo_toggle_hotkeys(&v);
		} else if (msg.wParam == NUM_DESKTOPS * 2 + 2) {
			virgo_toggle_pin(&v, GetForegroundWindow());
		} else if (msg.wParam % 2 == 0) {
			virgo_go_to_desk(&v, msg.wParam / 2);
		} else {
			virgo_move_to_desk(&v, (msg.wParam - 1) / 2);
		}
	}
	virgo_deinit(&v);
	ExitProcess(0);
}