Add a loader state for HiPE code loading

Just like the BEAM loader state (as returned by
erlang:prepare_loading/2), the HiPE loader state is contained in a magic
binary.

Eventually, we will separate HiPE loading into a prepare and a finalise
phase, like the BEAM loader, where the prepare phase will be implemented
by hipe_unified_loader and the finalise phase be implemented in C by
hipe_load.c and beam_load.c, making prepare side-effect free and
finalise atomic. The finalise phase will be exposed through the
erlang:finish_loading/1 API, just like the BEAM loader, as this will
allow HiPE and BEAM modules to be mixed in the same atomic "commit".

The usage of a loader state makes it easier to keep track of all
resources allocated during loading, and will not only make it easy to
prevent leaks when hipe_unified_loader crashes, but also paves the way
for proper, leak-free, unloading of HiPE modules.

14 files changed, 387 insertions, 75 deletions

diff --git a/bootstrap/lib/kernel/ebin/hipe_unified_loader.beam b/bootstrap/lib/kernel/ebin/hipe_unified_loader.beam
index f560389756..9877283cd2 100644
--- a/bootstrap/lib/kernel/ebin/hipe_unified_loader.beam
+++ b/bootstrap/lib/kernel/ebin/hipe_unified_loader.beam
diff --git a/erts/emulator/Makefile.in b/erts/emulator/Makefile.in
index 21bcbbab27..fa4de5805e 100644
--- a/erts/emulator/Makefile.in
+++ b/erts/emulator/Makefile.in
@@ -884,7 +884,9 @@ HIPE_OBJS= \
 	$(OBJDIR)/hipe_bif2.o \
 	$(OBJDIR)/hipe_debug.o \
 	$(OBJDIR)/hipe_gc.o \
+	$(OBJDIR)/hipe_load.o \
 	$(OBJDIR)/hipe_mode_switch.o \
+	$(OBJDIR)/hipe_module.o \
 	$(OBJDIR)/hipe_native_bif.o \
 	$(OBJDIR)/hipe_stack.o $(HIPE_ARCH_OBJS)
 ifdef HIPE_ENABLED
