Skip to content

Commit 8a0447a

Browse files
mp911demp911de
committed
Experiment with better code blocks.
1 parent 9536eb7 commit 8a0447a

File tree

3 files changed

+297
-47
lines changed

3 files changed

+297
-47
lines changed

‎spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/aot/BetterCodeBlock.java‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
/*
2+
* Copyright 2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.jdbc.repository.aot;
17+
18+
import java.util.ArrayList;
19+
import java.util.List;
20+
import java.util.Map;
21+
import java.util.function.Consumer;
22+
23+
import org.springframework.javapoet.CodeBlock;
24+
25+
/**
26+
* @author Mark Paluch
27+
*/
28+
class BetterCodeBlock {
29+
30+
public static BetterBuilder of(CodeBlock.Builder builder) {
31+
return new BetterBuilder(builder);
32+
}
33+
34+
public static class BetterBuilder {
35+
36+
private final CodeBlock.Builder builder;
37+
38+
BetterBuilder(CodeBlock.Builder builder) {
39+
this.builder = builder;
40+
}
41+
42+
public boolean isEmpty() {
43+
return builder.isEmpty();
44+
}
45+
46+
public BetterBuilder indent() {
47+
builder.indent();
48+
return this;
49+
}
50+
51+
public BetterBuilder endControlFlow() {
52+
builder.endControlFlow();
53+
return this;
54+
}
55+
56+
public BetterBuilder nextControlFlow(String controlFlow, Object... args) {
57+
builder.nextControlFlow(controlFlow, args);
58+
return this;
59+
}
60+
61+
public BetterBuilder add(String format, Object... args) {
62+
builder.add(format, args);
63+
return this;
64+
}
65+
66+
public BetterBuilder addStatement(CodeBlock codeBlock) {
67+
builder.addStatement(codeBlock);
68+
return this;
69+
}
70+
71+
public BetterBuilder addStatement(Consumer<StepBuilder> consumer) {
72+
73+
StepBuilder statementBuilder = new StepBuilder();
74+
consumer.accept(statementBuilder);
75+
76+
if (!statementBuilder.isEmpty()) {
77+
78+
this.add("$[");
79+
80+
for (int i = 0; i < statementBuilder.formats.size(); i++) {
81+
builder.add(statementBuilder.formats.get(i), statementBuilder.args.get(i));
82+
}
83+
84+
this.add(";\n$]");
85+
86+
}
87+
return this;
88+
}
89+
90+
public BetterBuilder beginControlFlow(String controlFlow, Object... args) {
91+
builder.beginControlFlow(controlFlow, args);
92+
return this;
93+
}
94+
95+
public BetterBuilder add(CodeBlock codeBlock) {
96+
builder.add(codeBlock);
97+
return this;
98+
}
99+
100+
public BetterBuilder addStatement(String format, Object... args) {
101+
builder.addStatement(format, args);
102+
return this;
103+
}
104+
105+
public CodeBlock build() {
106+
return builder.build();
107+
}
108+
109+
public BetterBuilder clear() {
110+
builder.clear();
111+
return this;
112+
}
113+
114+
public BetterBuilder addNamed(String format, Map<String, ?> arguments) {
115+
builder.addNamed(format, arguments);
116+
return this;
117+
}
118+
119+
public BetterBuilder endControlFlow(String controlFlow, Object... args) {
120+
builder.endControlFlow(controlFlow, args);
121+
return this;
122+
}
123+
124+
public BetterBuilder unindent() {
125+
builder.unindent();
126+
return this;
127+
}
128+
}
129+
130+
public static class StepBuilder {
131+
132+
List<String> formats = new ArrayList<>();
133+
List<Object[]> args = new ArrayList<>();
134+
135+
public boolean isEmpty() {
136+
return formats.isEmpty();
137+
}
138+
139+
public ConditionalStatementStep when(boolean state) {
140+
return whenNot(!state);
141+
}
142+
143+
public StepBuilder add(String format, Object... args) {
144+
formats.add(format);
145+
this.args.add(args);
146+
return this;
147+
}
148+
149+
public ConditionalStatementStep whenNot(boolean state) {
150+
151+
return (format, args) -> {
152+
153+
if (state) {
154+
formats.add(format);
155+
this.args.add(args);
156+
}
157+
return this;
158+
};
159+
}
160+
161+
public interface ConditionalStatementStep {
162+
163+
StepBuilder then(String format, Object... args);
164+
}
165+
}
166+
}

‎spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/aot/JdbcCodeBlocks.java‎

