1 files changed, 47 insertions, 36 deletions

diff --git a/erts/emulator/beam/beam_bif_load.c b/erts/emulator/beam/beam_bif_load.c
index b10250dc49..40d44dda4c 100644
--- a/erts/emulator/beam/beam_bif_load.c
+++ b/erts/emulator/beam/beam_bif_load.c
@@ -38,7 +38,7 @@
 #include "erl_thr_progress.h"
 
 static void set_default_trace_pattern(Eterm module);
-static Eterm check_process_code(Process* rp, Module* modp, Uint flags, int *redsp);
+static Eterm check_process_code(Process* rp, Module* modp, Uint flags, int *redsp, int fcalls);
 static void delete_code(Module* modp);
 static void decrement_refc(BeamCodeHeader*);
 static int any_heap_ref_ptrs(Eterm* start, Eterm* end, char* mod_start, Uint mod_size);
@@ -467,7 +467,7 @@ check_old_code_1(BIF_ALIST_1)
 }
 
 Eterm
-erts_check_process_code(Process *c_p, Eterm module, Uint flags, int *redsp)
+erts_check_process_code(Process *c_p, Eterm module, Uint flags, int *redsp, int fcalls)
 {
     Module* modp;
     Eterm res;
@@ -483,7 +483,7 @@ erts_check_process_code(Process *c_p, Eterm module, Uint flags, int *redsp)
 	return am_false;
     erts_rlock_old_code(code_ix);
     res = (!modp->old.code_hdr ? am_false :
-           check_process_code(c_p, modp, flags, redsp));
+           check_process_code(c_p, modp, flags, redsp, fcalls));
     erts_runlock_old_code(code_ix);
 
     return res;
@@ -506,7 +506,7 @@ BIF_RETTYPE erts_internal_check_process_code_2(BIF_ALIST_2)
         goto badarg;
     }
 
-    res = erts_check_process_code(BIF_P, BIF_ARG_1, flags, &reds);
+    res = erts_check_process_code(BIF_P, BIF_ARG_1, flags, &reds, BIF_P->fcalls);
 
     ASSERT(is_value(res));
 
