1 files changed, 135 insertions, 32 deletions

diff --git a/rts/CheckUnload.c b/rts/CheckUnload.c
index fa4843d8e4..de8180526b 100644
--- a/rts/CheckUnload.c
+++ b/rts/CheckUnload.c
@@ -38,30 +38,129 @@
 // object as referenced so that it won't get unloaded in this round.
 //
 
-static void checkAddress (HashTable *addrs, const void *addr)
+//
+// In certain circumstances, there may be a lot of unloaded ObjectCode structs
+// chained in `unloaded_objects` (such as when users `:load` a module in a very
+// big repo in GHCi). To speed up checking whether an address lies within any of
+// these objects, we populate the addresses of the their mapped sections in
+// an array sorted by their `start` address and do binary search for our address
+// on that array. Note that this works because the sections are mapped to mutual
+// exclusive memory regions, so we can simply find the largest lower bound among
+// the `start` addresses of the sections and then check if our address is inside
+// that section. In particular, we store the start address and end address of
+// each mapped section in a OCSectionIndex, arrange them all on a contiguous
+// memory range and then sort by start address. We then put this array in an
+// OCSectionIndices struct to be passed into `checkAddress` to do binary search
+// on.
+//
+
+typedef struct {
+    W_ start;
+    W_ end;
+    ObjectCode *oc;
+} OCSectionIndex;
+
+typedef struct {
+    int n_sections;
+    OCSectionIndex *indices;
+} OCSectionIndices;
+
+static OCSectionIndices *createOCSectionIndices(int n_sections)
+{
+    OCSectionIndices *s_indices;
+    s_indices = stgMallocBytes(sizeof(OCSectionIndices), "OCSectionIndices");
+    s_indices->n_sections = n_sections;
+    s_indices->indices = stgMallocBytes(n_sections*sizeof(OCSectionIndex),
+        "OCSectionIndices::indices");
+    return s_indices;
+}
+
+static int cmpSectionIndex(const void* indexa, const void *indexb)
+{
+    W_ s1 = ((OCSectionIndex*)indexa)->start;
+    W_ s2 = ((OCSectionIndex*)indexb)->start;
+    if (s1 < s2) {
+        return -1;
+    } else if (s1 > s2) {
+        return 1;
+    }
+    return 0;
+}
+
+static OCSectionIndices* buildOCSectionIndices(ObjectCode *ocs)
+{
+    int cnt_sections = 0;
+    ObjectCode *oc;
+    for (oc = ocs; oc; oc = oc->next) {
+        cnt_sections += oc->n_sections;
+    }
+    OCSectionIndices* s_indices = createOCSectionIndices(cnt_sections);
+    int s_i = 0, i;
+    for (oc = ocs; oc; oc = oc->next) {
+        for (i = 0; i < oc->n_sections; i++) {
+            if (oc->sections[i].kind != SECTIONKIND_OTHER) {
+                s_indices->indices[s_i].start = (W_)oc->sections[i].start;
+                s_indices->indices[s_i].end = (W_)oc->sections[i].start
+                    + oc->sections[i].size;
+                s_indices->indices[s_i].oc = oc;
+                s_i++;
+            }
+        }
+    }
+    s_indices->n_sections = s_i;
+    qsort(s_indices->indices,
+        s_indices->n_sections,
+        sizeof(OCSectionIndex),
+        cmpSectionIndex);
+    return s_indices;
+}
+
+static void freeOCSectionIndices(OCSectionIndices *section_indices)
+{
+    free(section_indices->indices);
+    free(section_indices);
+}
+
+static ObjectCode *findOC(OCSectionIndices *s_indices, const void *addr) {
+    W_ w_addr = (W_)addr;
+    if (s_indices->n_sections <= 0) return NULL;
+    if (w_addr < s_indices->indices[0].start) return NULL;
+
+    int left = 0, right = s_indices->n_sections;
+    while (left + 1 < right) {
+        int mid = (left + right)/2;
+        W_ w_mid = s_indices->indices[mid].start;
+        if (w_mid <= w_addr) {
+            left = mid;
+        } else {
+            right = mid;
+        }
+    }
+    ASSERT(w_addr >= s_indices->indices[left].start);
+    if (w_addr < s_indices->indices[left].end) {
+        return s_indices->indices[left].oc;
+    }
+    return NULL;
+}
+
+static void checkAddress (HashTable *addrs, const void *addr,
+        OCSectionIndices *s_indices)
 {
     ObjectCode *oc;
-    int i;
 
     if (!lookupHashTable(addrs, (W_)addr)) {
         insertHashTable(addrs, (W_)addr, addr);
 
-        for (oc = unloaded_objects; oc; oc = oc->next) {
-            for (i = 0; i < oc->n_sections; i++) {
-                if (oc->sections[i].kind != SECTIONKIND_OTHER) {
-                    if ((W_)addr >= (W_)oc->sections[i].start &&
-                        (W_)addr <  (W_)oc->sections[i].start
-                                    + oc->sections[i].size) {
-                        oc->referenced = 1;
-                        return;
-                    }
-                }
-            }
+        oc = findOC(s_indices, addr);
+        if (oc != NULL) {
+            oc->referenced = 1;
+            return;
         }
     }
 }
 
-static void searchStackChunk (HashTable *addrs, StgPtr sp, StgPtr stack_end)
+static void searchStackChunk (HashTable *addrs, StgPtr sp, StgPtr stack_end,
+        OCSectionIndices *s_indices)
 {
     StgPtr p;
     const StgRetInfoTable *info;
@@ -73,7 +172,7 @@ static void searchStackChunk (HashTable *addrs, StgPtr sp, StgPtr stack_end)
         switch (info->i.type) {
         case RET_SMALL:
         case RET_BIG:
-            checkAddress(addrs, (const void*)info);
+            checkAddress(addrs, (const void*)info, s_indices);
             break;
 
         default:
@@ -85,7 +184,8 @@ static void searchStackChunk (HashTable *addrs, StgPtr sp, StgPtr stack_end)
 }
 
 
-static void searchHeapBlocks (HashTable *addrs, bdescr *bd)
+static void searchHeapBlocks (HashTable *addrs, bdescr *bd,
+        OCSectionIndices *s_indices)
 {
     StgPtr p;
     const StgInfoTable *info;
@@ -189,7 +289,7 @@ static void searchHeapBlocks (HashTable *addrs, bdescr *bd)
                 prim = true;
                 size = ap_stack_sizeW(ap);
                 searchStackChunk(addrs, (StgPtr)ap->payload,
-                                 (StgPtr)ap->payload + ap->size);
+                                 (StgPtr)ap->payload + ap->size, s_indices);
                 break;
             }
 
@@ -223,7 +323,7 @@ static void searchHeapBlocks (HashTable *addrs, bdescr *bd)
                 StgStack *stack = (StgStack*)p;
                 prim = true;
                 searchStackChunk(addrs, stack->sp,
-                                 stack->stack + stack->stack_size);
+                                 stack->stack + stack->stack_size, s_indices);
                 size = stack_sizeW(stack);
                 break;
             }