Lines changed: 26 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
5959
import org.springframework.jdbc.core.namedparam.SqlParameterSource;
6060
import org.springframework.util.Assert;
61+
import org.springframework.util.ClassUtils;
6162
import org.springframework.util.ObjectUtils;
6263
import org.springframework.util.StringUtils;
6364

@@ -697,26 +698,19 @@ private CodeBlock update(Builder builder, Class<?> returnType) {
697698

698699
String result = context.localVariable("result");
699700

700-
builder.add("$[");
701+
BetterCodeBlock.BetterBuilder cb = BetterCodeBlock.of(builder);
701702

702-
if (!ReflectionUtils.isVoid(returnType)) {
703-
builder.add("int $L = ", result);
704-
}
703+
cb.addStatement(it -> {
704+
it.when(ReflectionUtils.isVoid(returnType)).then("int $L = ", result);
705705

706-
builder.add("getJdbcOperations().update($L, $L)", queryVariableName, parameterSourceVariableName);
707-
builder.add(";\n$]");
706+
it.add("getJdbcOperations().update($L, $L)", queryVariableName, parameterSourceVariableName);
707+
});
708708

709-
if (returnType == boolean.class || returnType == Boolean.class) {
710-
builder.addStatement("return $L != 0", result);
711-
} else if (returnType == Long.class) {
712-
builder.addStatement("return (long) $L", result);
713-
} else if (ReflectionUtils.isVoid(returnType)) {
714-
if (returnType == Void.class) {
715-
builder.addStatement("return null");
716-
}
717-
} else {
718-
builder.addStatement("return $L", result);
719-
}
709+
builder.add(ReturnStatement.returning(returnType) //
710+
.whenBoolean("$L != 0", result) //
711+
.whenBoxedLong("(long) $L", result) //
712+
.otherwise("$L", result)//
713+
.build());
720714

721715
return builder.build();
722716
}
@@ -732,22 +726,15 @@ private CodeBlock delete(Builder builder, String rowMapper, String result, TypeN
732726

733727
builder.addStatement("$L.forEach(getOperations()::delete)", result);
734728

735-
if (Collection.class.isAssignableFrom(context.getReturnType().toClass())) {
736-
builder.addStatement("return ($T) convertMany($L, $T.class)", context.getReturnTypeName(), result,
737-
queryResultType);
738-
} else if (returnType == context.getRepositoryInformation().getDomainType()) {
739-
builder.addStatement("return ($1T) ($2L.isEmpty() ? null : $2L.iterator().next())", actualReturnType, result);
740-
} else if (returnType == boolean.class || returnType == Boolean.class) {
741-
builder.addStatement("return !$L.isEmpty()", result);
742-
} else if (returnType == Long.class) {
743-
builder.addStatement("return (long) $L.size()", result);
744-
} else if (ReflectionUtils.isVoid(returnType)) {
745-
if (returnType == Void.class) {
746-
builder.addStatement("return null");
747-
}
748-
} else {
749-
builder.addStatement("return $L.size()", result);
750-
}
729+
builder.add(ReturnStatement.returning(returnType) //
730+
.when(Collection.class.isAssignableFrom(context.getReturnType().toClass()), "($T) convertMany($L, $T.class)",
731+
context.getReturnTypeName(), result, queryResultType) //
732+
.when(context.getRepositoryInformation().getDomainType(),
733+
"($1T) ($2L.isEmpty() ? null : $2L.iterator().next())", actualReturnType, result) //
734+
.whenBoolean("!$L.isEmpty()", result) //
735+
.whenBoxedLong("(long) $L.size()", result) //
736+
.otherwise("$L.size()", result) //
737+
.build());
751738

752739
return builder.build();
753740
}
@@ -757,18 +744,10 @@ private CodeBlock count(Builder builder, String result, Class<?> returnType, Typ
757744
builder.addStatement("$1T $2L = queryForObject($3L, $4L, new $5T<>($1T.class))", Number.class, result,
758745
queryVariableName, parameterSourceVariableName, SingleColumnRowMapper.class);
759746

760-
if (returnType == Long.class) {
761-
builder.addStatement("return $1L != null ? $1L.longValue() : null", result);
762-
} else if (returnType == Integer.class) {
763-
builder.addStatement("return $1L != null ? $1L.intValue() : null", result);
764-
} else if (returnType == Long.TYPE) {
765-
builder.addStatement("return $1L != null ? $1L.longValue() : 0L", result);
766-
} else if (returnType == Integer.TYPE) {
767-
builder.addStatement("return $1L != null ? $1L.intValue() : 0", result);
768-
} else {
769-
builder.addStatement("return ($T) convertOne($L, $T.class)", context.getReturnTypeName(), result,
770-
queryResultType);
771-
}
747+
builder.add(ReturnStatement.returning(returnType) //
748+
.number(result) //
749+
.otherwise("($T) convertOne($L, $T.class)", context.getReturnTypeName(), result, queryResultType) //
750+
.build());
772751

773752
return builder.build();
774753
}
@@ -783,8 +762,8 @@ private CodeBlock exists(Builder builder, TypeName queryResultType) {
783762

784763
public static boolean returnsModifying(Class<?> returnType) {
785764

786-
return returnType == int.class || returnType == long.class || returnType == Integer.class
787-
|| returnType == Long.class;
765+
return ClassUtils.resolvePrimitiveIfNecessary(returnType) == Integer.class
766+
|| ClassUtils.resolvePrimitiveIfNecessary(returnType) == Long.class;
788767
}
789768

790769
}

