Skip to content

Commit 1c347c8

Browse files
Bennett-LynchBennett-Lynch
authored
Improve performance of generated enum methods (#2693)
* Improve performance of generated enum methods ## Motivation Generated enums use a standard `fromValue` method to resolve a string value (usually received on the wire) to the corresponding enum. The existing implementation is not optimal for a few reasons: 1. It calls `Enum#values()`, which must allocate a new array every time 2. It iterates the entire array with O(N) complexity The generated `knownValues` follows a similar pattern, plus it does not take advantage of `EnumSet`, which is typically faster than HashSet. These are minor optimization concerns, however many code paths repeatedly call `fromValue()` (as opposed to caching it). ## Modifications * For each generated enum, store a value map which maps the string value to the corresponding enum constant value. Use this map for `fromValue`. * Use an EnumSet for `knownValues`. * Create new `CollectionUtils.index()` & `EnumUtils.index()` methods to facilitate the common use case of constructing a unique index lookup map ## Result * `fromValue` changes from O(N) to O(1) * Minor increase in memory overhead from storing a static map, but reduced allocation from creating new arrays every method call
1 parent 98eac61 commit 1c347c8

File tree

10 files changed

+170
-38
lines changed

10 files changed

+170
-38
lines changed

‎.changes/next-release/feature-AWSSDKforJavav2-0b3656e.json

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"category": "AWS SDK for Java v2",
3+
"contributor": "",
4+
"type": "feature",
5+
"description": "Improve performance of generated enum methods"
6+
}

‎codegen/src/main/java/software/amazon/awssdk/codegen/poet/common/AbstractEnumClass.java

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,24 +15,30 @@
1515

1616
package software.amazon.awssdk.codegen.poet.common;
1717

18+
import static javax.lang.model.element.Modifier.FINAL;
19+
import static javax.lang.model.element.Modifier.PRIVATE;
20+
import static javax.lang.model.element.Modifier.STATIC;
1821
import static software.amazon.awssdk.codegen.poet.PoetUtils.createEnumBuilder;
1922
import static software.amazon.awssdk.codegen.poet.PoetUtils.toStringBuilder;
2023

2124
import com.squareup.javapoet.ClassName;
25+
import com.squareup.javapoet.FieldSpec;
2226
import com.squareup.javapoet.MethodSpec;
2327
import com.squareup.javapoet.ParameterizedTypeName;
2428
import com.squareup.javapoet.TypeSpec;
2529
import com.squareup.javapoet.TypeSpec.Builder;
30+
import java.util.EnumSet;
31+
import java.util.Map;
2632
import java.util.Set;
27-
import java.util.stream.Collectors;
28-
import java.util.stream.Stream;
2933
import javax.lang.model.element.Modifier;
3034
import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
3135
import software.amazon.awssdk.codegen.poet.ClassSpec;
36+
import software.amazon.awssdk.utils.internal.EnumUtils;
3237

