support arrays of pointers in refcounting

2 files changed, 33 insertions, 21 deletions

diff --git a/python/examples/refcounting/dietpass b/python/examples/refcounting/dietpass
index 9a7b0e8..9d6a77f 100755
--- a/python/examples/refcounting/dietpass
+++ b/python/examples/refcounting/dietpass
@@ -67,8 +67,10 @@ def ptr_exprs(name, ty):
     elif isinstance(ty, (BasicType, FunctionType)):
         return []
     elif isinstance(ty, ArrayType):
-        # TODO
-        return []
+        result = []
+        for i in range(int(ty.length)):
+            result.extend(ptr_exprs(f"(({name})[{i}])", ty.base))
+        return result
     else:
         print(ty.classify(), file=sys.stderr)
         raise NotImplementedError
diff --git a/python/examples/refcounting/test.c b/python/examples/refcounting/test.c
index 212bb8b..cd53e8b 100644
--- a/python/examples/refcounting/test.c
+++ b/python/examples/refcounting/test.c
@@ -2,33 +2,43 @@
 #include <stdio.h>
 
 int *foo(void) {
-  int *ptr = 0;
-  for (int i = 0; i < 5; i++) {
+    int *ptr = 0;
+    for (int i = 0; i < 5; i++)
+        ptr = calloc(1, sizeof(*ptr));
     ptr = calloc(1, sizeof(*ptr));
-  }
-  ptr = calloc(1, sizeof(*ptr));
-  ptr = calloc(1, sizeof(*ptr));
-  *ptr = 5;
-  return ptr;
+    ptr = calloc(1, sizeof(*ptr));
+    *ptr = 5;
+    return ptr;
 }
 
 struct bar {
-  int x;
-  int *y;
+    int x;
+    int *y;
 };
 
 int *bar(void) {
-  struct bar bar;
-  bar.x = 5;
-  bar.y = calloc(1, sizeof(int));
-  *bar.y = bar.x;
-  return bar.y;
+    struct bar bar;
+    bar.x = 5;
+    bar.y = calloc(1, sizeof(int));
+    *bar.y = bar.x;
+    return bar.y;
+}
+
+int *arrays(void) {
+    int *foo[3];
+    foo[0] = calloc(1, sizeof(int));
+    foo[1] = calloc(1, sizeof(int));
+    foo[2] = calloc(1, sizeof(int));
+    *(foo[1]) = 123;
+    return foo[1];
 }
 
 int main(void) {
-  int *x = foo();
-  printf("Result of foo(): %d\n", *x);
-  x = bar();
-  printf("Result of bar(): %d\n", *x);
-  return 0;
+    int *x = foo();
+    printf("Result of foo(): %d\n", *x);
+    x = bar();
+    printf("Result of bar(): %d\n", *x);
+    x = arrays();
+    printf("Result of arrays(): %d\n", *x);
+    return 0;
 }