From 2b8985608d33abaae7b201a008e292cbbe2167ef Mon Sep 17 00:00:00 2001
From: Matthew Sotoudeh <matthew@masot.net>
Date: Thu, 27 Jul 2023 14:26:33 -0700
Subject: add automated refcounting pass

---
 python/examples/refcounting/runtime/refcounting.c | 114 ++++++++++++++++++++++
 1 file changed, 114 insertions(+)
 create mode 100644 python/examples/refcounting/runtime/refcounting.c

(limited to 'python/examples/refcounting/runtime/refcounting.c')

diff --git a/python/examples/refcounting/runtime/refcounting.c b/python/examples/refcounting/runtime/refcounting.c
new file mode 100644
index 0000000..ba8cacc
--- /dev/null
+++ b/python/examples/refcounting/runtime/refcounting.c
@@ -0,0 +1,114 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <assert.h>
+
+struct data {
+  void *start;
+  size_t len;
+  size_t refcount;
+};
+
+struct node {
+  struct data data;
+
+  struct node *left;
+  struct node *right;
+};
+
+static struct node *ROOT = 0;
+static int TOTAL = 0;
+
+static struct node *lookup_interval(struct node *root, void *ptr) {
+  if (!root) return 0;
+
+  if (ptr < root->data.start)
+    return lookup_interval(root->left, ptr);
+
+  if (ptr < (root->data.start + root->data.len))
+    return root;
+
+  // our last hope is something on the right
+  return lookup_interval(root->right, ptr);
+}
+
+// exact lookup, ignoring range len
+static struct node **lookup(void *ptr) {
+  struct node **node = &ROOT;
+  while (*node) {
+    if ((*node)->data.start == ptr)
+      return node;
+    if (ptr > (*node)->data.start)
+      node = &((*node)->right);
+    else
+      node = &((*node)->left);
+  }
+  return node;
+}
+
+static struct node *insert(void *ptr, size_t len) {
+  struct node **loc = lookup(ptr);
+  assert(!(*loc));
+  *loc = calloc(1, sizeof(struct node));
+  (*loc)->data.start = ptr;
+  (*loc)->data.len = len;
+  return *loc;
+}
+
+static void remove_root(struct node **tree) {
+  if (!(*tree)) {
+    return;
+  } else if (!((*tree)->right)) {
+    *tree = (*tree)->left;
+  } else if (!((*tree)->left)) {
+    *tree = (*tree)->right;
+  } else {
+    // swap *tree and (*tree)->right, then delete the root in tree->right
+    struct data root_data = (*tree)->data;
+    (*tree)->data = (*tree)->right->data;
+    (*tree)->right->data = root_data;
+    remove_root(&((*tree)->right));
+  }
+}
+
+static struct node *remove_node(void *ptr) {
+  remove_root(lookup(ptr));
+}
+
+void *refcount_calloc(unsigned long count, unsigned long size) {
+  void *data = calloc(count, size);
+  TOTAL += 1;
+  printf("Allocated: %p ; now tracking %d regions\n", data, TOTAL);
+  insert(data, count * size);
+  return data;
+}
+
+void refcount_write(void *new, void *old) {
+  // printf("Refcount write new: %p     old: %p\n", new, old);
+  if (new) {
+    struct node *new_node = lookup_interval(ROOT, new);
+    if (new_node) {
+      // printf("Incrementing refcount ...\n");
+      new_node->data.refcount++;
+    }
+  }
+  if (old) {
+    struct node *old_node = lookup_interval(ROOT, old);
+    if (old_node) {
+      // printf("Decrementing refcount ...\n");
+      old_node->data.refcount--;
+      if (!(old_node->data.refcount)) {
+        TOTAL -= 1;
+        printf("Freeing! Left: %d\n", TOTAL);
+        remove_node(old_node->data.start);
+        free(old);
+      }
+    }
+  }
+}
+
+void refcount_returning(void *ptr) {
+  if (!ptr) return;
+  struct node *ptr_node = lookup_interval(ROOT, ptr);
+  if (!ptr_node) return;
+  ptr_node->data.refcount--;
+}
-- 
cgit v1.2.3