rts: Refactor unloading of foreign export StablePtrs

Previously we would allocate a linked list cell for each foreign export.
Now we can avoid this by taking advantage of the fact that they are
already broken into groups.

4 files changed, 50 insertions, 37 deletions

diff --git a/includes/rts/ForeignExports.h b/includes/rts/ForeignExports.h
index f8828e59d4..aeb524aebf 100644
--- a/includes/rts/ForeignExports.h
+++ b/includes/rts/ForeignExports.h
@@ -29,6 +29,8 @@ struct ForeignExportsList {
       /* if the RTS linker loaded the module,
        * this points to an array of length ->n_entries
        * recording the StablePtr for each export. */
+    StgStablePtr **stable_ptrs;
+      /* the exported closures. of length ->exports. */
     StgPtr exports[];
 };
 
diff --git a/rts/ForeignExports.c b/rts/ForeignExports.c
index 3195b6e779..e218281b51 100644
--- a/rts/ForeignExports.c
+++ b/rts/ForeignExports.c
@@ -48,12 +48,14 @@ static ObjectCode *loading_obj = NULL;
  * For instance, doing exactly this resulted in #18548.
  *
  * Another consideration here is that the linker needs to know which
- * `StablePtr`s belong to each `ObjectCode` it loads for the sake of unloading.
- * For this reason, the linker informs us when it is loading an object by calling
- * `foreignExportsLoadingObject` and `foreignExportsFinishedLoadingObject`. We
- * take note of the `ObjectCode*` we are loading in `loading_obj` such that we
- * can associate the `StablePtr` with the `ObjectCode` in
- * `processForeignExports`.
+ * `StablePtr`s belong to each `ObjectCode` so it can free them when the module is
+ * unloaded.  For this reason, the linker informs us when it is loading an
+ * object by calling `foreignExportsLoadingObject` and
+ * `foreignExportsFinishedLoadingObject`. We take note of the `ObjectCode*` we
+ * are loading in `loading_obj` such that we can associate the `ForeignExportsList` with
+ * the `ObjectCode` in `processForeignExports`. We then record each of the
+ * StablePtrs we create in the ->stable_ptrs array of ForeignExportsList so
+ * they can be enumerated during unloading.
  *
  */
 
@@ -94,20 +96,35 @@ void foreignExportsFinishedLoadingObject()
 void processForeignExports()
 {
     while (pending) {
-        for (int i=0; i < pending->n_entries; i++) {
-            StgPtr p = pending->exports[i];
-            StgStablePtr *sptr = getStablePtr(p);
+        struct ForeignExportsList *cur = pending;
+        pending = cur->next;
 
-            if (loading_obj != NULL) {
-                ForeignExportStablePtr *fe_sptr = (ForeignExportStablePtr *)
-                  stgMallocBytes(sizeof(ForeignExportStablePtr),
-                                 "foreignExportStablePtr");
-                fe_sptr->stable_ptr = sptr;
-                fe_sptr->next = loading_obj->stable_ptrs;
-                pending->oc->stable_ptrs = fe_sptr;
+        /* sanity check */
+        ASSERT(cur->stable_ptrs == NULL);
+
+        /* N.B. We only need to populate the ->stable_ptrs
+         * array if the object might later be unloaded.
+         */
+        if (cur->oc != NULL) {
+            cur->stable_ptrs =
+              stgMallocBytes(sizeof(StgStablePtr*) * cur->n_entries,
+                             "foreignExportStablePtr");
+
+            for (int i=0; i < cur->n_entries; i++) {
+                StgStablePtr *sptr = getStablePtr(cur->exports[i]);
+
+                if (cur->oc != NULL) {
+                    cur->stable_ptrs[i] = sptr;
+                }
+            }
+            cur->next = cur->oc->foreign_exports;
+            cur->oc->foreign_exports = cur;
+        } else {
+            /* can't be unloaded, don't bother populating
+             * ->stable_ptrs array. */
+            for (int i=0; i < cur->n_entries; i++) {
+                getStablePtr(cur->exports[i]);
             }
         }
-
-        pending = pending->next;
     }
 }
diff --git a/rts/Linker.c b/rts/Linker.c
index fa38480fb6..2d54c29196 100644
--- a/rts/Linker.c
+++ b/rts/Linker.c
@@ -1239,14 +1239,18 @@ static void freeOcStablePtrs (ObjectCode *oc)
 {
     // Release any StablePtrs that were created when this
     // object module was initialized.
-    ForeignExportStablePtr *fe_ptr, *next;
+    struct ForeignExportsList *exports, *next;
 
-    for (fe_ptr = oc->stable_ptrs; fe_ptr != NULL; fe_ptr = next) {
-        next = fe_ptr->next;
-        freeStablePtr(fe_ptr->stable_ptr);
-        stgFree(fe_ptr);
+    for (exports = oc->foreign_exports; exports != NULL; exports = next) {
+        next = exports->next;
+        for (int i = 0; i < exports->n_entries; i++) {
+            freeStablePtr(exports->stable_ptrs[i]);
+        }
+        stgFree(exports->stable_ptrs);
+        exports->stable_ptrs = NULL;
+        exports->next = NULL;
     }
-    oc->stable_ptrs = NULL;
+    oc->foreign_exports = NULL;
 }
 
 static void
@@ -1404,7 +1408,7 @@ mkOc( pathchar *path, char *image, int imageSize,
    oc->n_segments        = 0;
    oc->segments          = NULL;
    oc->proddables        = NULL;
-   oc->stable_ptrs       = NULL;
+   oc->foreign_exports   = NULL;
 #if defined(NEED_SYMBOL_EXTRAS)
    oc->symbol_extras     = NULL;
 #endif
diff --git a/rts/LinkerInternals.h b/rts/LinkerInternals.h
index 1f63f0c485..2e76a888e4 100644
--- a/rts/LinkerInternals.h
+++ b/rts/LinkerInternals.h
@@ -135,17 +135,6 @@ typedef struct _Segment {
     int n_sections;
 } Segment;
 
-/*
- * We must keep track of the StablePtrs that are created for foreign
- * exports by constructor functions when the module is loaded, so that
- * we can free them again when the module is unloaded.  If we don't do
- * this, then the StablePtr will keep the module alive indefinitely.
- */
-typedef struct ForeignExportStablePtr_ {
-    StgStablePtr stable_ptr;
-    struct ForeignExportStablePtr_ *next;
-} ForeignExportStablePtr;
-
 #if defined(powerpc_HOST_ARCH) || defined(x86_64_HOST_ARCH)
 #define NEED_SYMBOL_EXTRAS 1
 #endif
@@ -240,7 +229,8 @@ typedef struct _ObjectCode {
     char* bssBegin;
     char* bssEnd;
 
-    ForeignExportStablePtr *stable_ptrs;
+    /* a list of all ForeignExportsLists owned by this object */
+    struct ForeignExportsList *foreign_exports;
 
     /* Holds the list of symbols in the .o file which
        require extra information.*/