Loading...
src/radix_tree.c /dev/null libmalloc-166.220.1
--- /dev/null
+++ libmalloc/libmalloc-166.220.1/src/radix_tree.c
@@ -0,0 +1,602 @@
+/*
+ * Copyright (c) 2016 Apple Inc. All rights reserved.
+ *
+ * @APPLE_LICENSE_HEADER_START@
+ *
+ * This file contains Original Code and/or Modifications of Original Code
+ * as defined in and that are subject to the Apple Public Source License
+ * Version 2.0 (the 'License'). You may not use this file except in
+ * compliance with the License. Please obtain a copy of the License at
+ * http://www.opensource.apple.com/apsl/ and read it before using this
+ * file.
+ *
+ * The Original Code and all software distributed under the License are
+ * distributed on an 'AS IS' basis, WITHOUT WARRANTY OF ANY KIND, EITHER
+ * EXPRESS OR IMPLIED, AND APPLE HEREBY DISCLAIMS ALL SUCH WARRANTIES,
+ * INCLUDING WITHOUT LIMITATION, ANY WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE, QUIET ENJOYMENT OR NON-INFRINGEMENT.
+ * Please see the License for the specific language governing rights and
+ * limitations under the License.
+ *
+ * @APPLE_LICENSE_HEADER_END@
+ */
+
+#include <assert.h>
+#include <mach/mach.h>
+#include <mach/mach_vm.h>
+#include <stdbool.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/reason.h>
+#include <unistd.h>
+
+#include <radix_tree.h>
+#include <radix_tree_internal.h>
+
+#if 0
+int radix_tree_indent = 0;
+bool radix_tree_should_print __attribute__((visibility("default")))= true;
+#include <stdio.h>
+#define D(s, ...)                                   \
+	if (radix_tree_should_print) {                  \
+		for (int i = 0; i < radix_tree_indent; i++) \
+			putchar(' ');                           \
+		printf(s, __VA_ARGS__);                     \
+	}
+#define DINDENT(x)                    \
+	if (radix_tree_should_print) {    \
+		for (int i = 0; i < (x); i++) \
+			putchar(' ');             \
+	}
+#define DINC(x) radix_tree_indent += x;
+#define DDEC(x) radix_tree_indent -= x;
+#else
+#define DINDENT(x)
+#define D(s, ...)
+#define DINC(x)
+#define DDEC(x)
+#endif
+
+static void __attribute__((noreturn)) radix_tree_panic(const char *fmt, ...)
+{
+	va_list args;
+	va_start(args, fmt);
+	char buf[256];
+	vsnprintf(buf, sizeof(buf), fmt, args);
+	va_end(args);
+	abort_with_reason(OS_REASON_TEST, 0, buf, 0);
+}
+
+struct interval {
+	uint64_t start;
+	uint64_t size;
+};
+
+struct answer {
+	struct interval interval;
+	uint64_t stackid;
+	uint64_t limit;
+};
+
+static inline bool
+in_interval(uint64_t x, struct interval interval)
+{
+	return x >= interval.start && ((x - interval.start) < interval.size);
+}
+
+static inline bool
+intervals_intersect(struct interval a, struct interval b)
+{
+	if (a.size == 0 || b.size == 0)
+		return false;
+	return (in_interval(a.start, b) || in_interval(a.start + a.size - 1, b) || in_interval(b.start, a) ||
+			in_interval(b.start + b.size - 1, a));
+}
+
+__unused static inline bool
+interval_is_subset(struct interval a, struct interval b)
+{
+	return in_interval(a.start, b) && in_interval(a.start + a.size - 1, b);
+}
+
+static inline struct interval
+truncate_interval(struct interval a, uint64_t limit)
+{
+	if (a.start >= limit) {
+		return (struct interval){.start = a.start, .size = 0};
+	} else {
+		return (struct interval){.start = a.start, .size = limit - a.start};
+	}
+}
+
+static inline bool
+answer_found(struct answer answer)
+{
+	return answer.stackid != radix_tree_invalid_value;
+}
+
+/*
+ * Modify the node to maintain the invariant that the lesser edge is first.
+ * Return true if node needed to be modified.
+ */
+static bool
+fixnode(struct radix_node *node)
+{
+	bool swap = false;
+	if (node->edges[0].labelBits && node->edges[1].labelBits) {
+		unsigned label0 = node->edges[0].label << (RADIX_LABEL_BITS - node->edges[0].labelBits);
+		unsigned label1 = node->edges[1].label << (RADIX_LABEL_BITS - node->edges[1].labelBits);
+		if (label1 < label0)
+			swap = true;
+	} else if (node->edges[0].labelBits == 0 && node->edges[1].labelBits != 0) {
+		swap = true;
+	}
+	if (swap) {
+		struct radix_edge edge0 = node->edges[0];
+		node->edges[0] = node->edges[1];
+		node->edges[1] = edge0;
+	}
+	return swap;
+}
+
+static struct answer
+radix_tree_lookup_recursive(struct radix_tree *tree,
+		struct interval keys,	 // keys we're looking for
+		struct interval nodekeys, // keys it is possible that we will find
+		struct radix_node *node,
+		int keyshift)
+{
+	DINDENT(keyshift);
+	D("LOOKUPREC %p keys=[%llx-%llx] nodekeys=[%llx-%llx]\n", node, keys.start, keys.start + keys.size, nodekeys.start,
+			nodekeys.start + nodekeys.size);
+
+	assert(node);
+	assert(intervals_intersect(nodekeys, keys));
+	assert(keyshift < RADIX_TREE_KEY_BITS);
+	assert(!fixnode(node));
+
+	if (keys.start < nodekeys.start) {
+		uint64_t diff = nodekeys.start - keys.start;
+		if (keys.size > diff) {
+			keys.start += diff;
+			keys.size -= diff;
+			assert(keys.start == nodekeys.start);
+		} else {
+			DINDENT(keyshift);
+			D("LOOKUPREC(<) quit keys.size=%llx diff=%llx\n", keys.size, diff);
+			return (struct answer){.limit = nodekeys.start, .stackid = radix_tree_invalid_value};
+		}
+		DINDENT(keyshift);
+		D("LOOKUPREC(<) %p keys=[%llx-%llx]\n", node, keys.start, keys.start + keys.size);
+	}
+
+	assert(keys.start >= nodekeys.start);
+	assert(intervals_intersect(nodekeys, keys));
+
+	for (int i = 1; i >= 0; i--) {
+		struct radix_edge *edge = &node->edges[i];
+		if (!edge_valid(edge)) {
+			continue;
+		}
+		uint64_t edgekeys_start = extend_key(nodekeys.start, edge->labelBits, keyshift, edge->label);
+		assert(edgekeys_start >= nodekeys.start);
+		struct interval edgekeys = {.start = edgekeys_start, .size = nodekeys.size - (edgekeys_start - nodekeys.start)};
+
+		DINDENT(keyshift);
+		D("LOOKUPREC(edge%d) edgekeys=[%llx-%llx] nodekeys=[%llx-%llx]\n", i, edgekeys.start, edgekeys.start + edgekeys.size,
+				nodekeys.start, nodekeys.start + nodekeys.size);
+
+		assert(interval_is_subset(edgekeys, nodekeys));
+		if (intervals_intersect(edgekeys, keys)) {
+			if (edge->isLeaf) {
+				struct radix_node *leaf = getnode(tree, edge->index);
+				assert(leaf);
+				assert(keyshift + edge->labelBits == RADIX_TREE_KEY_BITS);
+				uint64_t size = leaf_size(tree, leaf);
+				assert(size <= edgekeys.size);
+				edgekeys.size = size; // edgekeys is now exact.
+				if (intervals_intersect(edgekeys, keys)) {
+					DINDENT(keyshift);
+					D("LOOKUPREC(found) leaf=(%d)%p %llx\n", edge->index, leaf, leaf->stackid);
+					return (struct answer){.interval = edgekeys, .stackid = leaf->stackid};
+				}
+				nodekeys = truncate_interval(nodekeys, edgekeys.start);
+			} else {
+				struct answer answer =
+						radix_tree_lookup_recursive(tree, keys, edgekeys, getnode(tree, edge->index), keyshift + edge->labelBits);
+				if (answer_found(answer)) {
+					DINDENT(keyshift);
+					D("LOOKUPREC(found) %llx\n", answer.stackid);
+					return answer;
+				}
+				nodekeys = truncate_interval(nodekeys, answer.limit);
+			}
+		}
+	}
+
+	struct answer ans = {.limit = nodekeys.start + nodekeys.size, .stackid = radix_tree_invalid_value};
+	DINDENT(keyshift);
+	D("LOOKUPREC(notfound) limit=%llx\n", ans.limit);
+	return ans;
+}
+
+static struct answer
+radix_tree_lookup_interval(struct radix_tree *tree, struct interval keys)
+{
+	struct interval max_interval = {.start = 0, .size = (uint64_t)-1};
+	struct answer ans = radix_tree_lookup_recursive(tree, keys, max_interval, getnode(tree, 0), 0);
+	D("LOOKUP [%llx-%llx] -> [%llx, %llx] %llx\n", keys.start, keys.start + keys.size, ans.interval.start,
+			ans.interval.start + ans.interval.size, ans.stackid);
+	assert(!answer_found(ans) || intervals_intersect(keys, ans.interval));
+	return ans;
+}
+
+uint64_t
+radix_tree_lookup(struct radix_tree *tree, uint64_t key)
+{
+	return radix_tree_lookup_interval(tree, (struct interval){.start = key, .size = 1}).stackid;
+}
+
+static void radix_tree_grow(struct radix_tree **treep);
+
+static unsigned
+radix_tree_allocate_node(struct radix_tree **treep)
+{
+	if (!(*treep)->next_free)
+		radix_tree_grow(treep);
+
+	if (!(*treep)->next_free)
+		return 0;
+
+	struct radix_tree *tree = *treep;
+
+	unsigned ret = tree->next_free;
+	struct radix_node *node = getnode(tree, tree->next_free);
+	assert(node);
+	tree->next_free = node->next_free;
+	if (node->next_free && !node->next_free_is_initialized) {
+		struct radix_node *next = getnode(tree, node->next_free);
+		assert(next);
+		next->next_free = (node->next_free + 1 < tree->num_nodes) ? node->next_free + 1 : 0;
+	}
+	node->as_u64 = 0;
+	return ret;
+}
+
+static void
+radix_tree_free_node(struct radix_tree *tree, unsigned index)
+{
+	struct radix_node *node = getnode(tree, index);
+	assert(node);
+	node->next_free = tree->next_free;
+	node->next_free_is_initialized = true;
+	tree->next_free = index;
+}
+
+static bool
+radix_tree_insert_recursive(struct radix_tree **treep, struct interval keys, uint64_t value, unsigned node_index, int keyshift)
+{
+	struct radix_node *node = getnode(*treep, node_index);
+
+	DINDENT(keyshift);
+	D("INSERTREC %p keys=[%llx-%llx]\n", node, keys.start, keys.start + keys.size);
+
+	assert(keyshift < RADIX_TREE_KEY_BITS);
+	assert(node);
+
+	for (int i = 0; i < 2; i++) {
+		struct radix_edge *edge = &node->edges[i];
+		int matching_bits = count_matching_bits(edge, keys.start, keyshift);
+		if (matching_bits) {
+			if (matching_bits == edge->labelBits) {
+				if (edge->isLeaf) {
+					assert(false); // it should have been deleted before we got here
+					struct radix_node *leaf = getnode(*treep, edge->index);
+					assert(leaf);
+					assert(keyshift + edge->labelBits == RADIX_TREE_KEY_BITS);
+					leaf->stackid = value;
+					set_leaf_size(*treep, leaf, keys.size);
+					DINDENT(keyshift);
+					D("inserted %p\n", node);
+					return true;
+				} else {
+					return radix_tree_insert_recursive(treep, keys, value, edge->index, keyshift + edge->labelBits);
+				}
+			} else {
+				unsigned index = radix_tree_allocate_node(treep);
+				if (!index) {
+					DINDENT(keyshift);
+					D("FAILED! %p\n", node);
+					return false;
+				}
+				/* pointers may have changed */
+				node = getnode(*treep, node_index);
+				edge = &node->edges[i];
+
+				struct radix_node *newnode = getnode(*treep, index);
+				DINDENT(keyshift);
+				D("splitting edge newnode=%p isleaf=%s matching_bits=%d oldLabelBits=%d\n", newnode,
+						edge->isLeaf ? "true" : "false", matching_bits, edge->labelBits);
+				newnode->edges[0].labelBits = (edge->labelBits - matching_bits);
+				newnode->edges[0].isLeaf = edge->isLeaf;
+				newnode->edges[0].index = edge->index;
+				newnode->edges[0].label = edge->label & ((1 << (edge->labelBits - matching_bits)) - 1);
+
+				edge->label = edge->label >> (edge->labelBits - matching_bits);
+				edge->labelBits = matching_bits;
+				edge->isLeaf = false;
+				edge->index = index;
+
+				fixnode(node);
+				return radix_tree_insert_recursive(treep, keys, value, index, keyshift + matching_bits);
+			}
+		}
+		if (edge->labelBits == 0) {
+			if (RADIX_TREE_KEY_BITS - keyshift <= RADIX_LABEL_BITS) {
+				unsigned index = radix_tree_allocate_node(treep);
+				if (!index) {
+					DINDENT(keyshift);
+					D("FAILED! %p\n", node);
+					return false;
+				}
+				/* pointers may have changed */
+				node = getnode(*treep, node_index);
+				edge = &node->edges[i];
+
+				edge->labelBits = RADIX_TREE_KEY_BITS - keyshift;
+				edge->isLeaf = true;
+				edge->index = index;
+				edge->label = keybits(keys.start, RADIX_TREE_KEY_BITS - keyshift, keyshift);
+				struct radix_node *leaf = getnode(*treep, index);
+				DINDENT(keyshift);
+				D("new leaf node %p\n", leaf);
+				leaf->stackid = value;
+				set_leaf_size(*treep, leaf, keys.size);
+				fixnode(node);
+				return true;
+			} else {
+				unsigned index = radix_tree_allocate_node(treep);
+				if (!index) {
+					DINDENT(keyshift);
+					D("FAILED! %p\n", node);
+					return false;
+				}
+				/* pointers may have changed */
+				node = getnode(*treep, node_index);
+				edge = &node->edges[i];
+
+				edge->labelBits = RADIX_LABEL_BITS;
+				edge->isLeaf = false;
+				edge->index = index;
+				edge->label = keybits(keys.start, RADIX_LABEL_BITS, keyshift);
+
+				struct radix_node *newnode = getnode(*treep, index);
+				newnode->as_u64 = 0;
+				DINDENT(keyshift);
+				D("new internal node %p\n", newnode);
+
+				fixnode(node);
+				return radix_tree_insert_recursive(treep, keys, value, index, keyshift + RADIX_LABEL_BITS);
+			}
+		}
+	}
+
+	radix_tree_panic("MallocStackLogging INTERNAL ERROR: at least one edge must prefix-match or be unused");
+}
+
+bool
+radix_tree_insert(struct radix_tree **treep, uint64_t key, uint64_t size, uint64_t value)
+{
+	D("INSERT %llx-%llx\n", key, key + size);
+	DINC(4);
+	if (key + size < key) {
+		radix_tree_panic("MallocStackLogging INTERNAL ERROR: interval wraps around the end of the address space: %llx, size=%llx\n",
+				key, size);
+	}
+	struct radix_node node = {.stackid = value, .size = size >> (*treep)->leaf_size_shift};
+	if (node.stackid != value || (((uint64_t)node.size) << (*treep)->leaf_size_shift) != size) {
+		radix_tree_panic("MallocStackLogging INTERNAL ERROR: cannot represent value:%llx or size:%llx (key is %llx)\n", value, size, key);
+		return false;
+	}
+	uint64_t mask = ((uint64_t)-1) << (64 - RADIX_TREE_KEY_BITS);
+	if ((key & mask) != key) {
+		radix_tree_panic("MallocStackLogging INTERNAL ERROR: cannot represent key: %llx\n", key);
+	}
+	bool ok;
+	ok = radix_tree_delete(treep, key, size);
+	if (!ok) {
+		goto out;
+	}
+	ok = radix_tree_insert_recursive(treep, (struct interval){.start = key, .size = size}, value, 0, 0);
+out:
+	DDEC(4);
+	return ok;
+}
+
+static bool
+radix_tree_delete_recursive(struct radix_tree *tree, uint64_t key, struct radix_node *node, int keyshift)
+{
+	assert(keyshift < RADIX_TREE_KEY_BITS);
+	assert(node);
+
+	for (int i = 0; i < 2; i++) {
+		struct radix_edge *edge = &node->edges[i];
+		if (edge_matches(edge, key, keyshift)) {
+			if (edge->isLeaf) {
+				radix_tree_free_node(tree, edge->index);
+				if (i == 0) {
+					node->edges[0] = node->edges[1];
+					node->edges[1].labelBits = 0;
+				} else {
+					node->edges[1].labelBits = 0;
+				}
+				return true;
+			} else {
+				bool deleted = radix_tree_delete_recursive(tree, key, getnode(tree, edge->index), keyshift + edge->labelBits);
+				if (deleted) {
+					struct radix_node *child = getnode(tree, edge->index);
+					assert(child);
+					if (child->edges[0].labelBits == 0 && child->edges[1].labelBits == 0) {
+						radix_tree_free_node(tree, edge->index);
+						if (i == 0) {
+							node->edges[0] = node->edges[1];
+							node->edges[1].labelBits = 0;
+						} else {
+							node->edges[1].labelBits = 0;
+						}
+					}
+				}
+				return deleted;
+			}
+		}
+	}
+	return false;
+}
+
+bool
+radix_tree_delete(struct radix_tree **treep, uint64_t key, uint64_t size)
+{
+	D("BALETE %llx-%llx\n", key, key + size);
+	DINC(4);
+	struct interval keys = {.start = key, .size = size};
+	bool ok = true;
+	while (1) {
+		struct answer answer = radix_tree_lookup_interval(*treep, keys);
+		if (!answer_found(answer)) {
+			break;
+		}
+		ok = radix_tree_delete_recursive(*treep, answer.interval.start, getnode(*treep, 0), 0);
+		assert(ok);
+		D("BALETED %llx-%llx -> %llx\n", answer.interval.start, answer.interval.start + answer.interval.size, answer.stackid);
+		if (answer.interval.start < keys.start) {
+			D("REINSERTING %llx-%llx -> %llx\n", answer.interval.start,
+					answer.interval.start + (keys.start - answer.interval.start), answer.stackid);
+			ok = radix_tree_insert(treep, answer.interval.start, keys.start - answer.interval.start, answer.stackid);
+			if (!ok) {
+				goto out;
+			}
+		}
+		uint64_t answer_end = answer.interval.start + answer.interval.size;
+		uint64_t keys_end = keys.start + keys.size;
+		if (answer_end > keys_end) {
+			D("REINSERTING %llx-%llx -> %llx\n", keys_end, keys_end + (answer_end - keys_end), answer.stackid);
+			ok = radix_tree_insert(treep, keys_end, answer_end - keys_end, answer.stackid);
+			if (!ok) {
+				goto out;
+			}
+		}
+	}
+out:
+	DDEC(4);
+	return ok;
+}
+
+struct radix_tree *
+radix_tree_init(void *buf, size_t size)
+{
+	struct radix_tree *tree = buf;
+	memcpy(tree->header, "radixv2", 8);
+	void *nodestart = &tree->nodes[0];
+	void *nodesend = buf + size;
+	assert(nodestart < nodesend);
+	tree->num_nodes = (uint32_t)(nodesend - nodestart) / sizeof(struct radix_node);
+	assert(tree->num_nodes >= 3);
+	tree->next_free = 1;
+	tree->nodes[0].as_u64 = tree->nodes[1].as_u64 = 0;
+	tree->nodes[1].next_free = 2;
+	tree->leaf_size_shift = 12; // smallest size of a VM region is 4096
+	return tree;
+}
+
+struct radix_tree *
+radix_tree_create()
+{
+	mach_vm_size_t size = PAGE_SIZE;
+	mach_vm_address_t allocated;
+	kern_return_t kr =
+			mach_vm_allocate(mach_task_self(), &allocated, size, VM_FLAGS_ANYWHERE | VM_MAKE_TAG(VM_MEMORY_ANALYSIS_TOOL));
+	if (kr != KERN_SUCCESS) {
+		return NULL;
+	}
+	return radix_tree_init((void *)allocated, PAGE_SIZE);
+	return NULL;
+}
+
+static void
+radix_tree_grow(struct radix_tree **treep)
+{
+	mach_vm_size_t max_size = (1 << 16) * sizeof(struct radix_node);
+
+	assert((*treep)->next_free == 0);
+	mach_vm_size_t size = sizeof(struct radix_tree) + sizeof(struct radix_node) * (*treep)->num_nodes;
+	assert(size % PAGE_SIZE == 0);
+	mach_vm_size_t newsize = size * 2;
+	if (newsize > max_size) {
+		newsize = max_size;
+	}
+	if (newsize <= size) {
+		return;
+	}
+	mach_vm_address_t allocated;
+	kern_return_t kr =
+			mach_vm_allocate(mach_task_self(), &allocated, newsize, VM_FLAGS_ANYWHERE | VM_MAKE_TAG(VM_MEMORY_ANALYSIS_TOOL));
+	if (kr != KERN_SUCCESS) {
+		return;
+	}
+	D("GROW %p -> %p\n", *treep, (void *)allocated);
+	kr = mach_vm_copy(mach_task_self(), (mach_vm_address_t)*treep, size, allocated);
+	if (kr != KERN_SUCCESS) {
+		mach_vm_deallocate(mach_task_self(), allocated, newsize);
+		return;
+	}
+	uint32_t old_num_nodes = (*treep)->num_nodes;
+	mach_vm_deallocate(mach_task_self(), (mach_vm_address_t)*treep, size);
+	*treep = (void *)allocated;
+
+	void *nodestart = &(*treep)->nodes[0];
+	void *nodesend = ((void *)(*treep)) + newsize;
+	(*treep)->num_nodes = (uint32_t)(nodesend - nodestart) / sizeof(struct radix_node);
+	(*treep)->next_free = old_num_nodes;
+
+	(*treep)->nodes[old_num_nodes].next_free_is_initialized = 0;
+	(*treep)->nodes[old_num_nodes].next_free = old_num_nodes + 1;
+}
+
+void
+radix_tree_destory(struct radix_tree *tree)
+{
+	mach_vm_size_t size = sizeof(struct radix_tree) + sizeof(struct radix_node) * tree->num_nodes;
+	assert(size % PAGE_SIZE == 0);
+	mach_vm_deallocate(mach_task_self(), (mach_vm_address_t)tree, size);
+}
+
+static uint64_t
+radix_tree_count_recursive(struct radix_tree *tree, struct radix_node *node)
+{
+	uint64_t count = 0;
+	for (int i = 0; i < 2; i++) {
+		struct radix_edge *edge = &node->edges[i];
+		if (edge->labelBits == 0)
+			continue;
+		if (edge->isLeaf) {
+			count += leaf_size(tree, getnode(tree, edge->index));
+		} else {
+			count += radix_tree_count_recursive(tree, getnode(tree, edge->index));
+		}
+	}
+	return count;
+}
+
+uint64_t
+radix_tree_count(struct radix_tree *tree)
+{
+	return radix_tree_count_recursive(tree, getnode(tree, 0));
+}
+
+uint64_t
+radix_tree_size(struct radix_tree *tree)
+{
+	mach_vm_size_t size = sizeof(struct radix_tree) + sizeof(struct radix_node) * tree->num_nodes;
+	return size;
+}