More fixes for #5654

* In stg_ap_0_fast, if we're evaluating a thunk, the thunk might
  evaluate to a function in which case we may have to adjust its CCS.

* The interpreter has its own implementation of stg_ap_0_fast, so we
  have to do the same shenanigans with creating empty PAPs and copying
  PAPs there.

* GHCi creates Cost Centres as children of CCS_MAIN, which enterFunCCS()
  wrongly assumed to imply that they were CAFs.  Now we use the is_caf
  flag for this, which we have to correctly initialise when we create a
  Cost Centre in GHCi.

5 files changed, 113 insertions, 7 deletions

diff --git a/rts/Apply.cmm b/rts/Apply.cmm
index b18c347d40..b3a04ca58c 100644
--- a/rts/Apply.cmm
+++ b/rts/Apply.cmm
@@ -52,6 +52,10 @@ stg_ap_0_fast ( P_ fun )
   The mechanism we use to wrap the function is to create a
   zero-argument PAP as a proxy object to hold the new CCS, and return
   that.
+
+  If the closure we evaluated is itself a PAP, we cannot make a nested
+  PAP, so we copy the original PAP and set the CCS in the new PAP to
+  enterFunCCS(pap->header.prof.ccs).
 */
 
 again:
@@ -122,6 +126,8 @@ again:
                 CCS_ALLOC(BYTES_TO_WDS(SIZEOF_StgPAP), CCS_OVERHEAD);
                 P_ pap;
                 pap = Hp - size + WDS(1);
+                // We'll lose the original PAP, so we should enter its CCS
+                ccall enterFunCCS(BaseReg "ptr", StgHeader_ccs(fun) "ptr");
                 SET_HDR(pap, stg_PAP_info, CCCS);
                 StgPAP_arity(pap) = StgPAP_arity(fun);
                 StgPAP_n_args(pap) = StgPAP_n_args(fun);
@@ -137,6 +143,27 @@ again:
                 goto loop;
             }
         }