‎spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/aot/ReturnStatement.java‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
/*
2+
* Copyright 2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.jdbc.repository.aot;
17+
18+
import org.springframework.data.util.ReflectionUtils;
19+
import org.springframework.javapoet.CodeBlock;
20+
import org.springframework.util.ClassUtils;
21+
22+
/**
23+
* @author Mark Paluch
24+
*/
25+
class ReturnStatement {
26+
27+
public static ReturnStatementBuilder returning(Class<?> returnType) {
28+
return new ReturnStatementBuilder(returnType);
29+
}
30+
31+
static class ReturnStatementBuilder {
32+
33+
private final Class<?> returnType;
34+
private final CodeBlock.Builder builder = CodeBlock.builder();
35+
private boolean hasReturn = false;
36+
37+
private ReturnStatementBuilder(Class<?> returnType) {
38+
39+
this.returnType = returnType;
40+
whenBoxed(Void.class, "null");
41+
hasReturn = ReflectionUtils.isVoid(returnType);
42+
}
43+
44+
public ReturnStatementBuilder whenBoolean(String format, Object... args) {
45+
return when(returnType == boolean.class || returnType == Boolean.class, format, args);
46+
}
47+
48+
public ReturnStatementBuilder whenBoxedLong(String format, Object... args) {
49+
return whenBoxed(long.class, format, args);
50+
}
51+
52+
public ReturnStatementBuilder whenLong(String format, Object... args) {
53+
return when(returnType == long.class, format, args);
54+
}
55+
56+
public ReturnStatementBuilder whenBoxedInteger(String format, Object... args) {
57+
return whenBoxed(int.class, format, args);
58+
}
59+
60+
public ReturnStatementBuilder whenInt(String format, Object... args) {
61+
return when(returnType == int.class, format, args);
62+
}
63+
64+
public ReturnStatementBuilder whenBoxed(Class<?> primitiveOrWrapper, String format, Object... args) {
65+
66+
Class<?> primitiveWrapper = ClassUtils.resolvePrimitiveIfNecessary(primitiveOrWrapper);
67+
return when(returnType == primitiveWrapper, format, args);
68+
}
69+
70+
public ReturnStatementBuilder when(Class<?> returnType, String format, Object... args) {
71+
return when(this.returnType.isAssignableFrom(returnType), format, args);
72+
}
73+
74+
public ReturnStatementBuilder when(boolean condition, String format, Object... args) {
75+
76+
if (hasReturn) {
77+
return this;
78+
}
79+
80+
if (condition) {
81+
builder.addStatement("return " + format, args);
82+
hasReturn = true;
83+
}
84+
85+
return this;
86+
}
87+
88+
public ReturnStatementBuilder otherwise(String format, Object... args) {
89+
return when(!hasReturn, format, args);
90+
}
91+
92+
public ReturnStatementBuilder number(String resultVariableName) {
93+
94+
return whenBoxedLong("$1L != null ? $1L.longValue() : null", resultVariableName) //
95+
.whenLong("$1L != null ? $1L.longValue() : 0L", resultVariableName) //
96+
.whenBoxedInteger("$1L != null ? $1L.intValue() : null", resultVariableName) //
97+
.whenInt("$1L != null ? $1L.intValue() : 0", resultVariableName);
98+
}
99+
100+
public CodeBlock build() {
101+
return builder.build();
102+
}
103+
104+
}
105+
}

0 commit comments

Comments
 (0)