4 files changed, 46 insertions, 10 deletions

diff --git a/erts/emulator/beam/beam_bif_load.c b/erts/emulator/beam/beam_bif_load.c
index 8af7703f51..ea1323d651 100644
--- a/erts/emulator/beam/beam_bif_load.c
+++ b/erts/emulator/beam/beam_bif_load.c
@@ -1086,6 +1086,11 @@ check_process_code(Process* rp, Module* modp, Uint flags, int *redsp, int fcalls
 	|| ErtsInArea(rp->cp, mod_start, mod_size)) {
 	return am_true;
     }
+ 
+    *redsp += 1;
+
+    if (erts_check_nif_export_in_area(rp, mod_start, mod_size))
+	return am_true;
 
     *redsp += (STACK_START(rp) - rp->stop) / 32;
 
@@ -1161,6 +1166,12 @@ check_process_code(Process* rp, Module* modp, Uint flags, int *redsp, int fcalls
 	|| ErtsInArea(rp->cp, mod_start, mod_size)) {
 	return am_true;
     }
+ 
+    *redsp += 1;
+
+    if (erts_check_nif_export_in_area(rp, mod_start, mod_size))
+	return am_true;
+
 
     /*
      * Check all continuation pointers stored on the stack.
diff --git a/erts/emulator/beam/erl_gc.c b/erts/emulator/beam/erl_gc.c
index 6f641a1ea7..af799d09da 100644
--- a/erts/emulator/beam/erl_gc.c
+++ b/erts/emulator/beam/erl_gc.c
@@ -2325,6 +2325,11 @@ move_msgq_to_heap(Process *p)
 static Uint
 setup_rootset(Process *p, Eterm *objv, int nobj, Rootset *rootset)
 {
+    /*
+     * NOTE!
+     *   Remember to update offset_rootset() when changing
+     *   this function.
+     */
     Roots* roots;
     Uint n;
 
@@ -2969,6 +2974,12 @@ offset_one_rootset(Process *p, Sint offs, char* area, Uint area_size,
 	offset_heap_ptr(objv, nobj, offs, area, area_size);
     }
     offset_off_heap(p, offs, area, area_size);
+    if (ERTS_PROC_GET_NIF_TRAP_EXPORT(p)) {
+	Eterm* argv;
+	int argc;
+	if (erts_setup_nif_gc(p, &argv, &argc))
+	    offset_heap_ptr(argv, argc, offs, area, area_size);
+    }
 }
 
 static void
diff --git a/erts/emulator/beam/erl_nif.c b/erts/emulator/beam/erl_nif.c
index 3fe8bdb491..6b265a8b80 100644
--- a/erts/emulator/beam/erl_nif.c
+++ b/erts/emulator/beam/erl_nif.c
@@ -2309,6 +2309,17 @@ erts_setup_nif_gc(Process* proc, Eterm** objv, int* nobj)
     return gc;
 }
 
+int
+erts_check_nif_export_in_area(Process *p, char *start, Uint size)
+{
+    NifExport *nep = ERTS_PROC_GET_NIF_TRAP_EXPORT(p);
+    if (!nep || !nep->saved_current)
+	return 0;
+    if (ErtsInArea(nep->saved_current, start, size))
+	return 1;
+    return 0;
+}
+
 /*
  * Allocate a NifExport and set it in proc specific data
  */
@@ -2360,6 +2371,7 @@ init_nif_sched_data(ErlNifEnv* env, NativeFunPtr direct_fp, NativeFunPtr indirec
     Eterm* reg;
     NifExport* ep;
     int i, scheduler;
+    int orig_argc;
 
     execution_state(env, &proc, &scheduler);
 
@@ -2370,11 +2382,14 @@ init_nif_sched_data(ErlNifEnv* env, NativeFunPtr direct_fp, NativeFunPtr indirec
 
     reg = erts_proc_sched_data(proc)->x_reg_array;
 
+    ASSERT(!need_save || proc->current);
+    orig_argc = need_save ? (int) proc->current[2] : 0;
+
     ep = (NifExport*) ERTS_PROC_GET_NIF_TRAP_EXPORT(proc);
     if (!ep)
-	ep = allocate_nif_sched_data(proc, argc);
-    else if (need_save && ep->rootset_extra < argc) {
-	NifExport* new_ep = allocate_nif_sched_data(proc, argc);
+	ep = allocate_nif_sched_data(proc, orig_argc);
+    else if (need_save && ep->rootset_extra < orig_argc) {
+	NifExport* new_ep = allocate_nif_sched_data(proc, orig_argc);
 	destroy_nif_export(ep);
 	ep = new_ep;
     }
@@ -2387,16 +2402,14 @@ init_nif_sched_data(ErlNifEnv* env, NativeFunPtr direct_fp, NativeFunPtr indirec
     }
     if (scheduler > 0)
 	ERTS_VBUMP_ALL_REDS(proc);
-    for (i = 0; i < argc; i++) {
-	if (need_save)
-	    ep->rootset[i+1] = reg[i];
-	reg[i] = (Eterm) argv[i];
-    }
     if (need_save) {
-	ASSERT(proc->current);
 	ep->saved_current = proc->current;
-	ep->saved_argc = argc;
+	ep->saved_argc = orig_argc;
+	for (i = 0; i < orig_argc; i++)
+	    ep->rootset[i+1] = reg[i];
     }
+    for (i = 0; i < argc; i++)
+	reg[i] = (Eterm) argv[i];
     proc->i = (BeamInstr*) ep->exp.addressv[0];
     ep->exp.code[0] = (BeamInstr) proc->current[0];
     ep->exp.code[1] = (BeamInstr) proc->current[1];
diff --git a/erts/emulator/beam/erl_process.h b/erts/emulator/beam/erl_process.h
index 9e5121abd5..d7eff3605a 100644
--- a/erts/emulator/beam/erl_process.h
+++ b/erts/emulator/beam/erl_process.h
@@ -1586,6 +1586,7 @@ Uint64 erts_step_proc_interval(void);
 
 int erts_setup_nif_gc(Process* proc, Eterm** objv, int* nobj); /* see erl_nif.c */
 void erts_destroy_nif_export(void *); /* see erl_nif.c */
+int erts_check_nif_export_in_area(Process *p, char *start, Uint size);
 
 ErtsProcList *erts_proclist_create(Process *);
 ErtsProcList *erts_proclist_copy(ErtsProcList *);