ported rewrite2 to c3

src/vtree.c3

@@ -0,0 +1,333 @@
+module vtree(<ElemType>);
+
+import std::core::mem;
+import std::io;
+
+struct VTree {
+	usz elements;
+	ElemType[] vector; // vector of element ids
+	isz[] refs, ordered_refs;
+}
+
+fault VTreeError {
+    CANNOT_SHRINK,
+    INVALID_REFERENCE,
+    TREE_FULL,
+    REFERENCE_NOT_PRESENT,
+    INVALID_ARGUMENT,
+}
+
+macro VTree.ref_is_valid(&tree, isz ref) { return (ref >= 0 && ref < tree.refs.len); }
+macro VTree.ref_is_present(&tree, isz ref) { return tree.refs[ref] >= 0; }
+macro VTree.size(&tree) { return tree.refs.len; }
+
+// macro to zero an elemen
+macro @zero()
+{
+	$if $assignable(0, ElemType):
+		return 0;
+	$else
+		return ElemType{0};
+	$endif
+}
+
+fn void! VTree.init(&tree, usz size)
+{
+	tree.vector = mem::new_array(ElemType, size);
+	defer catch { (void)mem::free(tree.vector); }
+
+	tree.refs = mem::new_array(isz, size);
+	defer catch { (void)mem::free(tree.refs); }
+
+	tree.ordered_refs = mem::new_array(isz, size);
+	defer catch { (void)mem::free(tree.ordered_refs); }
+
+	// set all refs to -1, meaning invalid (free) element
+	tree.refs[..] = -1;
+
+	tree.elements = 0;
+}
+
+fn void VTree.free(&tree)
+{
+	(void)mem::free(tree.vector);
+	(void)mem::free(tree.refs);
+	(void)mem::free(tree.ordered_refs);
+}
+
+fn void VTree.pack(&tree)
+{
+	// TODO: add a PACKED flag to skip this
+
+	isz free_spot = -1;
+	for (usz i = 0; i < tree.size(); i++) {
+		if (tree.refs[i] == -1) {
+			free_spot = i;
+			continue;
+		}
+
+		// find a item that can be packed
+		if (free_spot >= 0 && tree.refs[i] >= 0) {
+			isz old_ref = i;
+
+			// move the item
+			tree.vector[free_spot] = tree.vector[i];
+			tree.refs[free_spot]   = tree.refs[i];
+
+			tree.vector[i] = @zero();
+			tree.refs[i]   = -1;
+
+			// and move all references
+			for (usz j = 0; j < tree.size(); j++) {
+				if (tree.refs[j] == old_ref) {
+				    tree.refs[j] = free_spot;
+				}
+			}
+
+			// mark the free spot as used
+			free_spot = -1;
+		}
+	}
+}
+
+fn void! VTree.resize(&tree, usz newsize)
+{
+	// return error when shrinking with too many elements
+	if (newsize < tree.elements) {
+		return VTreeError.CANNOT_SHRINK?;
+	}
+
+	// pack the vector when shrinking to avoid data loss
+	if ((int)newsize < tree.size()) {
+		// FIXME: packing destroys all references to elements of vec
+		//        so shrinking may cause dangling pointers
+		return VTreeError.CANNOT_SHRINK?;
+	}
+
+	usz old_size = tree.size();
+
+	tree.vector = ((ElemType*)mem::realloc(tree.vector, newsize*ElemType.sizeof))[:newsize];
+	defer catch { (void)mem::free(tree.vector); }
+
+	tree.refs = ((isz*)mem::realloc(tree.refs, newsize*isz.sizeof))[:newsize];
+	defer catch { (void)mem::free(tree.refs); }
+
+	tree.ordered_refs = ((isz*)mem::realloc(tree.ordered_refs, newsize*isz.sizeof))[:newsize];
+	defer catch { (void)mem::free(tree.ordered_refs); }
+
+	if (newsize > tree.size()) {
+	   tree.vector[old_size..newsize-1] = @zero();
+	   tree.refs[old_size..newsize-1] = -1;
+	}
+}
+
+// add an element to the tree, return it's ref
+fn isz! VTree.add(&tree, ElemType elem, isz parent)
+{
+	// invalid parent
+	if (!tree.ref_is_valid(parent)) {
+		return VTreeError.INVALID_REFERENCE?;
+	}
+
+	// no space left
+	if (tree.elements >= tree.size()) {
+	   return VTreeError.TREE_FULL?;
+	}
+
+	// check if the parent exists
+	// if there are no elements in the tree the first add will set the root
+	if (!tree.ref_is_present(parent) && tree.elements != 0) {
+	   return VTreeError.REFERENCE_NOT_PRESENT?;
+	}
+
+	// get the first free spot
+	isz free_spot = -1;
+	for (usz i = 0; i < tree.size(); i++) {
+		if (tree.refs[i] == -1) {
+			free_spot = i;
+			break;
+		}
+	}
+	if (free_spot < 0) {
+		return VTreeError.TREE_FULL?;
+	}
+
+	// finally add the element
+	tree.vector[free_spot] = elem;
+	tree.refs[free_spot]   = parent;
+	tree.elements++;
+
+	return free_spot;
+}
+
+// prune the tree starting from the ref
+// returns the number of pruned elements
+fn usz! VTree.prune(&tree, isz ref)
+{
+    if (!tree.ref_is_valid(ref)) {
+		return VTreeError.INVALID_REFERENCE?;
+	}
+
+	if (!tree.ref_is_present(ref)) {
+		return 0;
+	}
+
+	tree.vector[ref] = @zero();
+	tree.refs[ref]   = -1;
+	tree.elements--;
+
+	usz count = 1;
+	for (usz i = 0; tree.elements > 0 && i < tree.size(); i++) {
+		if (tree.refs[i] == ref) {
+			count += tree.prune(i)!;
+		}
+	}
+
+	return count;
+}
+
+// find the size of the subtree starting from ref
+fn usz! VTree.subtree_size(&tree, isz ref)
+{
+    if (!tree.ref_is_valid(ref)) {
+		return VTreeError.INVALID_REFERENCE?;
+	}
+
+	if (!tree.ref_is_present(ref)) {
+		return 0;
+	}
+
+	usz count = 1;
+	for (usz i = 0; i < tree.size(); i++) {
+		// only root has the reference to itself
+		if (tree.refs[i] == ref && ref != i) {
+			count += tree.subtree_size(i)!;
+		}
+	}
+
+	return count;
+}
+
+// iterate through the first level children, use a cursor like strtok_r
+fn isz! VTree.children_it(&tree, isz parent, isz *cursor)
+{
+	if (cursor == null) {
+		return VTreeError.INVALID_ARGUMENT?;
+	}
+
+	// if the cursor is out of bounds then we are done for sure
+	if (!tree.ref_is_valid(*cursor)) {
+		return VTreeError.INVALID_REFERENCE?;
+	}
+
+	// same for the parent, if it's invalid it can't have children
+	if (!tree.ref_is_valid(parent) || !tree.ref_is_present(parent)) {
+		return VTreeError.INVALID_REFERENCE?;
+	}
+
+	// find the first child, update the cursor and return the ref
+	for (isz i = *cursor; i < tree.size(); i++) {
+		if (tree.refs[i] == parent) {
+			*cursor = i + 1;
+			return i;
+		}
+	}
+
+	// if no children are found return -1
+	*cursor = -1;
+	return -1;
+}
+
+/* iterates trough every leaf of the subtree in the following manner
+ * node [x], x: visit order
+ *              [0]
+ *             / | \
+ *            / [2] [3]
+ *          [1]      |
+ *         /   \    [6]
+ *       [4]   [5]
+ */
+fn isz! VTree.level_order_it(&tree, isz ref, isz *cursor)
+{
+	if (cursor == null) {
+	   return VTreeError.INVALID_ARGUMENT?;
+	}
+
+	isz[] queue = tree.ordered_refs;
+
+	// TODO: this could also be done when adding or removing elements
+	// first call, create a ref array ordered like we desire
+	if (*cursor == -1) {
+		*cursor = 0;
+		queue[..] = -1;
+
+		// iterate through the queue appending found children
+		isz pos, off;
+		do {
+			// printf ("ref=%d\n", ref);
+			for (isz i = 0; i < tree.size(); i++) {
+				if (tree.refs[i] == ref) {
+					queue[pos++] = i;
+				}
+			}
+
+			for (; ref == queue[off] && off < tree.size(); off++);
+			ref = queue[off];
+
+		} while (tree.ref_is_valid(ref));
+		// This line is why tree.ordered_refs has to be size+1
+		queue[off + 1] = -1;
+	}
+
+	// PRINT_ARR(queue, tree.size());
+	// return -1;
+
+	// on successive calls just iterate through the queue until we find an
+	// invalid ref, if the user set the cursor to -1 it means it has found what
+	// he needed, so free
+	if (*cursor < 0) {
+		return -1;
+	} else if (tree.ref_is_valid(*cursor)) {
+		return queue[(*cursor)++];
+	}
+
+	return -1;
+}
+
+fn isz! VTree.parentof(&tree, isz ref)
+{
+	if (!tree.ref_is_valid(ref)) {
+		return VTreeError.INVALID_REFERENCE?;
+	}
+
+	if (!tree.ref_is_present(ref)) {
+	   return VTreeError.REFERENCE_NOT_PRESENT?;
+	}
+
+	return tree.refs[ref];
+}
+
+fn ElemType! VTree.get(&tree, isz ref)
+{
+    if (!tree.ref_is_valid(ref)) {
+		return VTreeError.INVALID_REFERENCE?;
+	}
+
+	if (!tree.ref_is_present(ref)) {
+	   return VTreeError.REFERENCE_NOT_PRESENT?;
+	}
+
+	return tree.vector[ref];
+}
+
+fn void VTree.print(&tree)
+{
+	for (isz i = 0; i < tree.size(); i++) {
+		if (tree.refs[i] == -1) {
+			continue;
+		}
+		io::printf("[%d] {parent=%d, data=", i, tree.refs[i]);
+		io::print(tree.vector[i]);
+		io::printn("}");
+	}
+}