3338
public abstract class AbstractEnumClass implements ClassSpec {
3439

3540
private static final String VALUE = "value";
41+
private static final String VALUE_MAP = "VALUE_MAP";
3642
private static final String UNKNOWN_TO_SDK_VERSION = "UNKNOWN_TO_SDK_VERSION";
3743
private final ShapeModel shape;
3844

@@ -43,6 +49,7 @@ public AbstractEnumClass(ShapeModel shape) {
4349
@Override
4450
public final TypeSpec poetSpec() {
4551
Builder enumBuilder = createEnumBuilder(className())
52+
.addField(valueMapField())
4653
.addField(String.class, VALUE, Modifier.PRIVATE, Modifier.FINAL)
4754
.addMethod(toStringBuilder().addStatement("return $T.valueOf($N)", String.class, VALUE).build())
4855
.addMethod(fromValueSpec())
@@ -67,6 +74,16 @@ protected final ShapeModel getShape() {
6774
protected abstract void addJavadoc(Builder enumBuilder);
6875

6976
protected abstract void addEnumConstants(Builder enumBuilder);
77+
78+
private FieldSpec valueMapField() {
79+
ParameterizedTypeName mapType = ParameterizedTypeName.get(ClassName.get(Map.class),
80+
ClassName.get(String.class),
81+
className());
82+
return FieldSpec.builder(mapType, VALUE_MAP)
83+
.addModifiers(PRIVATE, STATIC, FINAL)
84+
.initializer("$1T.uniqueIndex($2T.class, $2T::toString)", EnumUtils.class, className())
85+
.build();
86+
}
7087

7188
private MethodSpec createConstructor() {
7289
return MethodSpec.constructorBuilder()
@@ -88,13 +105,7 @@ private MethodSpec fromValueSpec() {
88105
.beginControlFlow("if ($N == null)", VALUE)
89106
.addStatement("return null")
90107
.endControlFlow()
91-
.addStatement("return $1T.of($2T.values())\n" +
92-
".filter(e -> e.toString().equals($3N))\n" +
93-
".findFirst()\n" +
94-
".orElse(UNKNOWN_TO_SDK_VERSION)",
95-
Stream.class,
96-
className(),
97-
VALUE)
108+
.addStatement("return $N.getOrDefault($N, $N)", VALUE_MAP, VALUE, UNKNOWN_TO_SDK_VERSION)
98109
.build();
99110
}
100111

@@ -106,10 +117,9 @@ private MethodSpec knownValuesSpec() {
106117
+ "SDK.\n"
107118
+ "This will return all known enum values except {@link #$N}.\n\n"
108119
+ "@return a {@link $T} of known {@link $T}s", UNKNOWN_TO_SDK_VERSION, Set.class, className())
109-
.addStatement("return $T.of(values()).filter(v -> v != $N).collect($T.toSet())",
110-
Stream.class,
111-
UNKNOWN_TO_SDK_VERSION,
112-
Collectors.class)
120+
.addStatement("$1T<$2T> knownValues = $3T.allOf($2T.class)", Set.class, className(), EnumSet.class)
121+
.addStatement("knownValues.remove($N)", UNKNOWN_TO_SDK_VERSION)
122+
.addStatement("return knownValues")
113123
.build();
114124
}
115125
}

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/common/test-enum-class.java

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
package software.amazon.awssdk.codegen.poet.common.model;
22

3+
import java.util.EnumSet;
4+
import java.util.Map;
35
import java.util.Set;
4-
import java.util.stream.Collectors;
5-
import java.util.stream.Stream;
66
import software.amazon.awssdk.annotations.Generated;
7+
import software.amazon.awssdk.utils.internal.EnumUtils;
78

89
/**
910
* Some comment on the class itself
@@ -16,6 +17,8 @@ public enum TestEnumClass {
1617

1718
UNKNOWN_TO_SDK_VERSION(null);
1819

20+
private static final Map<String, TestEnumClass> VALUE_MAP = EnumUtils.uniqueIndex(TestEnumClass.class, TestEnumClass::toString);
21+
1922
private final String value;
2023

2124
private TestEnumClass(String value) {
@@ -38,8 +41,7 @@ public static TestEnumClass fromValue(String value) {
3841
if (value == null) {
3942
return null;
4043
}
41-
return Stream.of(TestEnumClass.values()).filter(e -> e.toString().equals(value)).findFirst()
42-
.orElse(UNKNOWN_TO_SDK_VERSION);
44+
return VALUE_MAP.getOrDefault(value, UNKNOWN_TO_SDK_VERSION);
4345
}
4446

4547
/**
@@ -49,7 +51,9 @@ public static TestEnumClass fromValue(String value) {
4951
* @return a {@link Set} of known {@link TestEnumClass}s
5052
*/
5153
public static Set<TestEnumClass> knownValues() {
52-
return Stream.of(values()).filter(v -> v != UNKNOWN_TO_SDK_VERSION).collect(Collectors.toSet());
54+
Set<TestEnumClass> knownValues = EnumSet.allOf(TestEnumClass.class);
55+
knownValues.remove(UNKNOWN_TO_SDK_VERSION);
56+
return knownValues;
5357
}
5458
}
5559

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/eventstream.java

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
package software.amazon.awssdk.services.jsonprotocoltests.model;
22

33
import java.util.Collections;
4+
import java.util.EnumSet;
45
import java.util.List;
6+
import java.util.Map;
57
import java.util.Set;
6-
import java.util.stream.Collectors;
7-
import java.util.stream.Stream;
88
import software.amazon.awssdk.annotations.Generated;
99
import software.amazon.awssdk.annotations.SdkPublicApi;
1010
import software.amazon.awssdk.core.SdkField;
@@ -14,6 +14,7 @@
1414
import software.amazon.awssdk.services.jsonprotocoltests.model.eventstream.DefaultEventthree;
1515
import software.amazon.awssdk.services.jsonprotocoltests.model.eventstream.DefaultSecondEventOne;
1616
import software.amazon.awssdk.services.jsonprotocoltests.model.eventstream.DefaultSecondEventTwo;
17+
import software.amazon.awssdk.utils.internal.EnumUtils;
1718

1819
/**
1920
* Base interface for all event types in EventStream.
@@ -103,6 +104,8 @@ enum EventType {
103104

104105
UNKNOWN_TO_SDK_VERSION(null);
105106

107+
private static final Map<String, EventType> VALUE_MAP = EnumUtils.uniqueIndex(EventType.class, EventType::toString);
108+
106109
private final String value;
107110

108111
private EventType(String value) {
@@ -125,8 +128,7 @@ public static EventType fromValue(String value) {
125128
if (value == null) {
126129
return null;
127130
}
128-
return Stream.of(EventType.values()).filter(e -> e.toString().equals(value)).findFirst()
129-
.orElse(UNKNOWN_TO_SDK_VERSION);
131+
return VALUE_MAP.getOrDefault(value, UNKNOWN_TO_SDK_VERSION);
130132
}
131133

132134
/**
@@ -136,7 +138,9 @@ public static EventType fromValue(String value) {
136138
* @return a {@link Set} of known {@link EventType}s
137139
*/
138140
public static Set<EventType> knownValues() {
139-
return Stream.of(values()).filter(v -> v != UNKNOWN_TO_SDK_VERSION).collect(Collectors.toSet());
141+
Set<EventType> knownValues = EnumSet.allOf(EventType.class);
142+
knownValues.remove(UNKNOWN_TO_SDK_VERSION);
143+
return knownValues;
140144
}
141145
}
142146
}

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/inputeventstream.java

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
package software.amazon.awssdk.services.jsonprotocoltests.model;
22

