rt: Allow iteration over the dynastack

2 files changed, 109 insertions, 17 deletions

diff --git a/src/rt/rust_obstack.cpp b/src/rt/rust_obstack.cpp
index 46fca601ca9..81d0076a37c 100644
--- a/src/rt/rust_obstack.cpp
+++ b/src/rt/rust_obstack.cpp
@@ -25,28 +25,14 @@ const size_t DEFAULT_ALIGNMENT = 16;
 struct rust_obstack_alloc {
     size_t len;
     const type_desc *tydesc;
+    uint32_t pad0;  // FIXME: x86-specific
+    uint32_t pad1;
     uint8_t data[];
 
     rust_obstack_alloc(size_t in_len, const type_desc *in_tydesc)
     : len(in_len), tydesc(in_tydesc) {}
 };
 
-// A contiguous set of allocations.
-struct rust_obstack_chunk {
-    rust_obstack_chunk *prev;
-    size_t size;
-    size_t alen;
-    size_t pad;
-    uint8_t data[];
-
-    rust_obstack_chunk(rust_obstack_chunk *in_prev, size_t in_size)
-    : prev(in_prev), size(in_size), alen(0) {}
-
-    void *alloc(size_t len, type_desc *tydesc);
-    bool free(void *ptr);
-    void *mark();
-};
-
 void *
 rust_obstack_chunk::alloc(size_t len, type_desc *tydesc) {
     alen = align_to(alen, DEFAULT_ALIGNMENT);
@@ -132,3 +118,61 @@ rust_obstack::mark() {
     return chunk ? chunk->mark() : NULL;
 }
 
+
+// Iteration over self-describing obstacks
+
+std::pair<const type_desc *,void *>
+rust_obstack::iterator::operator*() const {
+    return std::make_pair(alloc->tydesc, alloc->data);
+}
+
+rust_obstack::iterator &
+rust_obstack::iterator::operator++() {
+    uint8_t *adata = align_to(alloc->data + alloc->len, DEFAULT_ALIGNMENT);
+    alloc = reinterpret_cast<rust_obstack_alloc *>(adata);
+    if (reinterpret_cast<uint8_t *>(alloc) >= chunk->data + chunk->alen) {
+        // We reached the end of this chunk; go on to the next one.
+        chunk = chunk->prev;
+        if (chunk)
+            alloc = reinterpret_cast<rust_obstack_alloc *>(chunk->data);
+        else
+            alloc = NULL;
+    }
+    return *this;
+}
+
+bool
+rust_obstack::iterator::operator==(const rust_obstack::iterator &other)
+        const {
+    return chunk == other.chunk && alloc == other.alloc;
+}
+
+bool
+rust_obstack::iterator::operator!=(const rust_obstack::iterator &other)
+        const {
+    return !(*this == other);
+}
+
+
+// Debugging
+
+void
+rust_obstack::dump() const {
+    iterator b = begin(), e = end();
+    while (b != e) {
+        std::pair<const type_desc *,void *> data = *b;
+        shape::arena arena;
+        shape::type_param *params = shape::type_param::from_tydesc(data.first,
+                                                                   arena);
+        shape::log log(task, true, data.first->shape, params,
+                       data.first->shape_tables,
+                       reinterpret_cast<uint8_t *>(data.second), std::cerr);
+        log.walk();
+        std::cerr << "\n";
+
+        ++b;
+    }
+
+    std::cerr << "end of dynastack dump\n";
+}
+
diff --git a/src/rt/rust_obstack.h b/src/rt/rust_obstack.h
index 6483a5f8d11..99d9ba801b4 100644
--- a/src/rt/rust_obstack.h
+++ b/src/rt/rust_obstack.h
@@ -3,10 +3,28 @@
 #ifndef RUST_OBSTACK_H
 #define RUST_OBSTACK_H
 
-struct rust_obstack_chunk;
+#include <utility>
+
+struct rust_obstack_alloc;
 struct rust_task;
 struct type_desc;
 
+// A contiguous set of allocations.
+struct rust_obstack_chunk {
+    rust_obstack_chunk *prev;
+    size_t size;
+    size_t alen;
+    size_t pad;
+    uint8_t data[];
+
+    rust_obstack_chunk(rust_obstack_chunk *in_prev, size_t in_size)
+    : prev(in_prev), size(in_size), alen(0) {}
+
+    void *alloc(size_t len, type_desc *tydesc);
+    bool free(void *ptr);
+    void *mark();
+};
+
 class rust_obstack {
     rust_obstack_chunk *chunk;
     rust_task *task;
@@ -15,12 +33,42 @@ class rust_obstack {
     void *alloc_new(size_t len, type_desc *tydesc);
 
 public:
+    class iterator {
+        rust_obstack_chunk *chunk;
+        rust_obstack_alloc *alloc;
+
+    public:
+        iterator(rust_obstack_chunk *in_chunk)
+        : chunk(in_chunk),
+          alloc(in_chunk
+                ? reinterpret_cast<rust_obstack_alloc *>(in_chunk->data)
+                : NULL) {}
+
+        std::pair<const type_desc *,void *> operator*() const;
+        iterator &operator++();
+        bool operator==(const iterator &other) const;
+        bool operator!=(const iterator &other) const;
+    };
+
     rust_obstack(rust_task *in_task) : chunk(NULL), task(in_task) {}
     ~rust_obstack();
 
+    inline iterator begin() const {
+        iterator it(chunk);
+        return it;
+    }
+
+    inline iterator end() const {
+        iterator it(NULL);
+        return it;
+    }
+
     void *alloc(size_t len, type_desc *tydesc);
     void free(void *ptr);
     void *mark();
+
+    /** Debugging tool: dumps the contents of this obstack to stderr. */
+    void dump() const;
 };
 
 #endif