+        case AP,
+             AP_STACK,
+             BLACKHOLE,
+             WHITEHOLE,
+             THUNK,
+             THUNK_1_0,
+             THUNK_0_1,
+             THUNK_2_0,
+             THUNK_1_1,
+             THUNK_0_2,
+             THUNK_STATIC,
+             THUNK_SELECTOR:
+        {
+            // The thunk might evaluate to a function, so we have to come
+            // back here again to adjust its CCS if necessary.  The
+            // stg_restore_ccs_eval stack frame does that.
+            STK_CHK_GEN();
+            jump %ENTRY_CODE(info)
+                (stg_restore_cccs_eval_info, CCCS)
+                (UNTAG(fun));
+        }
         default:
         {
             jump %ENTRY_CODE(info) (UNTAG(fun));
diff --git a/rts/Interpreter.c b/rts/Interpreter.c
index 7221ff64f9..5a395670b7 100644
--- a/rts/Interpreter.c
+++ b/rts/Interpreter.c
@@ -214,6 +214,48 @@ void interp_shutdown ( void )
 
 #endif
 
+#ifdef PROFILING
+
+//
+// Build a zero-argument PAP with the current CCS
+// See Note [Evaluating functions with profiling] in Apply.cmm
+//
+STATIC_INLINE
+StgClosure * newEmptyPAP (Capability *cap,
+                          StgClosure *tagged_obj, // a FUN or a BCO
+                          uint32_t arity)
+{
+    StgPAP *pap = (StgPAP *)allocate(cap, sizeofW(StgPAP));
+    SET_HDR(pap, &stg_PAP_info, cap->r.rCCCS);
+    pap->arity = arity;
+    pap->n_args = 0;
+    pap->fun = tagged_obj;
+    return (StgClosure *)pap;
+}
+
+//
+// Make an exact copy of a PAP, except that we combine the current CCS with the
+// CCS in the PAP.  See Note [Evaluating functions with profiling] in Apply.cmm
+//
+STATIC_INLINE
+StgClosure * copyPAP  (Capability *cap, StgPAP *oldpap)
+{
+    uint32_t size = PAP_sizeW(oldpap->n_args);
+    StgPAP *pap = (StgPAP *)allocate(cap, size);
+    enterFunCCS(&cap->r, oldpap->header.prof.ccs);
+    SET_HDR(pap, &stg_PAP_info, cap->r.rCCCS);
+    pap->arity = oldpap->arity;
+    pap->n_args = oldpap->n_args;
+    pap->fun = oldpap->fun;
+    uint32_t i;
+    for (i = 0; i < ((StgPAP *)pap)->n_args; i++) {
+        pap->payload[i] = oldpap->payload[i];
+    }
+    return (StgClosure *)pap;
+}
+
+#endif
+
 static StgWord app_ptrs_itbl[] = {
     (W_)&stg_ap_p_info,
     (W_)&stg_ap_pp_info,
@@ -343,6 +385,8 @@ eval_obj:
     case CONSTR_1_1:
     case CONSTR_0_2:
     case CONSTR_NOCAF:
+        break;
+
     case FUN:
     case FUN_1_0:
     case FUN_0_1:
@@ -350,15 +394,30 @@ eval_obj:
     case FUN_1_1:
     case FUN_0_2:
     case FUN_STATIC:
+#ifdef PROFILING
+        if (cap->r.rCCCS != obj->header.prof.ccs) {
+            tagged_obj =
+                newEmptyPAP(cap, tagged_obj, get_fun_itbl(obj)->f.arity);
+        }
+#endif
+        break;
+
     case PAP:
-        // already in WHNF
+#ifdef PROFILING
+        if (cap->r.rCCCS != obj->header.prof.ccs) {
+            tagged_obj = copyPAP(cap, (StgPAP *)obj);
+        }
+#endif
         break;
 
     case BCO:
-    {
         ASSERT(((StgBCO *)obj)->arity > 0);
+#ifdef PROFILING
+        if (cap->r.rCCCS != obj->header.prof.ccs) {
+            tagged_obj = newEmptyPAP(cap, tagged_obj, ((StgBCO *)obj)->arity);
+        }
+#endif
         break;
-    }
 
     case AP:    /* Copied from stg_AP_entry. */
     {
@@ -380,7 +439,7 @@ eval_obj:
         // restore the CCCS after evaluating the AP
         Sp -= 2;
         Sp[1] = (W_)cap->r.rCCCS;
-        Sp[0] = (W_)&stg_restore_cccs_info;
+        Sp[0] = (W_)&stg_restore_cccs_eval_info;
 #endif
 
         Sp -= sizeofW(StgUpdateFrame);
@@ -425,7 +484,7 @@ eval_obj:
         // restore the CCCS after evaluating the closure
         Sp -= 2;
         Sp[1] = (W_)cap->r.rCCCS;
-        Sp[0] = (W_)&stg_restore_cccs_info;
+        Sp[0] = (W_)&stg_restore_cccs_eval_info;
 #endif
         Sp -= 2;
         Sp[1] = (W_)tagged_obj;
@@ -465,7 +524,8 @@ do_return:
         // NOTE: not using get_itbl().
         info = ((StgClosure *)Sp)->header.info;
 
-        if (info == (StgInfoTable *)&stg_restore_cccs_info) {
+        if (info == (StgInfoTable *)&stg_restore_cccs_info ||
+            info == (StgInfoTable *)&stg_restore_cccs_eval_info) {
             cap->r.rCCCS = (CostCentreStack*)Sp[1];
             Sp += 2;
             goto do_return;
diff --git a/rts/Printer.c b/rts/Printer.c
index f23e0b0636..87b11e80d8 100644
--- a/rts/Printer.c
+++ b/rts/Printer.c
@@ -526,6 +526,11 @@ printStackChunk( StgPtr sp, StgPtr spBottom )
                 fprintCCS(stderr, (CostCentreStack*)sp[1]);
                 debugBelch("\n" );
                 continue;
+            } else if (c == (StgWord)&stg_restore_cccs_eval_info) {
+                debugBelch("stg_restore_cccs_eval_info\n" );
+                fprintCCS(stderr, (CostCentreStack*)sp[1]);
+                debugBelch("\n" );
+                continue;
 #endif
             } else {
                 debugBelch("RET_SMALL (%p)\n", info);
diff --git a/rts/Profiling.c b/rts/Profiling.c
index 952785be18..94ec55582b 100644
--- a/rts/Profiling.c
+++ b/rts/Profiling.c
@@ -236,6 +236,10 @@ CostCentre *mkCostCentre (char *label, char *module, char *srcloc)
     cc->label = label;
     cc->module = module;
     cc->srcloc = srcloc;
+    cc->is_caf = 0;
+    cc->mem_alloc = 0;
+    cc->time_ticks = 0;
+    cc->link = NULL;
     return cc;
 }
 
@@ -379,7 +383,7 @@ void enterFunCCS (StgRegTable *reg, CostCentreStack *ccsfn)
     }
 
     // common case 2: the function stack is empty, or just CAF
-    if (ccsfn->prevStack == CCS_MAIN) {
+    if (ccsfn->cc->is_caf) {
         return;
     }
 
diff --git a/rts/StgMiscClosures.cmm b/rts/StgMiscClosures.cmm
index aa22c99be4..e8a5b8fed5 100644
--- a/rts/StgMiscClosures.cmm
+++ b/rts/StgMiscClosures.cmm
@@ -55,6 +55,16 @@ INFO_TABLE_RET (stg_restore_cccs, RET_SMALL, W_ info_ptr, W_ cccs)
     jump %ENTRY_CODE(Sp(0)) [*]; // NB. all registers live!
 }
 
+
+INFO_TABLE_RET (stg_restore_cccs_eval, RET_SMALL, W_ info_ptr, W_ cccs)
+    return (P_ ret)
+{
+#if defined(PROFILING)
+    CCCS = cccs;
+#endif
+    jump stg_ap_0_fast(ret);
+}
+
 /* ----------------------------------------------------------------------------
    Support for the bytecode interpreter.
    ------------------------------------------------------------------------- */