Skip to content

Commit

Permalink
Two minor memory improvements to Dict.
Browse files Browse the repository at this point in the history 
(1) When we are constructing a new `Dict` instance and are in the `Mutability.IMMUTABLE` situation, store `contents` as a `ImmutableMap` not a `LinkedHashMap`. The former suffices and uses less memory. There was already a TODO for this from adonovan@. Unfortunately, this optimization is minor in practice because the common case is that a `Dict` is constructed without `Mutability.IMMUTABLE`, but then `Mutability#freeze` gets called later. Therefore I think the bigger optimization opportunity would be to make it so that call to `Mutability#freeze` would magically go in and replace the `LinkedHashMap` instance with a `ImmutableMap` instance (this is not currently mechanically possible; see the javadoc for `Mutability#freeze`). I added a TODO for this.

(2) Use `Dict#EMPTY` appropriately in all the code paths that construct `Dict` instances to avoid having multiple physical objects all for the same logical concept of "immutable empty dict".

PiperOrigin-RevId: 432973333
  • Loading branch information
@haxorz @copybara-github
haxorz authored and copybara-github committed Mar 7, 2022
1 parent 3b57ffa commit 24ec910
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 31 deletions.
116 changes: 90 additions & 26 deletions src/main/java/net/starlark/java/eval/Dict.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
package net.starlark.java.eval;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import java.util.ArrayList;
Expand All @@ -24,6 +25,7 @@
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
import java.util.function.BiConsumer;
import javax.annotation.Nullable;
import net.starlark.java.annot.Param;
import net.starlark.java.annot.StarlarkBuiltin;
Expand Down Expand Up @@ -109,29 +111,43 @@ public class Dict<K, V>
StarlarkIndexable,
StarlarkIterable<K> {

// TODO(adonovan): for dicts that are born frozen, use ImmutableMap, which is also
// insertion-ordered and has smaller Entries (singly linked, no hash).
private final LinkedHashMap<K, V> contents;
private final Map<K, V> contents;
private int iteratorCount; // number of active iterators (unused once frozen)

/** Final except for {@link #unsafeShallowFreeze}; must not be modified any other way. */
private Mutability mutability;

private Dict(@Nullable Mutability mutability, LinkedHashMap<K, V> contents) {
this.mutability = mutability == null ? Mutability.IMMUTABLE : mutability;
private Dict(Mutability mutability, LinkedHashMap<K, V> contents) {
Preconditions.checkNotNull(mutability);
Preconditions.checkState(mutability != Mutability.IMMUTABLE);
this.mutability = mutability;
// TODO(bazel-team): Memory optimization opportunity: Make it so that a call to
// `mutability.freeze()` causes `contents` here to become an ImmutableMap.
this.contents = contents;
}

private Dict(@Nullable Mutability mutability) {
this(mutability, new LinkedHashMap<>());
private Dict(ImmutableMap<K, V> contents) {
// An immutable dict might as well store its contents as an ImmutableMap, since ImmutableMap
// both is more memory-efficient than LinkedHashMap and also it has the requisite deterministic
// iteration order.
this.mutability = Mutability.IMMUTABLE;
this.contents = contents;
}

/**
* Takes ownership of the supplied LinkedHashMap and returns a new Dict that wraps it. The caller
* must not subsequently modify the map, but the Dict may do so.
*/
static <K, V> Dict<K, V> wrap(@Nullable Mutability mutability, LinkedHashMap<K, V> contents) {
return new Dict<>(mutability, contents);
static <K, V> Dict<K, V> wrap(@Nullable Mutability mu, LinkedHashMap<K, V> contents) {
if (mu == null) {
mu = Mutability.IMMUTABLE;
}
if (mu == Mutability.IMMUTABLE && contents.isEmpty()) {
return empty();
}
// #wrap is used in situations where the resulting Dict isn't necessarily retained [forever].
// So, don't make an ImmutableMap copy of `contents`, as #copyOf would do.
return new Dict<>(mu, contents);
}

@Override
Expand Down Expand Up @@ -394,7 +410,7 @@ public StarlarkList<?> keys(StarlarkThread thread) throws EvalException {
return StarlarkList.wrap(thread.mutability(), array);
}

private static final Dict<?, ?> EMPTY = of(Mutability.IMMUTABLE);
private static final Dict<?, ?> EMPTY = new Dict<>(ImmutableMap.of());

/** Returns an immutable empty dict. */
// Safe because the empty singleton is immutable.
Expand All @@ -405,24 +421,49 @@ public static <K, V> Dict<K, V> empty() {

/** Returns a new empty dict with the specified mutability. */
public static <K, V> Dict<K, V> of(@Nullable Mutability mu) {
return new Dict<>(mu);
if (mu == null) {
mu = Mutability.IMMUTABLE;
}
if (mu == Mutability.IMMUTABLE) {
return empty();
} else {
return new Dict<>(mu, new LinkedHashMap<>());
}
}

/** Returns a new dict with the specified mutability containing the entries of {@code m}. */
public static <K, V> Dict<K, V> copyOf(@Nullable Mutability mu, Map<? extends K, ? extends V> m) {
if (mu == null && m instanceof Dict && ((Dict) m).isImmutable()) {
if (mu == null) {
mu = Mutability.IMMUTABLE;
}

if (mu == Mutability.IMMUTABLE && m instanceof Dict && ((Dict) m).isImmutable()) {
@SuppressWarnings("unchecked")
Dict<K, V> dict = (Dict<K, V>) m; // safe
return dict;
}

Dict<K, V> dict = new Dict<>(mu);
for (Map.Entry<? extends K, ? extends V> e : m.entrySet()) {
dict.contents.put(
Starlark.checkValid(e.getKey()), //
Starlark.checkValid(e.getValue()));
if (mu == Mutability.IMMUTABLE) {
if (m.isEmpty()) {
return empty();
}
ImmutableMap.Builder<K, V> immutableMapBuilder =
ImmutableMap.builderWithExpectedSize(m.size());
for (Map.Entry<? extends K, ? extends V> e : m.entrySet()) {
immutableMapBuilder.put(
Starlark.checkValid(e.getKey()), //
Starlark.checkValid(e.getValue()));
}
return new Dict<>(immutableMapBuilder.buildOrThrow());
} else {
LinkedHashMap<K, V> linkedHashMap = new LinkedHashMap<>();
for (Map.Entry<? extends K, ? extends V> e : m.entrySet()) {
linkedHashMap.put(
Starlark.checkValid(e.getKey()), //
Starlark.checkValid(e.getValue()));
}
return new Dict<>(mu, linkedHashMap);
}
return dict;
}

/** Returns an immutable dict containing the entries of {@code m}. */
Expand Down Expand Up @@ -462,27 +503,50 @@ public Dict<K, V> buildImmutable() {

/** Returns a new {@link ImmutableKeyTrackingDict} containing the entries added so far. */
public ImmutableKeyTrackingDict<K, V> buildImmutableWithKeyTracking() {
return new ImmutableKeyTrackingDict<>(buildMap());
return new ImmutableKeyTrackingDict<>(buildImmutableMap());
}

/**
* Returns a new Dict containing the entries added so far. The result has the specified
* mutability; null means immutable.
*/
public Dict<K, V> build(@Nullable Mutability mu) {
return wrap(mu, buildMap());
if (mu == null) {
mu = Mutability.IMMUTABLE;
}

if (mu == Mutability.IMMUTABLE) {
if (items.isEmpty()) {
return empty();
}
return new Dict<>(buildImmutableMap());
} else {
return new Dict<>(mu, buildLinkedHashMap());
}
}

private LinkedHashMap<K, V> buildMap() {
int n = items.size() / 2;
LinkedHashMap<K, V> map = Maps.newLinkedHashMapWithExpectedSize(n);
private void populateMap(int n, BiConsumer<K, V> mapEntryConsumer) {
for (int i = 0; i < n; i++) {
@SuppressWarnings("unchecked")
K k = (K) items.get(2 * i); // safe
@SuppressWarnings("unchecked")
V v = (V) items.get(2 * i + 1); // safe
map.put(k, v);
mapEntryConsumer.accept(k, v);
}
}

private ImmutableMap<K, V> buildImmutableMap() {
int n = items.size() / 2;
ImmutableMap.Builder<K, V> immutableMapBuilder = ImmutableMap.builderWithExpectedSize(n);
populateMap(n, immutableMapBuilder::put);
// Respect the desired semantics of Builder#put.
return immutableMapBuilder.buildKeepingLast();
}

private LinkedHashMap<K, V> buildLinkedHashMap() {
int n = items.size() / 2;
LinkedHashMap<K, V> map = Maps.newLinkedHashMapWithExpectedSize(n);
populateMap(n, map::put);
return map;
}
}
Expand Down Expand Up @@ -695,8 +759,8 @@ public V remove(Object key) {
public static final class ImmutableKeyTrackingDict<K, V> extends Dict<K, V> {
private final ImmutableSet.Builder<K> accessedKeys = ImmutableSet.builder();

private ImmutableKeyTrackingDict(LinkedHashMap<K, V> contents) {
super(Mutability.IMMUTABLE, contents);
private ImmutableKeyTrackingDict(ImmutableMap<K, V> contents) {
super(contents);
}

public ImmutableSet<K> getAccessedKeys() {
Expand Down
12 changes: 7 additions & 5 deletions src/test/java/net/starlark/java/eval/StarlarkMutableTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ public void testDictBuilder() throws Exception {
Dict.<String, String>builder()
.put("one", "1")
.put("two", "2.0")
.put("two", "2") // overrwrites previous entry
.put("two", "2") // overwrites previous entry
.put("three", "3")
.buildImmutable();
assertThat(dict1.toString()).isEqualTo("{\"one\": \"1\", \"two\": \"2\", \"three\": \"3\"}");
Expand All @@ -125,19 +125,21 @@ public void testDictBuilder() throws Exception {

// builder reuse and mutability
Dict.Builder<String, String> builder = Dict.<String, String>builder().putAll(dict1);
builder.put("three", "33"); // overwrites previous entry
Mutability mu = Mutability.create("test");
Dict<String, String> dict3 = builder.build(mu);
dict3.putEntry("four", "4");
dict3.putEntry("four", "4"); // new entry
dict3.putEntry("two", "22"); // overwrites previous entry
assertThat(dict3.toString())
.isEqualTo("{\"one\": \"1\", \"two\": \"2\", \"three\": \"3\", \"four\": \"4\"}");
.isEqualTo("{\"one\": \"1\", \"two\": \"22\", \"three\": \"33\", \"four\": \"4\"}");
mu.close();
assertThrows(EvalException.class, dict1::clearEntries); // frozen
builder.put("five", "5"); // keep building
Dict<String, String> dict4 = builder.buildImmutable();
assertThat(dict4.toString())
.isEqualTo("{\"one\": \"1\", \"two\": \"2\", \"three\": \"3\", \"five\": \"5\"}");
.isEqualTo("{\"one\": \"1\", \"two\": \"2\", \"three\": \"33\", \"five\": \"5\"}");
assertThat(dict3.toString())
.isEqualTo(
"{\"one\": \"1\", \"two\": \"2\", \"three\": \"3\", \"four\": \"4\"}"); // unchanged
"{\"one\": \"1\", \"two\": \"22\", \"three\": \"33\", \"four\": \"4\"}"); // unchanged
}
}

0 comments on commit 24ec910

Please sign in to comment.