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

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

// SPDX-License-Identifier: MPL-2.0

use ds::map;

export type iterator = struct {
	vt: map::iterator,
	cur: nullable *node,
};

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 start: nullable *node = match (m.root) {
	case null => yield null;
	case let r: *node => yield subtree_min(r);
	};
	let it = alloc(iterator {
		vt = &_itvt,
		cur = start,
	})?;
	return (it: *map::iterator);
};

fn successor(n0: *node) nullable *node = {
	let n = n0;

	match (n.right) {
	case let r: *node =>
		return subtree_min(r);
	case null =>
		void;
	};

	let mut_p: nullable *node = n.parent;
	let child: nullable *node = n;

	for (mut_p != null) {
		let p = match (mut_p) {
		case let pp: *node => yield pp;
		case null => break;
		};
		if (p.right != child) {
			return p;
		};
		child = p;
		mut_p = p.parent;
	};

	return null;
};

export fn next(it: *iterator) (([]u8, *opaque) | done) = {
	match (it.cur) {
	case null =>
		return done;
	case let n: *node =>
		let k = n.key;
		let v = n.val;
		it.cur = successor(n);
		return (k, v);
	};
};