3+
import java.util.EnumSet;
4+
import java.util.Map;
35
import java.util.Set;
4-
import java.util.stream.Collectors;
5-
import java.util.stream.Stream;
66
import software.amazon.awssdk.annotations.Generated;
77
import software.amazon.awssdk.annotations.SdkPublicApi;
88
import software.amazon.awssdk.services.jsonprotocoltests.model.inputeventstream.DefaultInputEvent;
9+
import software.amazon.awssdk.utils.internal.EnumUtils;
910

1011
/**
1112
* Base interface for all event types in InputEventStream.
@@ -36,6 +37,8 @@ enum EventType {
3637

3738
UNKNOWN_TO_SDK_VERSION(null);
3839

40+
private static final Map<String, EventType> VALUE_MAP = EnumUtils.uniqueIndex(EventType.class, EventType::toString);
41+
3942
private final String value;
4043

4144
private EventType(String value) {
@@ -58,8 +61,7 @@ public static EventType fromValue(String value) {
5861
if (value == null) {
5962
return null;
6063
}
61-
return Stream.of(EventType.values()).filter(e -> e.toString().equals(value)).findFirst()
62-
.orElse(UNKNOWN_TO_SDK_VERSION);
64+
return VALUE_MAP.getOrDefault(value, UNKNOWN_TO_SDK_VERSION);
6365
}
6466

6567
/**
@@ -69,7 +71,9 @@ public static EventType fromValue(String value) {
6971
* @return a {@link Set} of known {@link EventType}s
7072
*/
7173
public static Set<EventType> knownValues() {
72-
return Stream.of(values()).filter(v -> v != UNKNOWN_TO_SDK_VERSION).collect(Collectors.toSet());
74+
Set<EventType> knownValues = EnumSet.allOf(EventType.class);
75+
knownValues.remove(UNKNOWN_TO_SDK_VERSION);
76+
return knownValues;
7377
}
7478
}
7579
}

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/inputeventstreamtwo.java

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
package software.amazon.awssdk.services.jsonprotocoltests.model;
22

3+
import java.util.EnumSet;
4+
import java.util.Map;
35
import java.util.Set;
4-
import java.util.stream.Collectors;
5-
import java.util.stream.Stream;
66
import software.amazon.awssdk.annotations.Generated;
77
import software.amazon.awssdk.annotations.SdkPublicApi;
88
import software.amazon.awssdk.services.jsonprotocoltests.model.inputeventstreamtwo.DefaultInputEventTwo;
9+
import software.amazon.awssdk.utils.internal.EnumUtils;
910

1011
/**
1112
* Base interface for all event types in InputEventStreamTwo.
@@ -36,6 +37,8 @@ enum EventType {
3637

3738
UNKNOWN_TO_SDK_VERSION(null);
3839

40+
private static final Map<String, EventType> VALUE_MAP = EnumUtils.uniqueIndex(EventType.class, EventType::toString);
41+
3942
private final String value;
4043

4144
private EventType(String value) {
@@ -58,8 +61,7 @@ public static EventType fromValue(String value) {
5861
if (value == null) {
5962
return null;
6063
}
61-
return Stream.of(EventType.values()).filter(e -> e.toString().equals(value)).findFirst()
62-
.orElse(UNKNOWN_TO_SDK_VERSION);
64+
return VALUE_MAP.getOrDefault(value, UNKNOWN_TO_SDK_VERSION);
6365
}
6466

6567
/**
@@ -69,7 +71,9 @@ public static EventType fromValue(String value) {
6971
* @return a {@link Set} of known {@link EventType}s
7072
*/
7173
public static Set<EventType> knownValues() {
72-
return Stream.of(values()).filter(v -> v != UNKNOWN_TO_SDK_VERSION).collect(Collectors.toSet());
74+
Set<EventType> knownValues = EnumSet.allOf(EventType.class);
75+
knownValues.remove(UNKNOWN_TO_SDK_VERSION);
76+
return knownValues;
7377
}
7478
}
7579
}