@@ -625,8 +625,8 @@ BIF_RETTYPE call_on_load_function_1(BIF_ALIST_1)
 {
     Module* modp = erts_get_module(BIF_ARG_1, erts_active_code_ix());
 
-    if (modp && modp->curr.code_hdr) {
-	BIF_TRAP_CODE_PTR_0(BIF_P, modp->curr.code_hdr->on_load_function_ptr);
+    if (modp && modp->old.code_hdr) {
+	BIF_TRAP_CODE_PTR_0(BIF_P, modp->old.code_hdr->on_load_function_ptr);
     }
     else {
 	BIF_ERROR(BIF_P, BADARG);
@@ -651,14 +651,14 @@ BIF_RETTYPE finish_after_on_load_2(BIF_ALIST_2)
     code_ix = erts_active_code_ix();
     modp = erts_get_module(BIF_ARG_1, code_ix);
 
-    if (!modp || !modp->curr.code_hdr) {
+    if (!modp || !modp->old.code_hdr) {
     error:
 	erts_smp_thr_progress_unblock();
         erts_smp_proc_lock(BIF_P, ERTS_PROC_LOCK_MAIN);
 	erts_release_code_write_permission();
 	BIF_ERROR(BIF_P, BADARG);
     }
-    if (modp->curr.code_hdr->on_load_function_ptr == NULL) {
+    if (modp->old.code_hdr->on_load_function_ptr == NULL) {
 	goto error;
     }
     if (BIF_ARG_2 != am_false && BIF_ARG_2 != am_true) {
@@ -667,44 +667,55 @@ BIF_RETTYPE finish_after_on_load_2(BIF_ALIST_2)
 
     if (BIF_ARG_2 == am_true) {
 	int i;
+	struct erl_module_instance t;
+
+	/*
+	 * Swap old and new code.
+	 */
+	t = modp->curr;
+	modp->curr = modp->old;
+	modp->old = t;
 
 	/*
 	 * The on_load function succeded. Fix up export entries.
 	 */
 	for (i = 0; i < export_list_size(code_ix); i++) {
 	    Export *ep = export_list(i,code_ix);
-	    if (ep != NULL &&
-		ep->code[0] == BIF_ARG_1 &&
-		ep->code[4] != 0) {
+	    if (ep == NULL || ep->code[0] != BIF_ARG_1) {
+		continue;
+	    }
+	    if (ep->code[4] != 0) {
 		ep->addressv[code_ix] = (void *) ep->code[4];
 		ep->code[4] = 0;
+	    } else {
+		if (ep->addressv[code_ix] == ep->code+3 &&
+		    ep->code[3] == (BeamInstr) em_apply_bif) {
+		    continue;
+		}
+		ep->addressv[code_ix] = ep->code+3;
+		ep->code[3] = (BeamInstr) em_call_error_handler;
 	    }
 	}
 	modp->curr.code_hdr->on_load_function_ptr = NULL;
 	set_default_trace_pattern(BIF_ARG_1);
     } else if (BIF_ARG_2 == am_false) {
-	BeamInstr* code;
-	BeamInstr* end;
+	int i;
 
 	/*
-	 * The on_load function failed. Remove the loaded code.
-	 * This is an combination of delete and purge. We purge
-	 * the current code; the old code is not touched.
+	 * The on_load function failed. Remove references to the
+	 * code that is about to be purged from the export entries.
 	 */
-	erts_total_code_size -= modp->curr.code_length;
-	code = (BeamInstr*) modp->curr.code_hdr;
-	end = (BeamInstr *) ((char *)code + modp->curr.code_length);
-	erts_cleanup_funs_on_purge(code, end);
-	beam_catches_delmod(modp->curr.catches, code, modp->curr.code_length,
-			    erts_active_code_ix());
-        if (modp->curr.code_hdr->literals_start) {
-            erts_free(ERTS_ALC_T_LITERAL, modp->curr.code_hdr->literals_start);
-        }
-	erts_free(ERTS_ALC_T_CODE, modp->curr.code_hdr);
-	modp->curr.code_hdr = NULL;
-	modp->curr.code_length = 0;
-	modp->curr.catches = BEAM_CATCHES_NIL;
-	erts_remove_from_ranges(code);
+
+	for (i = 0; i < export_list_size(code_ix); i++) {
+	    Export *ep = export_list(i,code_ix);
+	    if (ep == NULL || ep->code[0] != BIF_ARG_1) {
+		continue;
+	    }
+	    if (ep->code[3] == (BeamInstr) em_apply_bif) {
+		continue;
+	    }
+	    ep->code[4] = 0;
+	}
     }
     erts_smp_thr_progress_unblock();
     erts_smp_proc_lock(BIF_P, ERTS_PROC_LOCK_MAIN);
@@ -753,7 +764,7 @@ check_mod_funs(Process *p, ErlOffHeap *off_heap, char *area, size_t area_size)
 
 
 static Eterm
-check_process_code(Process* rp, Module* modp, Uint flags, int *redsp)
+check_process_code(Process* rp, Module* modp, Uint flags, int *redsp, int fcalls)
 {
     BeamInstr* start;
     char* literals;
@@ -832,7 +843,7 @@ check_process_code(Process* rp, Module* modp, Uint flags, int *redsp)
 
     /*
      * Message queue can contains funs, but (at least currently) no
-     * constants. If we got references to this module from the message
+     * literals. If we got references to this module from the message
      * queue, a GC cannot remove these...
      */
 
@@ -853,7 +864,7 @@ check_process_code(Process* rp, Module* modp, Uint flags, int *redsp)
 	for (; hfrag; hfrag = hfrag->next) {
 	    if (check_mod_funs(rp, &hfrag->off_heap, mod_start, mod_size))
 		return am_true;
-	    /* Should not contain any constants... */
+	    /* Should not contain any literals... */
 	    ASSERT(!any_heap_refs(&hfrag->mem[0],
 				  &hfrag->mem[hfrag->used_size],
                                   literals,
@@ -908,7 +919,7 @@ check_process_code(Process* rp, Module* modp, Uint flags, int *redsp)
 #ifdef DEBUG
 	/*
 	 * Message buffer fragments should not have any references
-	 * to constants, and off heap lists should already have
+	 * to literals, and off heap lists should already have
 	 * been moved into process off heap structure.
 	 */
 	for (msgp = rp->msg_frag; msgp; msgp = msgp->next) {
@@ -945,7 +956,7 @@ check_process_code(Process* rp, Module* modp, Uint flags, int *redsp)
 	need_gc &= ~done_gc;
 
 	/*
-	 * Try to get rid of constants by by garbage collecting.
+	 * Try to get rid of literals by by garbage collecting.
 	 * Clear both fvalue and ftrace.
 	 */
 
@@ -955,7 +966,7 @@ check_process_code(Process* rp, Module* modp, Uint flags, int *redsp)
 
 	if (need_gc & ERTS_ORDINARY_GC__) {
 	    FLAGS(rp) |= F_NEED_FULLSWEEP;
-	    *redsp += erts_garbage_collect_nobump(rp, 0, rp->arg_reg, rp->arity);
+	    *redsp += erts_garbage_collect_nobump(rp, 0, rp->arg_reg, rp->arity, fcalls);
 	    done_gc |= ERTS_ORDINARY_GC__;
 	}
 	if (need_gc & ERTS_LITERAL_GC__) {