Warning: Due to various recent migrations, viewing non-HEAD refs may be broken.

/ds/map/btree/iter.ha (raw)

// SPDX-License-Identifier: MPL-2.0

use ds::map;

export type iterator = struct {
	vt: map::iterator,
	nodes: []*node,
	idxs: []size,
	visit_key: []bool,
	finished: bool,
};

const _itvt: map::vtable_iterator = map::vtable_iterator {
	nexter = &vt_next,
};

fn vt_next(it: *map::iterator) (([]u8, *opaque) | done) = next(it: *iterator);

export fn iter(m: *map) (*map::iterator | nomem) = {
	let it = alloc(iterator {
		vt = &_itvt,
		nodes = [],
		idxs = [],
		visit_key = [],
		finished = false,
	})?;

	match (append(it.nodes, m.root)) {
	case void => void;
	case nomem => return nomem;
	};
	match (append(it.idxs, 0z)) {
	case void => void;
	case nomem => return nomem;
	};
	match (append(it.visit_key, false)) {
	case void => void;
	case nomem => return nomem;
	};

	return (it: *map::iterator);
};

export fn next(it: *iterator) (([]u8, *opaque) | done) = {
	if (it.finished) {
		return done;
	};
	for (len(it.nodes) != 0) {
		let top = len(it.nodes) - 1;
		let x = it.nodes[top];
		let i = it.idxs[top];
		let vk = it.visit_key[top];

		if (x.leaf) {
			if (i < len(x.keys)) {
				it.idxs[top] = i + 1;
				return (x.keys[i], x.vals[i]);
			};
			delete(it.nodes[top]);
			delete(it.idxs[top]);
			delete(it.visit_key[top]);
			continue;
		};

		if (vk) {
			if (i >= len(x.keys)) {
				it.visit_key[top] = false;
				continue;
			};
			it.visit_key[top] = false;
			it.idxs[top] = i + 1;
			return (x.keys[i], x.vals[i]);
		};

		if (i < len(x.keys)) {
			match (append(it.nodes, x.children[i])) {
			case void => void;
			case nomem => abort("btree::iter: nomem");
			};
			match (append(it.idxs, 0z)) {
			case void => void;
			case nomem => abort("btree::iter: nomem");
			};
			match (append(it.visit_key, false)) {
			case void => void;
			case nomem => abort("btree::iter: nomem");
			};
			it.visit_key[top] = true;
			continue;
		};

		if (i == len(x.keys)) {
			match (append(it.nodes, x.children[i])) {
			case void => void;
			case nomem => abort("btree::iter: nomem");
			};
			match (append(it.idxs, 0z)) {
			case void => void;
			case nomem => abort("btree::iter: nomem");
			};
			match (append(it.visit_key, false)) {
			case void => void;
			case nomem => abort("btree::iter: nomem");
			};
			it.idxs[top] = i + 1;
			continue;
		};

		delete(it.nodes[top]);
		delete(it.idxs[top]);
		delete(it.visit_key[top]);
	};

	free(it.nodes);
	free(it.idxs);
	free(it.visit_key);
	it.nodes = [];
	it.idxs = [];
	it.visit_key = [];
	it.finished = true;
	return done;
};