Optimize d3.{map,set}.

mbostock · mbostock · commit 9a3de7d775e4 · 2014-10-14T10:13:12.000-07:00
Rather than always escaping keys with a null-prefix, only escape keys that can
conflict with built-in properties (either on Object, d3_Map or d3_Set). By only
prefixing these special keys, we can avoid the cost of constructing new strings
in the common case.

To check whether a key is special, we check whether it is in an empty map
instance. In addition, keys that are already null-prefixed must be prefixed a
second time to correctly unescape.
diff --git a/d3.js b/d3.js
@@ -245,58 +245,81 @@
   d3_class(d3_Map, {
     has: d3_map_has,
     get: function(key) {
-      return this[d3_map_prefix + key];
+      return this[d3_map_escape(key)];
     },
     set: function(key, value) {
-      return this[d3_map_prefix + key] = value;
+      return this[d3_map_escape(key)] = value;
     },
     remove: d3_map_remove,
     keys: d3_map_keys,
     values: function() {
       var values = [];
-      this.forEach(function(key, value) {
-        values.push(value);
-      });
+      for (var key in this) {
+        if (d3_map_hasOwnProperty.call(this, key)) {
+          values.push(this[key]);
+        }
+      }
       return values;
     },
     entries: function() {
       var entries = [];
-      this.forEach(function(key, value) {
-        entries.push({
-          key: key,
-          value: value
-        });
-      });
+      for (var key in this) {
+        if (d3_map_hasOwnProperty.call(this, key)) {
+          entries.push({
+            key: d3_map_unescape(key),
+            value: this[key]
+          });
+        }
+      }
       return entries;
     },
     size: d3_map_size,
     empty: d3_map_empty,
     forEach: function(f) {
-      for (var key in this) if (key.charCodeAt(0) === d3_map_prefixCode) f.call(this, key.slice(1), this[key]);
+      for (var key in this) {
+        if (d3_map_hasOwnProperty.call(this, key)) {
+          f.call(this, d3_map_unescape(key), this[key]);
+        }
+      }
     }
   });
