shithub: snippets

ref: 6669ea8ed594d206d16171b7034125deeca6b806
dir: /treewalk.c/

/*
 * Demonstrates how simple recursive traversing a non-balanced tree can
 * eat your stack (fast).  Demonstrates an iterative O(N) complexity +
 * O(1) space traversal to be used instead if your tree may look like a
 * very long stick.
 *
 * Read the code.
 *
 * 7c -FTVw treewalk.c && 7l treewalk.7 && ./7.out [options]
 */
#include <u.h>
#include <libc.h>
#include <thread.h>
#include <bio.h>

typedef struct Tree Tree;

struct Tree {
	int v;
	Tree *left;
	Tree *right;
};

enum {
	PreOrder,
	InOrder,
	PostOrder,
};

static Biobuf *b;
static int simple;
static int mode = PreOrder;

/*
 * A tree that is on https://en.wikipedia.org/wiki/Tree_traversal#Depth-first_search
 * (as of 2024/10/24). It's to verify the correctness with.
 */
static Tree stC = {'C'}, stE = {'E'}, stH = {'H'};
static Tree stA = {'A'}, stD = {'D', &stC, &stE}, stI = {'I', &stH};
static Tree stB = {'B', &stA, &stD}, stG = {'G', nil, &stI};
static Tree st = {'F', &stB, &stG};

static void
visit(Tree *t)
{
	Bprint(b, simple ? "%c " : "%d\n", t->v);
}

/* nothing special here */
static void
walkrecursive(Tree *t)
{
	if(mode == PreOrder)
		visit(t);
	if(t->left)
		walkrecursive(t->left);
	if(mode == InOrder)
		visit(t);
	if(t->right)
		walkrecursive(t->right);
	if(mode == PostOrder)
		visit(t);
}

/*
* To morris-traverse in post-order we have to visit nodes in reverse
* upon discovering an earlier created thread.  This function flips the
* subtree in-place, then visits each node while flipping the subtree
* back to what it was.  It still does not allocate any extra space.
*/
static void
postwalk(Tree *a, Tree *b)
{
	Tree *v, *vr, *vr2;

	if(a != b){
		for(v = a, vr = v->right; vr != b; v = vr, vr = vr2){
			vr2 = vr->right;
			vr->right = v;
		}
		vr->right = v;
		for(v = b, vr = v->right; vr != a; v = vr, vr = vr2){
			visit(v);
			vr2 = vr->right;
			vr->right = v;
		}
		visit(v);
		vr->right = v;
	}
	visit(a);
}

/* this one is more interesting */
static void
walkmorris(Tree *t)
{
	Tree *p, t0;

	if(mode == PostOrder){
		/* in post-order we need to have a fake root to return to */
		t0.v = 'X';
		t0.left = t;
		t0.right = nil;
		t = &t0;
	}

	for(; t ;){
		if(t->left){
			for(p = t->left; p->right && p->right != t; p = p->right);
			if(p->right){
				if(mode == InOrder)
					visit(t);
				else if(mode == PostOrder)
					postwalk(t->left, p);
				p->right = nil;
				t = t->right;
			}else{
				if(mode == PreOrder)
					visit(t);
				p->right = t;
				t = t->left;
			}
		}else{
			if(mode != PostOrder)
				visit(t);
			t = t->right;
		}
	}
}

void
threadmain(int argc, char **argv)
{
	Tree *t;
	int n, i, num = 1<<16;
	int nol, nor, morris;

	nol = nor = simple = morris = 0;
	ARGBEGIN{
	case 'l': /* no left nodes will be set */
		nol++;
		break;
	case 'r': /* no right nodes will be set */
		nor++;
		break;
	case 'm': /* use iterative traversal */
		morris++;
		break;
	case 'i':
		mode = InOrder;
		break;
	case 'p':
		mode = PostOrder;
		break;
	case 's': /* use a small and simple tree */
		simple++;
		break;
	default:
		fprint(2, "usage: %s [-i | -p] [-l | -r] [-m] [-s]\n", argv0);
		threadexitsall("usage");
	}ARGEND

	if(nol && nor)
		sysfatal("it wouldn't even be a list, dumbass");

	b = Bfdopen(1, OWRITE);
	if(simple){
		t = &st;
	}else{
		t = calloc(num, sizeof(*t));
		for(i = 0, n = 1; i < num; i++){
			if(!nol && n < num)
				t[i].left = &t[n++];
			if(!nor && n < num)
				t[i].right = &t[n++];
			t[i].v = i;
		}
	}

	if(morris)
		walkmorris(t);
	else
		walkrecursive(t);
	if(simple)
		Bprint(b, "\n");

	Bterm(b);
	threadexitsall(nil);
}