2 files changed, 219 insertions, 93 deletions

diff --git a/erts/emulator/beam/beam_emu.c b/erts/emulator/beam/beam_emu.c
index 1cad31be2c..b666b7c3f7 100644
--- a/erts/emulator/beam/beam_emu.c
+++ b/erts/emulator/beam/beam_emu.c
@@ -959,8 +959,10 @@ static BeamInstr* apply_fun(Process* p, Eterm fun,
 static Eterm new_fun(Process* p, Eterm* reg,
 		     ErlFunEntry* fe, int num_free) NOINLINE;
 static Eterm new_map(Process* p, Eterm* reg, BeamInstr* I) NOINLINE;
-static Eterm update_map(Process* p, Eterm* reg,
-			Eterm map, BeamInstr* I) NOINLINE;
+static Eterm update_map_assoc(Process* p, Eterm* reg,
+			      Eterm map, BeamInstr* I) NOINLINE;
+static Eterm update_map_exact(Process* p, Eterm* reg,
+			      Eterm map, BeamInstr* I) NOINLINE;
 static int has_not_map_field(Eterm map, Eterm key);
 static Eterm get_map_element(Eterm map, Eterm key);
 
@@ -2353,21 +2355,39 @@ void process_main(void)
      Next(4+Arg(3));
  }
 