-  var d3_map_prefix = "\x00", d3_map_prefixCode = d3_map_prefix.charCodeAt(0);
+  var d3_map_prefix = "\x00", d3_map_builtin = new d3_Map(), d3_map_hasOwnProperty = Object.prototype.hasOwnProperty;
+  function d3_map_escape(key) {
+    return (key += "") in d3_map_builtin || key[0] === d3_map_prefix ? d3_map_prefix + key : key;
+  }
+  function d3_map_unescape(key) {
+    return (key += "")[0] === d3_map_prefix ? key.slice(1) : key;
+  }
   function d3_map_has(key) {
-    return d3_map_prefix + key in this;
+    return d3_map_escape(key) in this;
   }
   function d3_map_remove(key) {
-    key = d3_map_prefix + key;
-    return key in this && delete this[key];
+    return (key = d3_map_escape(key)) in this && delete this[key];
   }
   function d3_map_keys() {
     var keys = [];
-    this.forEach(function(key) {
-      keys.push(key);
-    });
+    for (var key in this) {
+      if (d3_map_hasOwnProperty.call(this, key)) {
+        keys.push(d3_map_unescape(key));
+      }
+    }
     return keys;
   }
   function d3_map_size() {
     var size = 0;
-    for (var key in this) if (key.charCodeAt(0) === d3_map_prefixCode) ++size;
+    for (var key in this) {
+      if (d3_map_hasOwnProperty.call(this, key)) {
+        ++size;
+      }
+    }
     return size;
   }
   function d3_map_empty() {
-    for (var key in this) if (key.charCodeAt(0) === d3_map_prefixCode) return false;
+    for (var key in this) {
+      if (d3_map_hasOwnProperty.call(this, key)) {
+        return false;
+      }
+    }
     return true;
   }
   d3.nest = function() {
@@ -370,21 +393,23 @@
   function d3_Set() {}
   d3_class(d3_Set, {
     has: d3_map_has,
-    add: function(value) {
-      this[d3_map_prefix + value] = true;
-      return value;
-    },
-    remove: function(value) {
-      value = d3_map_prefix + value;
-      return value in this && delete this[value];
+    add: function(key) {
+      this[d3_map_escape(key)] = true;
+      return key;
     },
+    remove: d3_map_remove,
     values: d3_map_keys,
     size: d3_map_size,
     empty: d3_map_empty,
     forEach: function(f) {
-      for (var value in this) if (value.charCodeAt(0) === d3_map_prefixCode) f.call(this, value.slice(1));
+      for (var key in this) {
+        if (d3_map_hasOwnProperty.call(this, key)) {
+          f.call(this, d3_map_unescape(key));
+        }
+      }
     }
   });
+  for (var key in new d3_Set()) d3_map_builtin[key] = true;
   d3.behavior = {};
   d3.rebind = function(target, source) {
     var i = 1, n = arguments.length, method;
diff --git a/d3.min.js b/d3.min.js
diff --git a/src/arrays/map.js b/src/arrays/map.js
@@ -12,55 +12,87 @@ function d3_Map() {}
 d3_class(d3_Map, {
   has: d3_map_has,
   get: function(key) {
-    return this[d3_map_prefix + key];
+    return this[d3_map_escape(key)];
   },
   set: function(key, value) {
-    return this[d3_map_prefix + key] = value;
+    return this[d3_map_escape(key)] = value;
   },
   remove: d3_map_remove,
   keys: d3_map_keys,
   values: function() {
     var values = [];
-    this.forEach(function(key, value) { values.push(value); });
+    for (var key in this) {
+      if (d3_map_hasOwnProperty.call(this, key)) {
+        values.push(this[key]);
+      }
+    }
     return values;
   },
   entries: function() {
     var entries = [];
-    this.forEach(function(key, value) { entries.push({key: key, value: value}); });
+    for (var key in this) {
+      if (d3_map_hasOwnProperty.call(this, key)) {
+        entries.push({key: d3_map_unescape(key), value: this[key]});
+      }
+    }
     return entries;
   },
   size: d3_map_size,
   empty: d3_map_empty,
   forEach: function(f) {
-    for (var key in this) if (key.charCodeAt(0) === d3_map_prefixCode) f.call(this, key.slice(1), this[key]);
+    for (var key in this) {
+      if (d3_map_hasOwnProperty.call(this, key)) {
+        f.call(this, d3_map_unescape(key), this[key]);
+      }
+    }
   }
 });
 
 var d3_map_prefix = "\0", // prevent collision with built-ins
-    d3_map_prefixCode = d3_map_prefix.charCodeAt(0);
+    d3_map_builtin = new d3_Map,
+    d3_map_hasOwnProperty = Object.prototype.hasOwnProperty;
+
+function d3_map_escape(key) {
+  return (key += "") in d3_map_builtin || key[0] === d3_map_prefix ? d3_map_prefix + key : key;
+}
+
+function d3_map_unescape(key) {
+  return (key += "")[0] === d3_map_prefix ? key.slice(1) : key;
+}
 
 function d3_map_has(key) {
-  return d3_map_prefix + key in this;
+  return d3_map_escape(key) in this;
 }
 
 function d3_map_remove(key) {
-  key = d3_map_prefix + key;
-  return key in this && delete this[key];
+  return (key = d3_map_escape(key)) in this && delete this[key];
 }
 
 function d3_map_keys() {
   var keys = [];
-  this.forEach(function(key) { keys.push(key); });
+  for (var key in this) {
+    if (d3_map_hasOwnProperty.call(this, key)) {
+      keys.push(d3_map_unescape(key));
+    }
+  }
   return keys;
 }
 
 function d3_map_size() {
   var size = 0;
-  for (var key in this) if (key.charCodeAt(0) === d3_map_prefixCode) ++size;
+  for (var key in this) {
+    if (d3_map_hasOwnProperty.call(this, key)) {
+      ++size;
+    }
+  }
   return size;
 }
 
 function d3_map_empty() {
-  for (var key in this) if (key.charCodeAt(0) === d3_map_prefixCode) return false;
+  for (var key in this) {
+    if (d3_map_hasOwnProperty.call(this, key)) {
+      return false;
+    }
+  }
   return true;
 }
diff --git a/src/arrays/set.js b/src/arrays/set.js
@@ -11,18 +11,22 @@ function d3_Set() {}
 
 d3_class(d3_Set, {
   has: d3_map_has,
-  add: function(value) {
-    this[d3_map_prefix + value] = true;
-    return value;
-  },
-  remove: function(value) {
-    value = d3_map_prefix + value;
-    return value in this && delete this[value];
+  add: function(key) {
+    this[d3_map_escape(key)] = true;
+    return key;
   },
+  remove: d3_map_remove,
   values: d3_map_keys,
   size: d3_map_size,
   empty: d3_map_empty,
   forEach: function(f) {
-    for (var value in this) if (value.charCodeAt(0) === d3_map_prefixCode) f.call(this, value.slice(1));
+    for (var key in this) {
+      if (d3_map_hasOwnProperty.call(this, key)) {
+        f.call(this, d3_map_unescape(key));
+      }
+    }
   }
 });
+
+// In case Object.defineProperty is not supported…
+for (var key in new d3_Set) d3_map_builtin[key] = true;
diff --git a/test/arrays/map-test.js b/test/arrays/map-test.js
@@ -119,6 +119,12 @@ suite.addBatch({
       "returns an array of string keys": function(map) {
         var m = map({foo: 1, bar: "42"});
         assert.deepEqual(m.keys().sort(), ["bar", "foo"]);
+      },
+      "properly unescapes zero-prefixed keys": function(map) {
+        var m = map();
+        m.set("__proto__", 42);
+        m.set("\0weird", 42);
+        assert.deepEqual(m.keys().sort(), ["\0weird", "__proto__"]);
       }
     },
     "values": {
@@ -246,6 +252,11 @@ suite.addBatch({
         m.set("__proto__", 42);
         assert.equal(m.get("__proto__"), 42);
       },
+      "can set keys using zero-prefixed names": function(map) {
+        var m = map();
+        m.set("\0weird", 42);
+        assert.equal(m.get("\0weird"), 42);
+      },
       "coerces keys to strings": function(map) {
         var m = map();
         m.set(42, 1);
diff --git a/test/arrays/set-test.js b/test/arrays/set-test.js
@@ -108,6 +108,10 @@ suite.addBatch({
         var s = set(["bar", "foo"]);
         assert.deepEqual(s.values().sort(_.ascending), ["bar", "foo"]);
       },
+      "properly unescapes zero-prefixed keys": function(set) {
+        var s = set(["__proto__", "\0weird"]);
+        assert.deepEqual(s.values().sort(_.ascending), ["\0weird", "__proto__"]);
+      },
       "observes changes via add and remove": function(set) {
         var s = set(["foo", "bar"]);
         assert.deepEqual(s.values().sort(_.ascending), ["bar", "foo"]);
@@ -162,6 +166,11 @@ suite.addBatch({
         s.add("__proto__");
         assert.isTrue(s.has("__proto__"));
       },
+      "can add values using zero-prefixed names": function(set) {
+        var s = set();
+        s.add("\0weird");
+        assert.isTrue(s.has("\0weird"));
+      },
       "coerces values to strings": function(set) {
         var s = set();
         s.add(42);
