From 95dec7497abbdb06defcb23e99fa31e06597496e Mon Sep 17 00:00:00 2001
From: Sverker Eriksson <sverker@erlang.org>
Date: Wed, 3 May 2017 18:22:45 +0200
Subject: erts: Fix ets:select_replace with {const, NewTuple}

Enable ets:select_replace to do a generic single object
compare-and-swap operation of any ets-tuple using
a matchspec like this:

[{Old, [], [{const, New}]}]

The only exception when this does not work is if the key
contains maps or atoms looking like variables (like '$1').
---
 erts/emulator/beam/erl_db_util.c | 41 +++++++++++++++++++++++++---------------
 lib/stdlib/test/ets_SUITE.erl    | 20 ++++++++++++++++++++
 2 files changed, 46 insertions(+), 15 deletions(-)

diff --git a/erts/emulator/beam/erl_db_util.c b/erts/emulator/beam/erl_db_util.c
index 24b22eafb8..d1130fac87 100644
--- a/erts/emulator/beam/erl_db_util.c
+++ b/erts/emulator/beam/erl_db_util.c
@@ -1125,12 +1125,11 @@ error:
  * Returns true if 'b' is guaranteed to always construct
  * the same term as 'a' has matched.
  */
-static int db_match_eq_body(Eterm a, Eterm b)
+static int db_match_eq_body(Eterm a, Eterm b, int const_mode)
 {
     DECLARE_ESTACK(s);
     Uint arity;
     Eterm *ap, *bp;
-    int const_mode = 0;
     const Eterm CONST_MODE_OFF = THE_NON_VALUE;
 
     while (1) {
@@ -1224,6 +1223,7 @@ int db_match_keeps_key(int keypos, Eterm match, Eterm guard, Eterm body)
     Eterm single_body_subterm;
     Eterm single_body_subterm_key;
     Eterm* single_body_subterm_key_tpl;
+    int const_mode;
 
     if (!is_list(body)) {
         return 0;
@@ -1244,42 +1244,53 @@ int db_match_keeps_key(int keypos, Eterm match, Eterm guard, Eterm body)
         return 0;
     }
 
-    single_body_term_tpl = tuple_val(single_body_term);
-    if (arityval(*single_body_term_tpl) != 1) {
-        // not the 1-element tuple we're expecting
-        return 0;
-    }
-
     match_key = db_getkey(keypos, match);
     if (!is_value(match_key)) {
         // can't get key out of match
         return 0;
     }
 
-    single_body_subterm = single_body_term_tpl[1];
+    single_body_term_tpl = tuple_val(single_body_term);
+    if (single_body_term_tpl[0] == make_arityval(2) &&
+        single_body_term_tpl[1] == am_const) {
+        /* {const, {"ets-tuple constant"}} */
+        single_body_subterm = single_body_term_tpl[2];
+        const_mode = 1;
+    }
+    else if (*single_body_term_tpl == make_arityval(1)) {
+        /* {{"ets-tuple construction"}} */
+        single_body_subterm = single_body_term_tpl[1];
+        const_mode = 0;
+    }
+    else {
+        /* not a tuple construction */
+        return 0;
+    }
+
     single_body_subterm_key = db_getkey(keypos, single_body_subterm);
     if (!is_value(single_body_subterm_key)) {
         // can't get key out of single body subterm
         return 0;
     }
 
-    if (db_match_eq_body(match_key, single_body_subterm_key)) {
+    if (db_match_eq_body(match_key, single_body_subterm_key, const_mode)) {
         /* tuple with same key is returned */
         return 1;
     }
 
-    if (!is_tuple(single_body_subterm_key)) {
-        /* can't possibly be an element instruction */
+    if (const_mode) {
+        /* constant key did not match */
         return 0;
     }
 
-    single_body_subterm_key_tpl = tuple_val(single_body_subterm_key);
-    if (arityval(*single_body_subterm_key_tpl) != 3) {
+    if (!is_tuple(single_body_subterm_key)) {
         /* can't possibly be an element instruction */
         return 0;
     }
 
-    if (single_body_subterm_key_tpl[1] == am_element &&
+    single_body_subterm_key_tpl = tuple_val(single_body_subterm_key);
+    if (single_body_subterm_key_tpl[0] == make_arityval(3) &&
+        single_body_subterm_key_tpl[1] == am_element &&
         single_body_subterm_key_tpl[3] == am_DollarUnderscore &&
         single_body_subterm_key_tpl[2] == make_small(keypos))
     {
diff --git a/lib/stdlib/test/ets_SUITE.erl b/lib/stdlib/test/ets_SUITE.erl
index 2dfd481665..13a3f99e7e 100644
--- a/lib/stdlib/test/ets_SUITE.erl
+++ b/lib/stdlib/test/ets_SUITE.erl
@@ -1349,6 +1349,26 @@ t_select_replace(Config) when is_list(Config) ->
                    NeqPairs(X,Y)) || X <- Terms, Y <- Terms],
 
 
+    %% Wrap entire tuple with 'const'
+    [[begin
+          Old = {Key, 1, 2},
+          ets:insert(T2, Old),
+          1 = ets:select_replace(T2, [{Old, [], [{const, New}]}]),
+          [New] = ets:lookup(T2, Key),
+          ets:delete(T2, Key)
+      end || New <- [{Key, 1, 2}, {Key, 3, 4}, {Key, 1}, {Key, 1, 2, 3}, {Key}]
+     ]
+     || Key <- [{1, tuple}, {nested, {tuple, {a,b}}} | Terms]],
+
+    %% 'const' wrap does not work with maps or variables in keys
+    [[begin
+          Old = {Key, 1, 2},
+          {'EXIT',{badarg,_}} = (catch ets:select_replace(T2, [{Old, [], [{const, New}]}]))
+      end || New <- [{Key, 1, 2}, {Key, 3, 4}, {Key, 1}, {Key, 1, 2, 3}, {Key}]
+     ]
+     || Key <- [#{a => 1}, {nested, #{a => 1}}, '$1']],
+
+
     ets:delete(T2),
 
     verify_etsmem(EtsMem).
-- 
cgit v1.2.3