fast-memchr.c

#include <stdint.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>

#define MIN(a,b) (((a)<(b))?(a):(b))
#define LO ((uint64_t)0x0101010101010101L)
#define HI ((uint64_t)0x8080808080808080L)
#define LOOP_SIZE (2 * sizeof(uint64_t))

static bool
contains_zero_byte(uint64_t x) {
	return(((x - LO) & ~x & HI) != 0);
}

static uint64_t
repeat_byte(uint8_t b) {
	return(((uint64_t) b) * (UINT64_MAX / 255));
}

static uint64_t
read_unaligned64(uint8_t *ptr) {
	uint64_t ret;
	memcpy(&ret, ptr, sizeof(uint64_t));
	return ret;
}

static void *
forward_search(uint8_t *start_ptr, uint8_t *end_ptr, uint8_t *ptr, uint8_t n) {
	while (ptr <= end_ptr) {
		if (*ptr == n) {
			return (void *)ptr;
		}
		ptr++;
	}

	return NULL;
}

void *
fast_memchr(void *haystack, int n, size_t len) {
	uint8_t n1 = (uint8_t) n;
	uint64_t vn1 = repeat_byte(n1);
	uint32_t loop_size = MIN(LOOP_SIZE, len);
	uint64_t align = sizeof(uint64_t) - 1;

	uint8_t *start_ptr = haystack;
	uint8_t *end_ptr = haystack + len;
	uint8_t *ptr = start_ptr;

	if (len < sizeof(uint64_t)) {
		return forward_search(start_ptr, end_ptr, ptr, n1);
	}

	uint64_t chunk = read_unaligned64(ptr);
	if (contains_zero_byte(chunk ^ vn1)) {
		return forward_search(start_ptr, end_ptr, ptr, n1);
	}

	ptr += sizeof(uint64_t) - (((uint64_t)start_ptr) & align);

	while (loop_size == LOOP_SIZE && ptr <= (end_ptr - loop_size)) {
		uint64_t a = *((uint64_t *)ptr);
		uint64_t b = *((uint64_t *)(ptr + sizeof(uint64_t)));
		bool eqa = contains_zero_byte(a ^ vn1);
		bool eqb = contains_zero_byte(b ^ vn1);

		if (eqa || eqb) {
			break;
		}
		ptr += LOOP_SIZE;
	}

	return forward_search(start_ptr, end_ptr, ptr, n1);
}

#include <unistd.h> 
#include <fcntl.h>
#define SIZE (1024 * 64)
#define LOOPS (1024 * 64)

int main(void) {
	int fd = open("test.file", O_RDONLY);

	char *buf[SIZE];
	size_t len = read(fd, &buf, SIZE);

	if (len != SIZE) {
		return 64;
	}

	int loop = LOOPS;
	int count = 0;

	while (loop--) {
		for (int i=0; i < 255; i++) {
			if (fast_memchr(&buf, i, len) != NULL) {
				count++;
			}
		}
	}

	printf("found %d\n", count);
	return 0;
}