diff --git a/erts/emulator/beam/beam_bif_load.c b/erts/emulator/beam/beam_bif_load.c
index 237513095a..3302de0787 100644
--- a/erts/emulator/beam/beam_bif_load.c
+++ b/erts/emulator/beam/beam_bif_load.c
@@ -150,7 +150,13 @@ BIF_RETTYPE code_is_module_native_1(BIF_ALIST_1)
 BIF_RETTYPE code_make_stub_module_3(BIF_ALIST_3)
 {
     Module* modp;
-    Eterm res;
+    Eterm res, mod;
+
+    if (!ERTS_TERM_IS_MAGIC_BINARY(BIF_ARG_1) ||
+	is_not_atom(mod = erts_module_for_prepared_code
+		    (((ProcBin*)binary_val(BIF_ARG_1))->val))) {
+	BIF_ERROR(BIF_P, BADARG);
+    }
 
     if (!erts_try_seize_code_write_permission(BIF_P)) {
 	ERTS_BIF_YIELD3(bif_export[BIF_code_make_stub_module_3],
@@ -160,7 +166,7 @@ BIF_RETTYPE code_make_stub_module_3(BIF_ALIST_3)
     erts_smp_proc_unlock(BIF_P, ERTS_PROC_LOCK_MAIN);
     erts_smp_thr_progress_block();
 
-    modp = erts_get_module(BIF_ARG_1, erts_active_code_ix());
+    modp = erts_get_module(mod, erts_active_code_ix());
 
     if (modp && modp->curr.num_breakpoints > 0) {
 	ASSERT(modp->curr.code_hdr != NULL);
@@ -172,12 +178,12 @@ BIF_RETTYPE code_make_stub_module_3(BIF_ALIST_3)
 
     res = erts_make_stub_module(BIF_P, BIF_ARG_1, BIF_ARG_2, BIF_ARG_3);
 
-    if (res == BIF_ARG_1) {
+    if (res == mod) {
 	erts_end_staging_code_ix();
 	erts_commit_staging_code_ix();
 #ifdef HIPE
         if (!modp)
-	    modp = erts_get_module(BIF_ARG_1, erts_active_code_ix());
+	    modp = erts_get_module(mod, erts_active_code_ix());
         hipe_redirect_to_module(modp);
 #endif
     }
diff --git a/erts/emulator/beam/beam_load.c b/erts/emulator/beam/beam_load.c
index 9b206f9a23..7d29f393e5 100644
--- a/erts/emulator/beam/beam_load.c
+++ b/erts/emulator/beam/beam_load.c
@@ -44,6 +44,7 @@
 #include "hipe_bif0.h"
 #include "hipe_mode_switch.h"
 #include "hipe_arch.h"
+#include "hipe_load.h"
 #endif
 
 ErlDrvBinary* erts_gzinflate_buffer(char*, int);
@@ -482,8 +483,8 @@ static void free_literal_fragment(ErlHeapFragment*);
 static void loader_state_dtor(Binary* magic);
 static Eterm stub_insert_new_code(Process *c_p, ErtsProcLocks c_p_locks,
 				  Eterm group_leader, Eterm module,
-				  BeamCodeHeader* code, Uint size,
-				  void* hipe_code_start, UWord hipe_code_size);
+				  BeamCodeHeader* code_hdr, Uint size,
+				  HipeModule *hipe_code);
 static int init_iff_file(LoaderState* stp, byte* code, Uint size);
 static int scan_iff_file(LoaderState* stp, Uint* chunk_types,
 			 Uint num_types, Uint num_mandatory);
@@ -942,6 +943,13 @@ erts_module_for_prepared_code(Binary* magic)
     LoaderState* stp;
 
     if (ERTS_MAGIC_BIN_DESTRUCTOR(magic) != loader_state_dtor) {
+#ifdef HIPE
+	HipeLoaderState *hipe_stp;
+	if ((hipe_stp = hipe_get_loader_state(magic))
+	    && hipe_stp->text_segment != 0) {
+	    return hipe_stp->module;
+	}
+#endif
 	return NIL;
     }
     stp = ERTS_MAGIC_BIN_DATA(magic);
@@ -1094,7 +1102,7 @@ static Eterm
 stub_insert_new_code(Process *c_p, ErtsProcLocks c_p_locks,
 		     Eterm group_leader, Eterm module,
 		     BeamCodeHeader* code_hdr, Uint size,
-		     void* hipe_code_start, UWord hipe_code_size)
+		     HipeModule *hipe_code)
 {
     Module* modp;
     Eterm retval;
@@ -1121,12 +1129,9 @@ stub_insert_new_code(Process *c_p, ErtsProcLocks c_p_locks,
     DBG_TRACE_MFA(make_atom(modp->module), 0, 0, "insert_new_code new_hipe_refs = %p", modp->new_hipe_refs);
     modp->curr.first_hipe_ref = modp->new_hipe_refs;
     modp->curr.first_hipe_sdesc = modp->new_hipe_sdesc;
-    modp->curr.hipe_code_start = hipe_code_start;
     modp->new_hipe_refs = NULL;
     modp->new_hipe_sdesc = NULL;
-# ifdef DEBUG
-    modp->curr.hipe_code_size  = hipe_code_size;
-# endif
+    modp->curr.hipe_code = hipe_code;
 #endif
 
     /*
@@ -6271,10 +6276,13 @@ patch_funentries(Eterm Patchlist)
  */
 
 Eterm
-erts_make_stub_module(Process* p, Eterm Mod, Eterm Beam, Eterm Info)
+erts_make_stub_module(Process* p, Eterm hipe_magic_bin, Eterm Beam, Eterm Info)
 {
     Binary* magic;
+    Binary* hipe_magic;
     LoaderState* stp;
+    HipeLoaderState* hipe_stp;
+    HipeModule *hipe_code;
     BeamInstr Funcs;
     BeamInstr Patchlist;
     Eterm MD5Bin;
@@ -6290,7 +6298,6 @@ erts_make_stub_module(Process* p, Eterm Mod, Eterm Beam, Eterm Info)
     byte* temp_alloc = NULL;
     byte* bytes;
     Uint size;
-    UWord hipe_code_start = NULL, hipe_code_size = 0;
 
     /*
      * Must initialize stp->lambdas here because the error handling code
@@ -6298,15 +6305,19 @@ erts_make_stub_module(Process* p, Eterm Mod, Eterm Beam, Eterm Info)
      */
     magic = erts_alloc_loader_state();
     stp = ERTS_MAGIC_BIN_DATA(magic);
+    hipe_code = erts_alloc(ERTS_ALC_T_HIPE, sizeof(*hipe_code));
 
-    if (is_not_atom(Mod)) {
+    if (!ERTS_TERM_IS_MAGIC_BINARY(hipe_magic_bin) ||
+	!(hipe_magic = ((ProcBin*)binary_val(hipe_magic_bin))->val,
+	  hipe_stp = hipe_get_loader_state(hipe_magic)) ||
+	hipe_stp->module == NIL || hipe_stp->text_segment == 0) {
 	goto error;
     }
     if (is_not_tuple(Info)) {
 	goto error;
     }
     tp = tuple_val(Info);
-    if (tp[0] != make_arityval(5)) {
+    if (tp[0] != make_arityval(3)) {
       goto error;
     }
     Funcs = tp[1];
@@ -6323,20 +6334,11 @@ erts_make_stub_module(Process* p, Eterm Mod, Eterm Beam, Eterm Info)
     }
     size = binary_size(Beam);
 
-#ifdef HIPE
-    if (!term_to_Uint(tp[4], &hipe_code_start))
-        goto error;
-# ifdef DEBUG
-    if (!term_to_Uint(tp[5], &hipe_code_size))
-        goto error;
-# endif
-#endif
-
     /*
      * Scan the Beam binary and read the interesting sections.
      */
 
-    stp->module = Mod;
+    stp->module = hipe_stp->module;
     stp->group_leader = p->group_leader;
     stp->num_functions = n;
     if (!init_iff_file(stp, bytes, size)) {
@@ -6450,7 +6452,8 @@ erts_make_stub_module(Process* p, Eterm Mod, Eterm Beam, Eterm Info)
 #else
 	op = (Eterm) BeamOpCode(op_move_return_n);
 #endif
-	fp = make_stub(fp, Mod, func, arity, (Uint)native_address, op);
+	fp = make_stub(fp, hipe_stp->module, func, arity, (Uint)native_address,
+		       op);
     }
 
     /*
@@ -6490,12 +6493,18 @@ erts_make_stub_module(Process* p, Eterm Mod, Eterm Beam, Eterm Info)
     }
 
     /*
+     * Initialise HiPE module
+     */
+    hipe_code->text_segment = hipe_stp->text_segment;
+    hipe_code->text_segment_size = hipe_stp->text_segment_size;
+    hipe_code->data_segment = hipe_stp->data_segment;
+
+    /*
      * Insert the module in the module table.
      */
 
-    rval = stub_insert_new_code(p, 0, p->group_leader, Mod,
-				code_hdr, code_size,
-				(void*)hipe_code_start, hipe_code_size);
+    rval = stub_insert_new_code(p, 0, p->group_leader, hipe_stp->module,
+				code_hdr, code_size, hipe_code);
     if (rval != NIL) {
 	goto error;
     }
@@ -6511,12 +6520,20 @@ erts_make_stub_module(Process* p, Eterm Mod, Eterm Beam, Eterm Info)
     }
 
     if (patch_funentries(Patchlist)) {
+	Eterm mod = hipe_stp->module;
+	/* Prevent code from being freed */
+	hipe_stp->text_segment = 0;
+	hipe_stp->data_segment = 0;
+
 	erts_free_aligned_binary_bytes(temp_alloc);
 	free_loader_state(magic);
-	return Mod;
+	hipe_free_loader_state(hipe_magic);
+
+	return mod;
     }
 
  error:
+    erts_free(ERTS_ALC_T_HIPE, hipe_code);
     erts_free_aligned_binary_bytes(temp_alloc);
     free_loader_state(magic);
     BIF_ERROR(p, BADARG);
diff --git a/erts/emulator/beam/module.c b/erts/emulator/beam/module.c
index 93ce8fcf9f..5c56ad20df 100644
--- a/erts/emulator/beam/module.c
+++ b/erts/emulator/beam/module.c
@@ -79,10 +79,7 @@ void erts_module_instance_init(struct erl_module_instance* modi)
 #ifdef HIPE
     modi->first_hipe_ref = NULL;
     modi->first_hipe_sdesc = NULL;
-    modi->hipe_code_start = NULL;
-# ifdef DEBUG
-    modi->hipe_code_size = 0;
-# endif
+    modi->hipe_code = NULL;
 #endif
 }
 
diff --git a/erts/emulator/beam/module.h b/erts/emulator/beam/module.h
index 694583597b..c3adaf12fc 100644
--- a/erts/emulator/beam/module.h
+++ b/erts/emulator/beam/module.h
@@ -23,6 +23,10 @@
 
 #include "index.h"
 
+#ifdef HIPE
+#include "hipe_module.h"
+#endif
+
 struct erl_module_instance {
     BeamCodeHeader* code_hdr;
     int code_length;		/* Length of loaded code in bytes. */
@@ -33,10 +37,7 @@ struct erl_module_instance {
 #ifdef HIPE
     struct hipe_ref* first_hipe_ref;  /* all external hipe calls from this module */
     struct hipe_sdesc* first_hipe_sdesc;  /* all stack descriptors for this module */
-    void* hipe_code_start;
-# ifdef DEBUG
-    UWord hipe_code_size;
-# endif
+    HipeModule *hipe_code;
 #endif
 };
 
diff --git a/erts/emulator/hipe/hipe_bif0.c b/erts/emulator/hipe/hipe_bif0.c
index 6f9bfe2f40..daa198d695 100644
--- a/erts/emulator/hipe/hipe_bif0.c
+++ b/erts/emulator/hipe/hipe_bif0.c
@@ -44,6 +44,7 @@
 #include "hipe_mode_switch.h"
 #include "hipe_native_bif.h"
 #include "hipe_bif0.h"
+#include "hipe_load.h"
 /* We need hipe_literals.h for HIPE_SYSTEM_CRC, but it redefines
    a few constants. #undef them here to avoid warnings. */
 #undef F_TIMO
@@ -375,15 +376,31 @@ BIF_RETTYPE hipe_bifs_ref_set_2(BIF_ALIST_2)
 }
 
 /*
+ * BIFs for loading code.
+ */
+
+static HipeLoaderState *get_loader_state(Eterm term)
+{
+    ProcBin *pb;
+
+    if (!ERTS_TERM_IS_MAGIC_BINARY(term)) return NULL;
+
+    pb = (ProcBin*) binary_val(term);
+    return hipe_get_loader_state(pb->val);
+}
+
+
+/*
  * Allocate memory and copy machine code to it.
  */
-BIF_RETTYPE hipe_bifs_enter_code_2(BIF_ALIST_2)
+BIF_RETTYPE hipe_bifs_enter_code_3(BIF_ALIST_3)
 {
     Uint nrbytes;
     void *bytes;
     void *address;
     Eterm trampolines;
     Eterm *hp;
+    HipeLoaderState *stp;
 #ifndef DEBUG
     ERTS_DECLARE_DUMMY(Uint bitoffs);
     ERTS_DECLARE_DUMMY(Uint bitsize);
@@ -392,13 +409,15 @@ BIF_RETTYPE hipe_bifs_enter_code_2(BIF_ALIST_2)
     Uint bitsize;
 #endif
 
-    if (is_not_binary(BIF_ARG_1))
+    if (is_not_binary(BIF_ARG_1) ||
+	(!(stp = get_loader_state(BIF_ARG_3))))
 	BIF_ERROR(BIF_P, BADARG);
     nrbytes = binary_size(BIF_ARG_1);
     ERTS_GET_BINARY_BYTES(BIF_ARG_1, bytes, bitoffs, bitsize);
     ASSERT(bitoffs == 0);
     ASSERT(bitsize == 0);
     trampolines = NIL;
+    // XXX: Trampolines are not tracked and freed
     address = hipe_alloc_code(nrbytes, BIF_ARG_2, &trampolines, BIF_P);
     if (!address) {
 	Uint nrcallees;
@@ -407,11 +426,15 @@ BIF_RETTYPE hipe_bifs_enter_code_2(BIF_ALIST_2)
 	    nrcallees = arityval(tuple_val(BIF_ARG_2)[0]);
 	else
 	    nrcallees = 0;
+	// XXX: Is there any reason to not just BIF_ERROR, so that the runtime
+	// survives?
 	erts_exit(ERTS_ERROR_EXIT, "%s: failed to allocate %lu bytes and %lu trampolines\r\n",
 		 __func__, (unsigned long)nrbytes, (unsigned long)nrcallees);
     }
     memcpy(address, bytes, nrbytes);
     hipe_flush_icache_range(address, nrbytes);
+    stp->text_segment = address;
+    stp->text_segment_size = nrbytes;
     hp = HAlloc(BIF_P, 3);
     hp[0] = make_arityval(2);
     hp[1] = address_to_term(address, BIF_P);
@@ -424,25 +447,34 @@ BIF_RETTYPE hipe_bifs_enter_code_2(BIF_ALIST_2)
 /*
  * Allocate memory for arbitrary non-Erlang data.
  */
-BIF_RETTYPE hipe_bifs_alloc_data_2(BIF_ALIST_2)
+BIF_RETTYPE hipe_bifs_alloc_data_3(BIF_ALIST_3)
 {
-    Uint align, nrbytes;
-    void *block;
+    Uint align;
+    HipeLoaderState *stp;
 
     if (is_not_small(BIF_ARG_1) || is_not_small(BIF_ARG_2) ||
+	(!(stp = get_loader_state(BIF_ARG_3))) ||
 	(align = unsigned_val(BIF_ARG_1), !IS_POWER_OF_TWO(align)))
 	BIF_ERROR(BIF_P, BADARG);
-    nrbytes = unsigned_val(BIF_ARG_2);
-    if (nrbytes == 0)
+
+    if (stp->data_segment_size || stp->data_segment)
+	BIF_ERROR(BIF_P, BADARG);
+
+    stp->data_segment_size = unsigned_val(BIF_ARG_2);
+    if (stp->data_segment_size == 0)
 	BIF_RET(make_small(0));
-    block = erts_alloc(ERTS_ALC_T_HIPE, nrbytes);
-    if ((unsigned long)block & (align-1)) {
-	fprintf(stderr, "%s: erts_alloc(%lu) returned %p which is not %lu-byte aligned\r\n",
-		__FUNCTION__, (unsigned long)nrbytes, block, (unsigned long)align);
-	erts_free(ERTS_ALC_T_HIPE, block);
+    stp->data_segment = erts_alloc(ERTS_ALC_T_HIPE, stp->data_segment_size);
+    if ((unsigned long)stp->data_segment & (align-1)) {
+	fprintf(stderr, "%s: erts_alloc(%lu) returned %p which is not %lu-byte "
+		"aligned\r\n",
+		__FUNCTION__, (unsigned long)stp->data_segment_size,
+		stp->data_segment, (unsigned long)align);
+	erts_free(ERTS_ALC_T_HIPE, stp->data_segment);
+	stp->data_segment = NULL;
+	stp->data_segment_size = 0;
 	BIF_ERROR(BIF_P, EXC_NOTSUP);
     }
-    BIF_RET(address_to_term(block, BIF_P));
+    BIF_RET(address_to_term(stp->data_segment, BIF_P));
 }
 
 /*
@@ -1711,12 +1743,14 @@ void hipe_purge_module(Module* modp)
                 prevp = &p->next_in_mod;
         }
     }
-    if (modp->old.hipe_code_start) {
+    if (modp->old.hipe_code && modp->old.hipe_code->text_segment) {
 #ifdef DEBUG
-        sys_memset(modp->old.hipe_code_start, 0xfe, modp->old.hipe_code_size);
+        sys_memset(modp->old.hipe_code->text_segment, 0xfe,
+                   modp->old.hipe_code->text_segment_size);
 #endif
-        hipe_free_code(modp->old.hipe_code_start);
-        modp->old.hipe_code_start = NULL;
+        hipe_free_code(modp->old.hipe_code->text_segment);
+        modp->old.hipe_code->text_segment = NULL;
+        modp->old.hipe_code->text_segment_size = 0;
     }
 }
 
@@ -1870,3 +1904,23 @@ BIF_RETTYPE hipe_bifs_patch_call_3(BIF_ALIST_3)
 	BIF_ERROR(BIF_P, BADARG);
     BIF_RET(NIL);
 }
+
+BIF_RETTYPE hipe_bifs_alloc_loader_state_1(BIF_ALIST_1)
+{
+    Binary *magic;
+    Eterm *hp;
+    Eterm res;
+
+    if (is_not_atom(BIF_ARG_1))
+	BIF_ERROR(BIF_P, BADARG);
+
+    magic = hipe_alloc_loader_state(BIF_ARG_1);
+
+    if (!magic)
+	BIF_ERROR(BIF_P, BADARG);
+
+    hp = HAlloc(BIF_P, PROC_BIN_SIZE);
+    res = erts_mk_magic_binary_term(&hp, &MSO(BIF_P), magic);
+    erts_refc_dec(&magic->refc, 1);
+    BIF_RET(res);
+}
diff --git a/erts/emulator/hipe/hipe_bif0.tab b/erts/emulator/hipe/hipe_bif0.tab
index ff9c1d06fb..1e4bde5ef2 100644
--- a/erts/emulator/hipe/hipe_bif0.tab
+++ b/erts/emulator/hipe/hipe_bif0.tab
@@ -44,8 +44,8 @@ bif hipe_bifs:ref/1
 bif hipe_bifs:ref_get/1
 bif hipe_bifs:ref_set/2
 
-bif hipe_bifs:enter_code/2
-bif hipe_bifs:alloc_data/2
+bif hipe_bifs:enter_code/3
+bif hipe_bifs:alloc_data/3
 bif hipe_bifs:constants_size/0
 bif hipe_bifs:merge_term/1
 
@@ -83,6 +83,8 @@ bif hipe_bifs:patch_call/3
 bif hipe_bifs:add_ref/2
 bif hipe_bifs:remove_refs_from/1
 
+bif hipe_bifs:alloc_loader_state/1
+
 # atoms used by add_ref/2
 atom call
 atom load_mfa
diff --git a/erts/emulator/hipe/hipe_load.c b/erts/emulator/hipe/hipe_load.c
new file mode 100644
index 0000000000..5bce3f1aee
--- /dev/null
+++ b/erts/emulator/hipe/hipe_load.c
@@ -0,0 +1,101 @@
+/*
+ * %CopyrightBegin%
+
+ *
+ * Copyright Ericsson AB 2016. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * %CopyrightEnd%
+ */
+/*
+ * hipe_load.c
+ *
+ * HiPE atomic code loader
+ */
+
+#ifdef HAVE_CONFIG_H
+#include "config.h"
+#endif
+#include "sys.h"
+#include "global.h"
+#include "erl_binary.h"
+#include "hipe_load.h"
+
+/*
+ * This destructor function can safely be called multiple times.
+ */
+static void
+hipe_loader_state_dtor(Binary* magic)
+{
+    HipeLoaderState* stp = ERTS_MAGIC_BIN_DATA(magic);
+    if (stp->module == NIL) return;
+
+    erts_fprintf(stderr, "Destroying HiPE loader state for module %T\n",
+		 stp->module);
+
+    // TODO: Needs to be freed separately. We'd like have a unified executable
+    // code allocator, so postpone this for now.
+    /* if (stp->text_segment) */
+    /* 	erts_free(ERTS_ALC_T_HIPE, stp->text_segment); */
+    stp->text_segment = NULL;
+    stp->text_segment_size = 0;
+
+    if (stp->data_segment)
+	erts_free(ERTS_ALC_T_HIPE, stp->data_segment);
+    stp->data_segment = NULL;
+    stp->data_segment_size = 0;
+
+    stp->module = NIL;
+}
+
+Binary *hipe_alloc_loader_state(Eterm module)
+{
+    HipeLoaderState *stp;
+    Binary *magic;
+
+    if (is_not_atom(module)) return NULL;
+
+    erts_fprintf(stderr, "Creating HiPE loader state for module %T\n", module);
+
+    magic = erts_create_magic_binary(sizeof(HipeLoaderState),
+				     hipe_loader_state_dtor);
+    erts_refc_inc(&magic->refc, 1);
+    stp = ERTS_MAGIC_BIN_DATA(magic);
+
+    stp->module = module;
+    stp->text_segment = NULL;
+    stp->text_segment_size = 0;
+    stp->data_segment = NULL;
+    stp->data_segment_size = 0;
+
+    return magic;
+}
+
+void hipe_free_loader_state(Binary *magic)
+{
+    if (ERTS_MAGIC_BIN_DESTRUCTOR(magic) != hipe_loader_state_dtor)
+	return;
+
+    /* Why does BEAM do a refc_dec here? What is holding that reference? */
+    hipe_loader_state_dtor(magic);
+}
+
+HipeLoaderState *
+hipe_get_loader_state(Binary *magic)
+{
+    if (ERTS_MAGIC_BIN_DESTRUCTOR(magic) != hipe_loader_state_dtor)
+	return NULL;
+
+    return (HipeLoaderState*) ERTS_MAGIC_BIN_DATA(magic);
+}
diff --git a/erts/emulator/hipe/hipe_load.h b/erts/emulator/hipe/hipe_load.h
new file mode 100644
index 0000000000..9e24e915ed
--- /dev/null
+++ b/erts/emulator/hipe/hipe_load.h
@@ -0,0 +1,45 @@
+/*
+ * %CopyrightBegin%
+ *
+ * Copyright Ericsson AB 2016. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * %CopyrightEnd%
+ */
+/*
+ * hipe_load.h
+ *
+ * HiPE atomic code loader
+ */
+#ifndef HIPE_LOAD_H
+#define HIPE_LOAD_H
+
+#include "global.h"
+
+typedef struct hipe_loader_state {
+    Eterm module;                   /* Module name, atom */
+
+    void *text_segment;
+    Uint text_segment_size;
+
+    void *data_segment;
+    Uint data_segment_size;
+
+} HipeLoaderState;
+
+extern Binary *hipe_alloc_loader_state(Eterm module);
+extern void hipe_free_loader_state(Binary *binary);
+extern HipeLoaderState *hipe_get_loader_state(Binary *binary);
+
+#endif /* HIPE_LOAD_H */
diff --git a/erts/emulator/hipe/hipe_module.c b/erts/emulator/hipe/hipe_module.c
new file mode 100644
index 0000000000..1553bf9087
--- /dev/null
+++ b/erts/emulator/hipe/hipe_module.c
@@ -0,0 +1,37 @@
+/*
+ * %CopyrightBegin%
+ *
+ * Copyright Ericsson AB 2016. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * %CopyrightEnd%
+ */
+
+#ifdef HAVE_CONFIG_H
+#include "config.h"
+#endif
+#include "sys.h"
+#include "hipe_module.h"
+
+void hipe_free_module(HipeModule *mod)
+{
+#ifdef DEBUG
+    sys_memzero(mod->text_segment, mod->text_segment_size);
+#endif
+    /* XXX: erts_free(ERTS_ALC_T_HIPE, mod->text_segment); */
+    if (mod->data_segment) /* Some modules lack data segments */
+	erts_free(ERTS_ALC_T_HIPE, mod->data_segment);
+
+    erts_free(ERTS_ALC_T_HIPE, mod);
+}
diff --git a/erts/emulator/hipe/hipe_module.h b/erts/emulator/hipe/hipe_module.h
new file mode 100644
index 0000000000..1e1302fc0a
--- /dev/null
+++ b/erts/emulator/hipe/hipe_module.h
@@ -0,0 +1,42 @@
+/*
+ * %CopyrightBegin%
+
+ *
+ * Copyright Ericsson AB 2016. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * %CopyrightEnd%
+ */
+/*
+ * hipe_module.h
+ *
+ *
+ */
+#ifndef HIPE_MODULE_H
+#define HIPE_MODULE_H
+
+/* Forward-declare type to resolve circular dependency with module.h */
+typedef struct hipe_module HipeModule;
+
+#include "global.h"
+
+struct hipe_module {
+    void *text_segment;
+    Uint text_segment_size;
+    void *data_segment;
+};
+
+extern void hipe_free_module(HipeModule *mod);
+
+#endif /* HIPE_MODULE_H */
diff --git a/lib/hipe/cerl/erl_bif_types.erl b/lib/hipe/cerl/erl_bif_types.erl
index 3e5d8e7b90..9a08305686 100644
--- a/lib/hipe/cerl/erl_bif_types.erl
+++ b/lib/hipe/cerl/erl_bif_types.erl
@@ -1004,8 +1004,8 @@ type(erlang, tuple_to_list, 1, Xs, Opaques) ->
 %%-- hipe_bifs ----------------------------------------------------------------
 type(hipe_bifs, add_ref, 2, Xs, Opaques) ->
   strict(hipe_bifs, add_ref, 2, Xs, fun (_) -> t_nil() end, Opaques);
-type(hipe_bifs, alloc_data, 2, Xs, Opaques) ->
-  strict(hipe_bifs, alloc_data, 2, Xs,
+type(hipe_bifs, alloc_data, 3, Xs, Opaques) ->
+  strict(hipe_bifs, alloc_data, 3, Xs,
 	 fun (_) -> t_integer() end, Opaques); % address
 type(hipe_bifs, array, 2, Xs, Opaques) ->
   strict(hipe_bifs, array, 2, Xs, fun (_) -> t_immarray() end, Opaques);
@@ -1052,8 +1052,8 @@ type(hipe_bifs, call_count_on, 1, Xs, Opaques) ->
 	 fun (_) -> t_sup(t_atom('true'), t_nil()) end, Opaques);
 type(hipe_bifs, check_crc, 1, Xs, Opaques) ->
   strict(hipe_bifs, check_crc, 1, Xs, fun (_) -> t_boolean() end, Opaques);
-type(hipe_bifs, enter_code, 2, Xs, Opaques) ->
-  strict(hipe_bifs, enter_code, 2, Xs,
+type(hipe_bifs, enter_code, 3, Xs, Opaques) ->
+  strict(hipe_bifs, enter_code, 3, Xs,
 	 fun (_) -> t_tuple([t_integer(),
 			     %% XXX: The tuple below contains integers and
 			     %% is of size same as the length of the MFA list
@@ -1111,6 +1111,8 @@ type(hipe_bifs, write_u32, 2, Xs, Opaques) ->
   strict(hipe_bifs, write_u32, 2, Xs, fun (_) -> t_nil() end, Opaques);
 type(hipe_bifs, write_u64, 2, Xs, Opaques) ->
   strict(hipe_bifs, write_u64, 2, Xs, fun (_) -> t_nil() end, Opaques);
+type(hipe_bifs, alloc_loader_state, 1, Xs, Opaques) ->
+  strict(hipe_bifs, alloc_loader_state, 1, Xs, fun (_) -> t_binary() end, Opaques);
 %%-- lists --------------------------------------------------------------------
 type(lists, all, 2, Xs, Opaques) ->
   strict(lists, all, 2, Xs,
@@ -2458,8 +2460,8 @@ arg_types(hipe_bifs, add_ref, 2) ->
 		     t_integer(),
 		     t_sup(t_atom('call'), t_atom('load_mfa')),
 		     t_trampoline()])];
-arg_types(hipe_bifs, alloc_data, 2) ->
-  [t_integer(), t_integer()];
+arg_types(hipe_bifs, alloc_data, 3) ->
+  [t_integer(), t_integer(), t_binary()];
 arg_types(hipe_bifs, array, 2) ->
   [t_non_neg_fixnum(), t_immediate()];
 arg_types(hipe_bifs, array_length, 1) ->
@@ -2494,8 +2496,8 @@ arg_types(hipe_bifs, call_count_on, 1) ->
   [t_mfa()];
 arg_types(hipe_bifs, check_crc, 1) ->
   [t_crc32()];
-arg_types(hipe_bifs, enter_code, 2) ->
-  [t_binary(), t_sup(t_nil(), t_tuple())];
+arg_types(hipe_bifs, enter_code, 3) ->
+  [t_binary(), t_sup(t_nil(), t_tuple()), t_binary()];
 arg_types(hipe_bifs, enter_sdesc, 1) ->
   [t_tuple([t_integer(), t_integer(), t_integer(), t_integer(), t_integer(), t_mfa()])];
 arg_types(hipe_bifs, find_na_or_make_stub, 1) ->
@@ -2540,6 +2542,9 @@ arg_types(hipe_bifs, write_u32, 2) ->
   [t_integer(), t_integer()];
 arg_types(hipe_bifs, write_u64, 2) ->
   [t_integer(), t_integer()];
+arg_types(hipe_bifs, alloc_loader_state, 1) ->
+  [t_atom()];
+
 %%------- lists ---------------------------------------------------------------
 arg_types(lists, all, 2) ->
   [t_fun([t_any()], t_boolean()), t_list()];
diff --git a/lib/kernel/src/hipe_unified_loader.erl b/lib/kernel/src/hipe_unified_loader.erl
index b3acd9ea74..6245646263 100644
--- a/lib/kernel/src/hipe_unified_loader.erl
+++ b/lib/kernel/src/hipe_unified_loader.erl
@@ -188,14 +188,16 @@ load_common(Mod, Bin, Beam, Architecture) ->
       put(closures_to_patch, []),
       WordSize = word_size(Architecture),
       WriteWord = write_word_fun(WordSize),
+      LoaderState = hipe_bifs:alloc_loader_state(Mod),
       %% Create data segment
       {ConstAddr,ConstMap2} =
-	create_data_segment(ConstAlign, ConstSize, ConstMap, WriteWord),
+	create_data_segment(ConstAlign, ConstSize, ConstMap, WriteWord,
+			    LoaderState),
       %% Find callees for which we may need trampolines.
       CalleeMFAs = find_callee_mfas(Refs, Architecture),
       %% Write the code to memory.
       {CodeAddress,Trampolines} =
-	enter_code(CodeSize, CodeBinary, CalleeMFAs),
+	enter_code(CodeSize, CodeBinary, CalleeMFAs, LoaderState),
       %% Construct CalleeMFA-to-trampoline mapping.
       TrampolineMap = mk_trampoline_map(CalleeMFAs, Trampolines,
                                         Architecture),
@@ -231,8 +233,8 @@ load_common(Mod, Bin, Beam, Architecture) ->
 	  AddressesOfClosuresToPatch =
 	    calculate_addresses(ClosurePatches, CodeAddress, FunDefs),
 	  export_funs(FunDefs),
-	  make_beam_stub(Mod, MD5, BeamBinary, FunDefs, AddressesOfClosuresToPatch,
-			 CodeAddress, byte_size(CodeBinary))
+	  make_beam_stub(Mod, LoaderState, MD5, BeamBinary, FunDefs,
+                         AddressesOfClosuresToPatch)
       end,
 
       %% Final clean up.
@@ -427,9 +429,9 @@ export_funs([FunDef | FunDefs]) ->
 export_funs([]) ->
   ok.
 
-make_beam_stub(Mod, MD5, Beam, FunDefs, ClosuresToPatch, CodeAddress, CodeSize) ->
+make_beam_stub(Mod, LoaderState, MD5, Beam, FunDefs, ClosuresToPatch) ->
   Fs = [{F,A,Address} || #fundef{address=Address, mfa={_M,F,A}} <- FunDefs],
-  Mod = code:make_stub_module(Mod, Beam, {Fs,ClosuresToPatch,MD5,CodeAddress,CodeSize}),
+  Mod = code:make_stub_module(LoaderState, Beam, {Fs,ClosuresToPatch,MD5}),
   ok.
 
 %%========================================================================
@@ -685,9 +687,9 @@ bif_address(Name) when is_atom(Name) ->
 %% memory, and produces a ConstMap2 mapping each constant's ConstNo to
 %% its runtime address, tagged if the constant is a term.
 %%
-create_data_segment(DataAlign, DataSize, DataList, WriteWord) ->
+create_data_segment(DataAlign, DataSize, DataList, WriteWord, LoaderState) ->
   %%io:format("create_data_segment: \nDataAlign: ~p\nDataSize: ~p\nDataList: ~p\n",[DataAlign,DataSize,DataList]),
-  DataAddress = hipe_bifs:alloc_data(DataAlign, DataSize),
+  DataAddress = hipe_bifs:alloc_data(DataAlign, DataSize, LoaderState),
   enter_data(DataList, [], DataAddress, DataSize, WriteWord).
 
 enter_data(List, ConstMap2, DataAddress, DataSize, WriteWord) ->
@@ -863,9 +865,10 @@ assert_local_patch(Address) when is_integer(Address) ->
 
 %% Beam: nil() | binary()  (used as a flag)
 
-enter_code(CodeSize, CodeBinary, CalleeMFAs) ->
+enter_code(CodeSize, CodeBinary, CalleeMFAs, LoaderState) ->
   true = byte_size(CodeBinary) =:= CodeSize,
-  {CodeAddress,Trampolines} = hipe_bifs:enter_code(CodeBinary, CalleeMFAs),
+  {CodeAddress,Trampolines} = hipe_bifs:enter_code(CodeBinary, CalleeMFAs,
+						   LoaderState),
   ?init_assert_patch(CodeAddress, byte_size(CodeBinary)),
   {CodeAddress,Trampolines}.