mandelbrot.c

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <png.h>
#include <pthread.h>
#include <unistd.h>
#include <semaphore.h>

// Modify these
static int xstart = -2;
static int xend = 1;
static int ystart = -1;
static int yend = 1;
static double resolution = 10000;

static int nthreads_per_cpu = 4;

static double min(double n1, double n2)
{
	return n1 > n2 ? n2 : n1;
}

typedef struct compnum
{
	double real;
	double imgn;
} compnum;

static void compnum_init(compnum* c, double real, double imgn)
{
	c->real = real;
	c->imgn = imgn;
}

static void compnum_mul(compnum *dst, compnum *c, compnum *num)
{
	double dr = (c->real * num->real) - (c->imgn * num->imgn);
	double di = (c->real * num->imgn) + (c->imgn * num->real);
	dst->real = dr;
	dst->imgn = di;
}

static void compnum_add(compnum *dst, compnum *c, compnum *num)
{
	dst->real = c->real + num->real;
	dst->imgn = c->imgn + num->imgn;
}

static double compnum_abs(compnum *c)
{
	return sqrt((c->real * c->real) + (c->imgn * c->imgn));
}

static void compnum_f(compnum *dst, compnum *c, compnum *num)
{
	compnum_mul(dst, num, num);
	compnum_add(dst, c, dst);
}

// Returns 0 if mandelbrot, or 1 or more if it's not mandelbrot.
// The number is the number of times you have to run compnum_f to get a number
// that's guaranteed to be trending towards infinity.
static int compnum_mandelbrot(compnum *c)
{
	static int lim = 100;

	compnum num;
	compnum zero;
	compnum_init(&zero, 0, 0);

	compnum_f(&num, c, &zero);

	for (int i = 0; i < lim; ++i)
	{
		compnum_f(&num, c, &num);

		if (compnum_abs(&num) > 2)
		{
			return i + 1;
		}
	}

	return 0;
}

// Never used libpng before, based on code from
// http://zarb.org/~gc/html/libpng.html
static char* write_png(char* path, png_bytep *row_pointers, int width, int height)
{
	png_byte bit_depth = 8;
	png_byte color_type = PNG_COLOR_TYPE_RGB;

	FILE *f = fopen(path, "wb");
	if (!f)
		return "Failed to open file for writing";

	// Create write struct

	png_structp png_ptr = png_create_write_struct(
		PNG_LIBPNG_VER_STRING,
		NULL, NULL, NULL);

	if (!png_ptr)
		return "png_create_write_struct failed";

	// Create info struct

	png_infop info_ptr = png_create_info_struct(png_ptr);

	if (!info_ptr)
		return "png_create_info_struct failed";

	// Init IO

	if (setjmp(png_jmpbuf(png_ptr)))
		return "Error during init_io";

	png_init_io(png_ptr, f);

	// Write header

	if (setjmp(png_jmpbuf(png_ptr)))
		return "Error during writing header";

	png_set_IHDR(png_ptr, info_ptr, width, height,
		bit_depth, color_type, PNG_INTERLACE_NONE,
		PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

	png_write_info(png_ptr, info_ptr);

	// Write bytes

	if (setjmp(png_jmpbuf(png_ptr)))
		return "Error during writing bytes";

	png_write_image(png_ptr, row_pointers);

	// End write

	if (setjmp(png_jmpbuf(png_ptr)))
		return "Error during end of write";

	png_write_end(png_ptr, NULL);

	fclose(f);

	return NULL;
}

struct draw_args {
	int xstart;
	int xend;
	int ystart;
	int yend;
	int yoffset;
	png_bytep *row_pointers;
};

void *draw_lines(void* args)
{
	struct draw_args dargs = *(struct draw_args *)args;
	int xstart = dargs.xstart;
	int xend = dargs.xend;
	int ystart = dargs.ystart;
	int yend = dargs.yend;
	int yoffset = dargs.yoffset;
	png_bytep *row_pointers = dargs.row_pointers;

	for (int y = ystart; y <= yend; ++y)
	{
		int ly = y - ystart + yoffset;
		png_byte* row = row_pointers[ly];

		for (int x = xstart; x <= xend; ++x)
		{
			compnum c;
			compnum_init(&c, x / resolution, y / resolution);
			int times = compnum_mandelbrot(&c);

			int lx = (x - xstart) * 3;
			if (times == 0)
			{
				row[lx] = 0;
				row[lx + 1] = 0;
				row[lx + 2] = 0;
			}
			else
			{
				png_byte color = min(((times / 100.0) * 230) + 90, 255);
				row[lx] = color;
				row[lx + 1] = color / 2;
				row[lx + 2] = color / 3;
			}
		}
	}

	printf("row %i to %i done\n", ystart, yend);

	return NULL;
}

void draw(
		int xstart, int xend,
		int ystart, int yend,
		int resolution, int nthreads)
{
	xstart *= resolution;
	xend *= resolution;
	ystart *= resolution;
	yend *= resolution;

	int width = xend - xstart + 1;
	int height = yend - ystart + 1;

	// Prepare the array of rows
	png_bytep *row_pointers = (png_bytep *)malloc(
		sizeof(*row_pointers) * height);

	// Prepare each row
	for (int i = 0; i < height; ++i)
	{
		row_pointers[i] = (png_byte *)malloc(
			sizeof(png_bytep *) * width * 3);
	}

	// Divide up work
	printf("creating %i threads\n", nthreads);

	struct draw_args *allargs[nthreads];
	pthread_t threads[nthreads];
	for (int i = 0; i < nthreads; ++i)
	{
		int start = ystart + ((height / nthreads) * i);
		int yoffset = start - ystart;

		int end;
		if (i == nthreads - 1)
			end = yend;
		else
			end = ystart + ((height / nthreads) * (i + 1)) - 1;

		struct draw_args *args = malloc(sizeof(*args));
		args->xstart = xstart;
		args->xend = xend;
		args->ystart = start;
		args->yend = end;
		args->yoffset = yoffset;
		args->row_pointers = row_pointers;
		allargs[i] = args;

		pthread_t pth;
		pthread_create(&pth, NULL, &draw_lines, (void*)args);
		threads[i] = pth;
	}

	// Wait for threads to complete
	for (int i = 0; i < nthreads; ++i)
		pthread_join(threads[i], NULL);

	// Free arguments
	for (int i = 0; i < nthreads; ++i)
		free(allargs[i]);

	printf("Done, writing png file...\n");

	char* err = write_png("fractal.png", row_pointers, width, height);
	if (err)
		fprintf(stderr, err);

	// Clean memory
	for (int i = 0; i < height; ++i)
	{
		free(row_pointers[i]);
	}
	free(row_pointers);
}

int main()
{
	int nthreads = sysconf(_SC_NPROCESSORS_ONLN) * nthreads_per_cpu;
	draw(xstart, xend, ystart, yend, resolution, nthreads);
}