‎codegen/src/test/resources/software/amazon/awssdk/codegen/poet/model/sharedstream/eventstream.java

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,16 @@
11
package software.amazon.awssdk.services.sharedeventstream.model;
22

33
import java.util.Collections;
4+
import java.util.EnumSet;
45
import java.util.List;
6+
import java.util.Map;
57
import java.util.Set;
6-
import java.util.stream.Collectors;
7-
import java.util.stream.Stream;
88
import software.amazon.awssdk.annotations.Generated;
99
import software.amazon.awssdk.annotations.SdkPublicApi;
1010
import software.amazon.awssdk.core.SdkField;
1111
import software.amazon.awssdk.core.SdkPojo;
1212
import software.amazon.awssdk.services.sharedeventstream.model.eventstream.DefaultPerson;
13+
import software.amazon.awssdk.utils.internal.EnumUtils;
1314

1415
/**
1516
* Base interface for all event types in EventStream.
@@ -76,6 +77,8 @@ enum EventType {
7677

7778
UNKNOWN_TO_SDK_VERSION(null);
7879

80+
private static final Map<String, EventType> VALUE_MAP = EnumUtils.uniqueIndex(EventType.class, EventType::toString);
81+
7982
private final String value;
8083

8184
private EventType(String value) {
@@ -98,8 +101,7 @@ public static EventType fromValue(String value) {
98101
if (value == null) {
99102
return null;
100103
}
101-
return Stream.of(EventType.values()).filter(e -> e.toString().equals(value)).findFirst()
102-
.orElse(UNKNOWN_TO_SDK_VERSION);
104+
return VALUE_MAP.getOrDefault(value, UNKNOWN_TO_SDK_VERSION);
103105
}
104106

105107
/**
@@ -109,7 +111,9 @@ public static EventType fromValue(String value) {
109111
* @return a {@link Set} of known {@link EventType}s
110112
*/
111113
public static Set<EventType> knownValues() {
112-
return Stream.of(values()).filter(v -> v != UNKNOWN_TO_SDK_VERSION).collect(Collectors.toSet());
114+
Set<EventType> knownValues = EnumSet.allOf(EventType.class);
115+
knownValues.remove(UNKNOWN_TO_SDK_VERSION);
116+
return knownValues;
113117
}
114118
}
115119
}

‎utils/src/main/java/software/amazon/awssdk/utils/CollectionUtils.java

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import java.util.ArrayList;
1919
import java.util.Collection;
20+
import java.util.HashMap;
2021
import java.util.LinkedHashMap;
2122
import java.util.LinkedList;
2223
import java.util.List;
@@ -139,4 +140,26 @@ public static <K, VInT, VOutT> Map<K, VOutT> mapValues(Map<K, VInT> inputMap, Fu
139140
public static <K, V> Map<K, V> inverseMap(Map<V, K> inputMap) {
140141
return inputMap.entrySet().stream().collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
141142
}
143+
144+
/**
145+
* For a collection of values of type {@code V} that can all be converted to type {@code K}, create a map that
146+
* indexes all of the values by {@code K}. This requires that no two values map to the same index.
147+
*
148+
* @param values the collection of values to index
149+
* @param indexFunction the function used to convert a value to its index
150+
* @param <K> the index (or key) type
151+
* @param <V> the value type
152+
* @return a (modifiable) map that indexes K to its unique value V
153+
* @throws IllegalArgumentException if any of the values map to the same index
154+
*/
155+
public static <K, V> Map<K, V> uniqueIndex(Iterable<V> values, Function<? super V, K> indexFunction) {
156+
Map<K, V> map = new HashMap<>();
157+
for (V value : values) {
158+
K index = indexFunction.apply(value);
159+
V prev = map.put(index, value);
160+
Validate.isNull(prev, "No duplicate indices allowed but both %s and %s have the same index: %s",
161+
prev, value, index);
162+
}
163+
return map;
164+
}
142165
}

0 commit comments

Comments
 (0)