make shape code use the tydesc found in the box, not the shape str

3 files changed, 88 insertions, 46 deletions

diff --git a/src/rt/rust_cc.cpp b/src/rt/rust_cc.cpp
index 03a4bed014e..ec54b1ef181 100644
--- a/src/rt/rust_cc.cpp
+++ b/src/rt/rust_cc.cpp
@@ -93,7 +93,11 @@ class irc : public shape::data<irc,shape::ptr> {
     }
 
     void walk_box2() {
-        shape::data<irc,shape::ptr>::walk_box_contents1();
+        // the box ptr can be NULL for env ptrs in closures and data
+        // not fully initialized
+        rust_opaque_box *box = *(rust_opaque_box**)dp;
+        if (box)
+            shape::data<irc,shape::ptr>::walk_box_contents1();
     }
 
     void walk_uniq2() {
@@ -103,8 +107,6 @@ class irc : public shape::data<irc,shape::ptr> {
     void walk_fn2(char code) {
         switch (code) {
           case shape::SHAPE_BOX_FN: {
-              // Record an irc for the environment box, but don't descend
-              // into it since it will be walked via the box's allocation
               shape::bump_dp<void*>(dp); // skip over the code ptr
               walk_box2();               // walk over the environment ptr
               break;
@@ -137,19 +139,19 @@ class irc : public shape::data<irc,shape::ptr> {
 
     void walk_uniq_contents2(irc &sub) { sub.walk(); }
 
-    void walk_box_contents2(irc &sub, shape::ptr &box_dp) {
-        maybe_record_irc(box_dp);
+    void walk_box_contents2(irc &sub) {
+        maybe_record_irc();
 
         // Do not traverse the contents of this box; it's in the allocation
         // somewhere, so we're guaranteed to come back to it (if we haven't
         // traversed it already).
     }
 
-    void maybe_record_irc(shape::ptr &box_dp) {
-        if (!box_dp)
-            return;
+    void maybe_record_irc() {
+        rust_opaque_box *box_ptr = *(rust_opaque_box **) dp;
 
-        rust_opaque_box *box_ptr = (rust_opaque_box *) box_dp;
+        if (!box_ptr)
+            return;
 
         // Bump the internal reference count of the box.
         if (ircs.find(box_ptr) == ircs.end()) {
@@ -326,7 +328,11 @@ class mark : public shape::data<mark,shape::ptr> {
     }
 
     void walk_box2() {
-        shape::data<mark,shape::ptr>::walk_box_contents1();
+        // the box ptr can be NULL for env ptrs in closures and data
+        // not fully initialized
+        rust_opaque_box *box = *(rust_opaque_box**)dp;
+        if (box)
+            shape::data<mark,shape::ptr>::walk_box_contents1();
     }
 
     void walk_uniq2() {
@@ -336,8 +342,6 @@ class mark : public shape::data<mark,shape::ptr> {
     void walk_fn2(char code) {
         switch (code) {
           case shape::SHAPE_BOX_FN: {
-              // Record an irc for the environment box, but don't descend
-              // into it since it will be walked via the box's allocation
               shape::data<mark,shape::ptr>::walk_fn_contents1();
               break;
           }
@@ -369,11 +373,11 @@ class mark : public shape::data<mark,shape::ptr> {
 
     void walk_uniq_contents2(mark &sub) { sub.walk(); }
 
-    void walk_box_contents2(mark &sub, shape::ptr &box_dp) {
-        if (!box_dp)
-            return;
+    void walk_box_contents2(mark &sub) {
+        rust_opaque_box *box_ptr = *(rust_opaque_box **) dp;
 
-        rust_opaque_box *box_ptr = (rust_opaque_box *) box_dp;
+        if (!box_ptr)
+            return;
 
         if (marked.find(box_ptr) != marked.end())
             return; // Skip to avoid chasing cycles.
@@ -516,7 +520,9 @@ class sweep : public shape::data<sweep,shape::ptr> {
     }
 
     void walk_box2() {
-        shape::data<sweep,shape::ptr>::walk_box_contents1();
+        // In sweep phase, do not walk the box contents.  There is an
+        // outer loop walking all remaining boxes, and this box may well
+        // have been freed already!
     }
 
     void walk_fn2(char code) {
@@ -524,14 +530,16 @@ class sweep : public shape::data<sweep,shape::ptr> {
           case shape::SHAPE_UNIQ_FN: {
               fn_env_pair pair = *(fn_env_pair*)dp;
 
-              // free closed over data:
-              shape::data<sweep,shape::ptr>::walk_fn_contents1();
-
-              // now free the embedded type descr:
-              upcall_s_free_shared_type_desc((type_desc*)pair.env->td);
-
-              // now free the ptr:
-              task->kernel->free(pair.env);
+              if (pair.env) {
+                  // free closed over data:
+                  shape::data<sweep,shape::ptr>::walk_fn_contents1();
+                  
+                  // now free the embedded type descr:
+                  upcall_s_free_shared_type_desc((type_desc*)pair.env->td);
+                  
+                  // now free the ptr:
+                  task->kernel->free(pair.env);
+              }
               break;
           }
           case shape::SHAPE_BOX_FN: {
@@ -579,10 +587,6 @@ class sweep : public shape::data<sweep,shape::ptr> {
 
     void walk_uniq_contents2(sweep &sub) { sub.walk(); }
 
-    void walk_box_contents2(sweep &sub, shape::ptr &box_dp) {
-        return;
-    }
-
     void walk_struct2(const uint8_t *end_sp) {
         while (this->sp != end_sp) {
             this->walk();
diff --git a/src/rt/rust_shape.cpp b/src/rt/rust_shape.cpp
index 20f8bef89aa..866bc4ac17a 100644
--- a/src/rt/rust_shape.cpp
+++ b/src/rt/rust_shape.cpp
@@ -264,7 +264,7 @@ private:
         result = sub.result;
     }
 
-    inline void walk_box_contents2(cmp &sub, ptr_pair &box_dp) {
+    inline void walk_box_contents2(cmp &sub) {
         sub.align = true;
         sub.walk();
         result = sub.result;
@@ -310,6 +310,15 @@ public:
       result(0) {}
 
     cmp(const cmp &other,
+        const uint8_t *in_sp,
+        const type_param *in_params,
+        const rust_shape_tables *in_tables,
+        ptr_pair &in_dp)
+    : data<cmp,ptr_pair>(other.task, other.align, in_sp, in_params, in_tables,
+                         in_dp),
+      result(0) {}
+
+    cmp(const cmp &other,
         const uint8_t *in_sp = NULL,
         const type_param *in_params = NULL,
         const rust_shape_tables *in_tables = NULL)
diff --git a/src/rt/rust_shape.h b/src/rt/rust_shape.h
index 92f660db741..675d67f12d4 100644
--- a/src/rt/rust_shape.h
+++ b/src/rt/rust_shape.h
@@ -712,6 +712,17 @@ public:
     inline operator bool() const { return p != NULL; }
     inline operator uintptr_t() const { return (uintptr_t)p; }
 
+    inline const type_desc *box_body_td() const {
+        rust_opaque_box *box = *reinterpret_cast<rust_opaque_box**>(p);
+        assert(box->ref_count >= 1);
+        return box->td;
+    }
+
+    inline ptr box_body() const {
+        rust_opaque_box *box = *reinterpret_cast<rust_opaque_box**>(p);
+        return make((uint8_t*)::box_body(box));
+    }
+
     static inline ptr make(uint8_t *in_p) {
         ptr self(in_p);
         return self;
@@ -746,7 +757,7 @@ public:
     inline void operator=(const T rhs) { fst = snd = rhs; }
 
     static data_pair<T> make(T &fst, T &snd) {
-        data_pair<T> data(fst, snd);
+          data_pair<T> data(fst, snd);
         return data;
     }
 };
@@ -792,6 +803,25 @@ public:
         ptr_pair self(pair.fst, pair.snd);
         return self;
     }
+
+    inline const type_desc *box_body_td() const {
+        // Here we assume that the two ptrs are both boxes with
+        // equivalent type descriptors.  This is safe because we only
+        // use ptr_pair in the cmp glue, and we only use the cmp glue
+        // when rust guarantees us that the boxes are of the same
+        // type.  As box types are not opaque to Rust, it is in a
+        // position to make this determination.
+        rust_opaque_box *box_fst = *reinterpret_cast<rust_opaque_box**>(fst);
+        assert(box_fst->ref_count >= 1);
+        return box_fst->td;
+    }
+
+    inline ptr_pair box_body() const {
+        rust_opaque_box *box_fst = *reinterpret_cast<rust_opaque_box**>(fst);
+        rust_opaque_box *box_snd = *reinterpret_cast<rust_opaque_box**>(snd);
+        return make((uint8_t*)::box_body(box_fst),
+                    (uint8_t*)::box_body(box_snd));
+    }
 };
 
 // NB: This function does not align.
@@ -933,18 +963,16 @@ public:
 template<typename T,typename U>
 void
 data<T,U>::walk_box_contents1() {
-    typename U::template data<uint8_t *>::t box_ptr = bump_dp<uint8_t *>(dp);
-    U box_dp(box_ptr);
-
-    // No need to worry about alignment so long as the box header is
-    // a multiple of 16 bytes.  We can just find the body by adding
-    // the size of header to box_dp.
-    assert ((sizeof(rust_opaque_box) % 16) == 0 ||
-            !"Must align to find the box body");
-
-    U body_dp = box_dp + sizeof(rust_opaque_box);
-    T sub(*static_cast<T *>(this), body_dp);
-    static_cast<T *>(this)->walk_box_contents2(sub, box_dp);
+    const type_desc *body_td = dp.box_body_td();
+    if (body_td) {
+        U body_dp(dp.box_body());
+        arena arena;
+        type_param *params = type_param::from_tydesc(body_td, arena);
+        T sub(*static_cast<T *>(this), body_td->shape, params,
+              body_td->shape_tables, body_dp);
+        sub.align = true;
+        static_cast<T *>(this)->walk_box_contents2(sub);
+    }
 }
 
 template<typename T,typename U>
@@ -1006,7 +1034,7 @@ data<T,U>::walk_tag1(tag_info &tinfo) {
 
 template<typename T,typename U>
 void
-data<T,U>::walk_fn_contents1() {
+  data<T,U>::walk_fn_contents1() {
     fn_env_pair pair = bump_dp<fn_env_pair>(dp);
     if (!pair.env)
         return;
@@ -1119,9 +1147,10 @@ private:
 
     void walk_subcontext2(log &sub) { sub.walk(); }
 
-    void walk_box_contents2(log &sub, ptr &ref_count_dp) {
+    void walk_box_contents2(log &sub) {
         out << prefix;
-        if (!ref_count_dp) {
+        rust_opaque_box *box_ptr = *(rust_opaque_box **) dp;
+        if (!box_ptr) {
             out << "(null)";
         } else {
             sub.align = true;