- OpCase(update_map_jddII): {
+ OpCase(update_map_assoc_jddII): {
      Eterm res;
      Eterm map;
 
      GetArg1(1, map);
      x(0) = r(0);
      SWAPOUT;
-     res = update_map(c_p, reg, map, I);
+     res = update_map_assoc(c_p, reg, map, I);
      SWAPIN;
-     if (res) {
+     if (is_value(res)) {
 	 r(0) = x(0);
 	 StoreResult(res, Arg(2));
 	 Next(5+Arg(4));
      } else {
-	 goto lb_Cl_error;
+	 goto badarg;
+     }
+ }
+
+ OpCase(update_map_exact_jddII): {
+     Eterm res;
+     Eterm map;
+
+     GetArg1(1, map);
+     x(0) = r(0);
+     SWAPOUT;
+     res = update_map_exact(c_p, reg, map, I);
+     SWAPIN;
+     if (is_value(res)) {
+	 r(0) = x(0);
+	 StoreResult(res, Arg(2));
+	 Next(5+Arg(4));
+     } else {
+	 goto badarg;
      }
  }
 
@@ -6387,18 +6407,10 @@ new_map(Process* p, Eterm* reg, BeamInstr* I)
     return make_map(mp);
 }
 
-
-/* This entire instruction will be split into two.
- * 1) update_map_exact (literals) <- this can be much more optimized
- * 2) update_map_assoc (literals)
- * Also update_map is pretty bad code as it stands now.
- */
-
 static Eterm
-update_map(Process* p, Eterm* reg, Eterm map, BeamInstr* I)
+update_map_assoc(Process* p, Eterm* reg, Eterm map, BeamInstr* I)
 {
     Uint n;
-    Uint i;
     Uint num_old;
     Uint num_updates;
     Uint need;
@@ -6413,7 +6425,7 @@ update_map(Process* p, Eterm* reg, Eterm map, BeamInstr* I)
     Eterm* kp;
 
     if (is_not_map(map)) {
-	return 0;
+	return THE_NON_VALUE;
     }
 
     old_mp = (map_t *) map_val(map);
@@ -6428,11 +6440,12 @@ update_map(Process* p, Eterm* reg, Eterm map, BeamInstr* I)
     }
 
     /*
-     * Allocate heap space for the worst case (i.e. all keys are new).
+     * Allocate heap space for the worst case (i.e. all keys in the
+     * update list are new).
      */
 
     num_updates = Arg(4) / 2;
-    need = 2*(num_old+num_updates) + 4;
+    need = 2*(num_old+num_updates) + 1 + sizeof(map_t) / sizeof(Eterm);
     if (HeapWordsLeft(p) < need) {
 	Uint live = Arg(3);
 	reg[live] = map;
@@ -6442,67 +6455,37 @@ update_map(Process* p, Eterm* reg, Eterm map, BeamInstr* I)
     }
 
     /*
-     * Update map, optimistically assuming that there are no
-     * new keys, allowing us to keep the old key tuple.
+     * Build the skeleton for the map, ready to be filled in.
+     *
+     * +-----------------------------------+
+     * | (Space for aritvyal for keys)     | <-----------+
+     * +-----------------------------------+		 |
+     * | (Space for key 1)		   |		 |    <-- kp
+     * +-----------------------------------+		 |
+     *        .				    		 |
+     *        .				    		 |
+     *        .				    		 |
+     * +-----------------------------------+		 |
+     * | (Space for last key)		   |		 |
+     * +-----------------------------------+		 |
+     * | MAP_HEADER			   |		 |
+     * +-----------------------------------+		 |
+     * | (Space for number of keys/values) |		 |
+     * +-----------------------------------+		 |
+     * | Boxed tuple pointer            >----------------+
+     * +-----------------------------------+
+     * | (Space for value 1)		   |                  <-- hp
+     * +-----------------------------------+
      */
 
-    hp = p->htop;
-    E  = p->stop;
-
-    old_vals = map_get_values(old_mp);
-    old_keys = map_get_keys(old_mp);
+    E = p->stop;
+    kp = p->htop + 1;		/* Point to first key */
+    hp = kp + num_old + num_updates;
 
     res = make_map(hp);
-    mp = (map_t *)hp; hp += 3;
+    mp = (map_t *)hp;
+    hp += sizeof(map_t) / sizeof(Eterm);
     mp->thing_word = MAP_HEADER;
-    mp->size = num_old;
-    mp->keys = old_mp->keys;
-
-    ASSERT(num_updates > 0);
-
-    /* Get array of key/value pairs to be updated */
-    new_p = &Arg(5);
-    GET_TERM(*new_p, new_key);
-    
-    n = num_updates;
-
-    for (i = 0; i < num_old; i++) {
-	if (new_key == THE_NON_VALUE || !eq(*old_keys, new_key)) {
-	    /* not same keys */
-	    *hp++ = *old_vals;
-	} else {
-	    GET_TERM(new_p[1], *hp);
-	    hp++;
-	    n--;
-	    if (n == 0) {
-		new_key = THE_NON_VALUE;
-	    } else {
-		new_p += 2;
-		GET_TERM(*new_p, new_key);
-	    }
-	}
-	old_vals++, old_keys++;
-    }
-
-    /*
-     * If we have exhausted the update list we are done.
-     */
-
-    if (n == 0) {
-	p->htop = hp;
-	return res;
-    }
-
-    /*
-     * There were some new keys. We'll have to start over and rebuild
-     * the key tuple too.
-     */
-
-    kp    = p->htop;
-    *kp++ = make_arityval(0);
-
-    res = make_map(hp);
-    mp  = (map_t *)hp; hp += 3;
     mp->keys = make_tuple(kp-1);
 
     old_vals = map_get_values(old_mp);
@@ -6512,20 +6495,28 @@ update_map(Process* p, Eterm* reg, Eterm map, BeamInstr* I)
     GET_TERM(*new_p, new_key);
     n = num_updates;
 
+    /*
+     * Fill in keys and values, until we run out of either updates
+     * or old values and keys.
+     */
+
     for (;;) {
 	Eterm key;
 	Sint c;
 
 	ASSERT(kp < (Eterm *)mp);
 	key = *old_keys;
-	if ((c = cmp(key, new_key)) < 0) {
+	if ((c = CMP(key, new_key)) < 0) {
+	    /* Copy old key and value */
 	    *kp++ = key;
 	    *hp++ = *old_vals;
 	    old_keys++, old_vals++, num_old--;
 	} else {		/* Replace or insert new */
-	    *kp++ = new_key;
 	    GET_TERM(new_p[1], *hp++);
-	    if (c == 0) {	/* If replacement */
+	    if (c > 0) {	/* If new new key */
+		*kp++ = new_key;
+	    } else {		/* If replacement */
+		*kp++ = key;
 		old_keys++, old_vals++, num_old--;
 	    }
 	    n--;
@@ -6541,32 +6532,162 @@ update_map(Process* p, Eterm* reg, Eterm map, BeamInstr* I)
 	}
     }
 
-    while (n-- > 0) {
-	GET_TERM(new_p[0], *kp++);
-	GET_TERM(new_p[1], *hp++);
-	new_p += 2;
-    }
-
     /*
-     * All updates done. Now copy the remaining part of the frame's
-     * keys and values.
+     * At this point, we have run out of either old keys and values,
+     * or the update list. In other words, at least of one n and
+     * num_old must be zero.
      */
 
-    while (num_old-- > 0) {
-	ASSERT(kp < (Eterm *)mp);
-	*kp++ = *old_keys++;
-	*hp++ = *old_vals++;
+    if (n > 0) {
+	/*
+	 * All old keys and values have been copied, but there
+	 * are still new keys and values in the update list that
+	 * must be copied.
+	 */
+	ASSERT(num_old == 0);
+	while (n-- > 0) {
+	    GET_TERM(new_p[0], *kp++);
+	    GET_TERM(new_p[1], *hp++);
+	    new_p += 2;
+	}
+    } else {
+	/*
+	 * All updates are now done. We may still have old
+	 * keys and values that we must copy.
+	 */
+	ASSERT(n == 0);
+	while (num_old-- > 0) {
+	    ASSERT(kp < (Eterm *)mp);
+	    *kp++ = *old_keys++;
+	    *hp++ = *old_vals++;
+	}
     }
+
+    /*
+     * Calculate how many values that are unused at the end of the
+     * key tuple and fill it out with a bignum header.
+     */
     if ((n = (Eterm *)mp - kp) > 0) {
 	*kp = make_pos_bignum_header(n-1);
     }
+
+    /*
+     * Fill in the size of the map in both the key tuple and in the map.
+     */
+
     n = kp - p->htop - 1;	/* Actual number of keys/values */
     *p->htop = make_arityval(n);
-    mp->thing_word = MAP_HEADER;
     mp->size = n;
     p->htop = hp;
     return res;
 }
+
+/*
+ * Update values for keys that already exist in the map.
+ */
+
+static Eterm
+update_map_exact(Process* p, Eterm* reg, Eterm map, BeamInstr* I)
+{
+    Uint n;
+    Uint i;
+    Uint num_old;
+    Uint need;
+    map_t *old_mp, *mp;
+    Eterm res;
+    Eterm* hp;
+    Eterm* E;
+    Eterm* old_keys;
+    Eterm* old_vals;
+    Eterm* new_p;
+    Eterm new_key;
+
+    if (is_not_map(map)) {
+	return THE_NON_VALUE;
+    }
+
+    old_mp = (map_t *) map_val(map);
+    num_old = map_get_size(old_mp);
+
+    /*
+     * If the old map is empty, create a new map.
+     */
+
+    if (num_old == 0) {
+	return new_map(p, reg, I+1);
+    }
+
+    /*
+     * Allocate the exact heap space needed.
+     */
+
+    need = num_old + sizeof(map_t) / sizeof(Eterm);
+    if (HeapWordsLeft(p) < need) {
+	Uint live = Arg(3);
+	reg[live] = map;
+	erts_garbage_collect(p, need, reg, live+1);
+	map      = reg[live];
+	old_mp   = (map_t *)map_val(map);
+    }
+
+    /*
+     * Update map, keeping the old key tuple.
+     */
+
+    hp = p->htop;
+    E = p->stop;
+
+    old_vals = map_get_values(old_mp);
+    old_keys = map_get_keys(old_mp);
+
+    res = make_map(hp);
+    mp = (map_t *)hp;
+    hp += sizeof(map_t) / sizeof(Eterm);
+    mp->thing_word = MAP_HEADER;
+    mp->size = num_old;
+    mp->keys = old_mp->keys;
+
+    /* Get array of key/value pairs to be updated */
+    new_p = &Arg(5);
+    GET_TERM(*new_p, new_key);
+
+    /* Update all values */
+    n = Arg(4) / 2;		/* Number of values to be updated */
+    ASSERT(n > 0);
+    for (i = 0; i < num_old; i++) {
+	if (!EQ(*old_keys, new_key)) {
+	    /* Not same keys */
+	    *hp++ = *old_vals;
+	} else {
+	    GET_TERM(new_p[1], *hp);
+	    hp++;
+	    n--;
+	    if (n == 0) {
+		/*
+		 * All updates done. Copy remaining values
+		 * and return the result.
+		 */
+		for (i++, old_vals++; i < num_old; i++) {
+		    *hp++ = *old_vals++;
+		}
+		ASSERT(hp == p->htop + need);
+		p->htop = hp;
+		return res;
+	    } else {
+		new_p += 2;
+		GET_TERM(*new_p, new_key);
+	    }
+	}
+	old_vals++, old_keys++;
+    }
+
+    /*
+     * Updates left. That means that at least one the keys in the
+     * update list did not previously exist.
+     */
+    ASSERT(hp == p->htop + need);
+    return THE_NON_VALUE;
+}
 #undef GET_TERM
 
 int catchlevel(Process *p)
diff --git a/erts/emulator/beam/ops.tab b/erts/emulator/beam/ops.tab
index b89bdb2e3a..f35997efee 100644
--- a/erts/emulator/beam/ops.tab
+++ b/erts/emulator/beam/ops.tab
@@ -1474,11 +1474,16 @@ apply_last I P
 # has_map_field Fail Src Key => jump Fail
 # get_map_element Fail Src Key Dst => jump Fail
 
-put_map F n Dst Live Size Rest=* => new_map F Dst Live Size Rest
-put_map F Src Dst Live Size Rest=* => update_map F Src Dst Live Size Rest
+put_map_assoc F n Dst Live Size Rest=* => new_map F Dst Live Size Rest
+put_map_exact F n Dst Live Size Rest=* => new_map F Dst Live Size Rest
+put_map_assoc F Src Dst Live Size Rest=* => \
+	       update_map_assoc F Src Dst Live Size Rest
+put_map_exact F Src Dst Live Size Rest=* => \
+	      update_map_exact F Src Dst Live Size Rest
 
 new_map j d I I
-update_map j d d I I
+update_map_assoc j d d I I
+update_map_exact j d d I I
 
 is_map Fail cq => jump Fail