@@ -238,7 +338,7 @@ static void searchHeapBlocks (HashTable *addrs, bdescr *bd)
             }
 
             if (!prim) {
-                checkAddress(addrs,info);
+                checkAddress(addrs,info, s_indices);
             }
 
             p += size;
@@ -251,15 +351,16 @@ static void searchHeapBlocks (HashTable *addrs, bdescr *bd)
 // Do not unload the object if the CCS tree refers to a CCS or CC which
 // originates in the object.
 //
-static void searchCostCentres (HashTable *addrs, CostCentreStack *ccs)
+static void searchCostCentres (HashTable *addrs, CostCentreStack *ccs,
+        OCSectionIndices* s_indices)
 {
     IndexTable *i;
 
-    checkAddress(addrs, ccs);
-    checkAddress(addrs, ccs->cc);
+    checkAddress(addrs, ccs, s_indices);
+    checkAddress(addrs, ccs->cc, s_indices);
     for (i = ccs->indexTable; i != NULL; i = i->next) {
         if (!i->back_edge) {
-            searchCostCentres(addrs, i->ccs);
+            searchCostCentres(addrs, i->ccs, s_indices);
         }
     }
 }
@@ -288,6 +389,7 @@ void checkUnload (StgClosure *static_objects)
 
   ACQUIRE_LOCK(&linker_unloaded_mutex);
 
+  OCSectionIndices *s_indices = buildOCSectionIndices(unloaded_objects);
   // Mark every unloadable object as unreferenced initially
   for (oc = unloaded_objects; oc; oc = oc->next) {
       IF_DEBUG(linker, debugBelch("Checking whether to unload %" PATH_FMT "\n",
@@ -299,7 +401,7 @@ void checkUnload (StgClosure *static_objects)
 
   for (p = static_objects; p != END_OF_STATIC_OBJECT_LIST; p = link) {
       p = UNTAG_STATIC_LIST_PTR(p);
-      checkAddress(addrs, p);
+      checkAddress(addrs, p, s_indices);
       info = get_itbl(p);
       link = *STATIC_LINK(info, p);
   }
@@ -309,32 +411,33 @@ void checkUnload (StgClosure *static_objects)
        p != END_OF_CAF_LIST;
        p = ((StgIndStatic *)p)->static_link) {
       p = UNTAG_STATIC_LIST_PTR(p);
-      checkAddress(addrs, p);
+      checkAddress(addrs, p, s_indices);
   }
 
   for (g = 0; g < RtsFlags.GcFlags.generations; g++) {
-      searchHeapBlocks (addrs, generations[g].blocks);
-      searchHeapBlocks (addrs, generations[g].large_objects);
+      searchHeapBlocks (addrs, generations[g].blocks, s_indices);
+      searchHeapBlocks (addrs, generations[g].large_objects, s_indices);
 
       for (n = 0; n < n_capabilities; n++) {
           ws = &gc_threads[n]->gens[g];
-          searchHeapBlocks(addrs, ws->todo_bd);
-          searchHeapBlocks(addrs, ws->part_list);
-          searchHeapBlocks(addrs, ws->scavd_list);
+          searchHeapBlocks(addrs, ws->todo_bd, s_indices);
+          searchHeapBlocks(addrs, ws->part_list, s_indices);
+          searchHeapBlocks(addrs, ws->scavd_list, s_indices);
       }
   }
 
 #if defined(PROFILING)
   /* Traverse the cost centre tree, calling checkAddress on each CCS/CC */
-  searchCostCentres(addrs, CCS_MAIN);
+  searchCostCentres(addrs, CCS_MAIN, s_indices);
 
   /* Also check each cost centre in the CC_LIST */
   CostCentre *cc;
   for (cc = CC_LIST; cc != NULL; cc = cc->link) {
-      checkAddress(addrs, cc);
+      checkAddress(addrs, cc, s_indices);
   }
 #endif /* PROFILING */
 
+  freeOCSectionIndices(s_indices);
   // Look through the unloadable objects, and any object that is still
   // marked as unreferenced can be physically unloaded, because we
   // have no references to it.