aoc2021/day3/cmkinput.c

//! -pipe -std=gnu17 -O3 -msse -fwhole-program -flto -fno-strict-aliasing -Wall -Wstrict-aliasing -Werror -Wl,-flto -Wl,-O3 -o

#include <stdio.h>
#include <assert.h>
#include <stdlib.h>
#include <stdint.h>
#include <fcntl.h>
#include <stdbool.h>
#include <unistd.h>
#include <errno.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>

#define expect(expr, msg) ({ __auto_type _expect__expr = (expr); \
			if(__builtin_expect(!_expect__expr, 0)) { perror("Fatal error: `" #expr "` failed: " msg); exit(1); } \
			_expect__expr; })

#define report(...) ({ fprintf(stderr, __VA_ARGS__); fputc('\n', stderr); false; })
#define throw(...) return report(__VA_ARGS__)

typedef struct mappedfd {
	int fd;
	size_t len;

	void* origin;
} mmap_t;

bool filesz(int fd, off_t* restrict size)
{
	if(fd<0) return false;
	struct stat st;
	if(fstat(fd, &st)) return false;
	*size = st.st_size;
	return true;
}

bool map_stdin(mmap_t* restrict map, size_t sz)
{
	size_t size;
	if(! (size = sz)) {
		off_t osz;
		if(!filesz(STDIN_FILENO, &osz)) return false;
		size = (size_t)osz;
	}
	void* origin = mmap(NULL, size, PROT_READ, MAP_PRIVATE, STDIN_FILENO, 0);
	if(origin == MAP_FAILED) return false;
	*map = (mmap_t) {
		.fd = STDIN_FILENO,
		.len = size,
		.origin = origin,
	};	
	return true;
}

bool map_input(const char* file, mmap_t *restrict map)
{
	int fd = open(file, O_RDONLY);
	if(fd<0) return false;

	off_t _sz;
	if(!filesz(fd, &_sz)) return (close(fd), false);
	size_t sz = (size_t)_sz;

	void* origin = mmap(NULL, sz, PROT_READ, MAP_PRIVATE, fd, 0);
	if(origin == MAP_FAILED) return (close(fd), false);
	
	*map = (mmap_t){
		.fd = fd,
		.len = sz,
		.origin = origin,
	};

	return true;
}

bool map_stdout(mmap_t* restrict map, size_t size)
{
	char link[64];
	int fd;
	snprintf(link, sizeof(link), "/proc/%d/fd/1", getpid());

	fd = open(link, O_RDWR);
	if(fd < 0) return report("failed to open fd 1 for rdwr (path was `%s')", link);

	//dup2(fd, STDOUT_FILENO);
	//close(fd);
	// !!! normal stdout dead beyond this point
	if(ftruncate(fd, size)<0) return (close(fd), report("failed to truncate stdout (temp %d) to size %lu", fd, size));
	if(dup2(fd, STDOUT_FILENO)<0) return (close(fd), report("failed to dup2 %d -> %d", fd, STDOUT_FILENO));
	close(fd);

	void* origin = mmap(NULL, size, PROT_WRITE | PROT_READ, MAP_SHARED, STDOUT_FILENO, 0);
	if(origin == MAP_FAILED) return report("failed to map stdout as write + shared at size %lu", size);

	*map = (mmap_t){
		.fd = STDOUT_FILENO,
		.len = size,
		.origin = origin,
	};
	return true;
}

void map_close(mmap_t map)
{
	munmap(map.origin, map.len);
	if(map.fd >= 0) close(map.fd);
}

const unsigned char VALID[255] = {
	['0' ... '9'] = 1,
	['\n'] = 2,
};

inline bool is_number(const char* str)
{
	register unsigned char vl;
	while( (vl = (unsigned char)*str++) ) if(VALID[vl] != 1) return false;
	return true;
}

size_t get_width(const char* str)
{

	size_t i=0;
	unsigned char v;
	while( (v = VALID[(unsigned char) *str++]) ) {
		i+=1;
		if(v == 2) break;
	}
	return i;
}

bool process(const mmap_t* imap, const mmap_t* omap, size_t *restrict osize)
{
	const char* input = imap->origin;
	char* output = omap->origin;
	uintptr_t _start = (uintptr_t)output;
	size_t size = omap->len;

	size_t width = get_width(input) - 1; //XXX: This isn't working, because of newlines and bullshit like that. Idk how to fix it...
	size_t height = size / (width + 1);
	fprintf(stderr, "size: %lu, w: %lu, h: %lu\n", size, width, height);
#define getat(ar, x, y) ((ar)[(((y) * (width+1)) + (x))])
	assert(size % (width + 1)== 0);

	for(size_t x = 0;x<width;x++) {
		for(size_t y = 0; y<height; y++)
			*output++ = getat(input, x, y);
		*output++ = '\n';
	}
	*osize = (size_t)(((uintptr_t)output) - _start);
	return true;
#undef getat
}

mmap_t open_input(char*const* argv)
{
	const char* in = argv[1];
	mmap_t input;
	if(!in) expect(map_stdin(&input, 0), "failed to map stdin with no size hint");
	else if(is_number(in)) expect(map_stdin(&input, strtoul(in, NULL, 10)), "failed to map stdin at that size");
	else expect(map_input(in, &input), "failed to map input file");
	return input;
}

int main(int argc, char** argv)
{
	int rc = 0;

	mmap_t input;
	mmap_t output;

	input = open_input(argv); //expect(map_stdin(&input), "failed to map stdin");
	expect(map_stdout(&output, input.len), "failed to map stdout");

	expect(input.len == output.len, "invalid i/o lengths");
	size_t fout;
	expect(process(&input, &output, &fout), "failed to process");
	fprintf(stderr, "output size %lu -> %lu\n", output.len, fout);
	map_close(input);

	if(msync(output.origin, fout, MS_SYNC)) perror("failed to sync. output may be corrupted");
	munmap(output.origin, output.len);
	if(ftruncate(output.fd, fout)) perror("failed to truncate file to correct output size");
	close(output.fd);

	return rc;
}