2323import javax .lang .model .type .TypeMirror ;
2424import javax .lang .model .util .Elements ;
2525
26- import static org .elasticsearch .compute .gen .Methods .appendMethod ;
2726import static org .elasticsearch .compute .gen .Methods .buildFromFactory ;
2827import static org .elasticsearch .compute .gen .Methods .getMethod ;
2928import static org .elasticsearch .compute .gen .Types .ABSTRACT_CONVERT_FUNCTION_EVALUATOR ;
4140public class ConvertEvaluatorImplementer {
4241
4342 private final TypeElement declarationType ;
44- private final ExecutableElement processFunction ;
43+ private final EvaluatorImplementer . ProcessFunction processFunction ;
4544 private final String extraName ;
4645 private final ClassName implementation ;
4746 private final TypeName argumentType ;
48- private final TypeName resultType ;
4947 private final List <TypeMirror > warnExceptions ;
5048
5149 public ConvertEvaluatorImplementer (
5250 Elements elements ,
51+ javax .lang .model .util .Types types ,
5352 ExecutableElement processFunction ,
5453 String extraName ,
5554 List <TypeMirror > warnExceptions
5655 ) {
5756 this .declarationType = (TypeElement ) processFunction .getEnclosingElement ();
58- this .processFunction = processFunction ;
59- if (processFunction .getParameters ().size () != 1 ) {
60- throw new IllegalArgumentException ("processing function should have exactly one parameter" );
57+ this .processFunction = new EvaluatorImplementer .ProcessFunction (types , processFunction , warnExceptions );
58+
59+ if (this .processFunction .args .getFirst () instanceof EvaluatorImplementer .StandardProcessFunctionArg == false ) {
60+ throw new IllegalArgumentException ("first argument must be the field to process" );
61+ }
62+ for (int a = 1 ; a < this .processFunction .args .size (); a ++) {
63+ if (this .processFunction .args .get (a ) instanceof EvaluatorImplementer .FixedProcessFunctionArg == false ) {
64+ throw new IllegalArgumentException ("fixed function args supported after the first" );
65+ // TODO support more function types when we need them
66+ }
6167 }
68+
6269 this .extraName = extraName ;
6370 this .argumentType = TypeName .get (processFunction .getParameters ().get (0 ).asType ());
64- this .resultType = TypeName .get (processFunction .getReturnType ());
6571 this .warnExceptions = warnExceptions ;
6672
6773 this .implementation = ClassName .get (
@@ -87,29 +93,36 @@ private TypeSpec type() {
8793 builder .addModifiers (Modifier .PUBLIC , Modifier .FINAL );
8894 builder .superclass (ABSTRACT_CONVERT_FUNCTION_EVALUATOR );
8995
96+ for (EvaluatorImplementer .ProcessFunctionArg a : processFunction .args ) {
97+ a .declareField (builder );
98+ }
9099 builder .addMethod (ctor ());
91- builder .addMethod (name ());
100+ builder .addMethod (next ());
92101 builder .addMethod (evalVector ());
93102 builder .addMethod (evalValue (true ));
94103 builder .addMethod (evalBlock ());
95104 builder .addMethod (evalValue (false ));
105+ builder .addMethod (processFunction .toStringMethod (implementation ));
106+ builder .addMethod (processFunction .close ());
96107 builder .addType (factory ());
97108 return builder .build ();
98109 }
99110
100111 private MethodSpec ctor () {
101112 MethodSpec .Builder builder = MethodSpec .constructorBuilder ().addModifiers (Modifier .PUBLIC );
102- builder .addParameter (EXPRESSION_EVALUATOR , "field" );
103113 builder .addParameter (SOURCE , "source" );
114+ builder .addStatement ("super(driverContext, source)" );
115+ for (EvaluatorImplementer .ProcessFunctionArg a : processFunction .args ) {
116+ a .implementCtor (builder );
117+ }
104118 builder .addParameter (DRIVER_CONTEXT , "driverContext" );
105- builder .addStatement ("super(driverContext, field, source)" );
106119 return builder .build ();
107120 }
108121
109- private MethodSpec name () {
110- MethodSpec .Builder builder = MethodSpec .methodBuilder ("name" ).addModifiers (Modifier .PUBLIC );
111- builder .addAnnotation ( Override . class ). returns (String . class );
112- builder .addStatement ("return $S " , declarationType . getSimpleName () + extraName );
122+ private MethodSpec next () {
123+ MethodSpec .Builder builder = MethodSpec .methodBuilder ("next" ). addAnnotation ( Override . class ).addModifiers (Modifier .PUBLIC );
124+ builder .returns (EXPRESSION_EVALUATOR );
125+ builder .addStatement ("return $N " , (( EvaluatorImplementer . StandardProcessFunctionArg ) processFunction . args . getFirst ()). name () );
113126 return builder .build ();
114127 }
115128
@@ -129,7 +142,7 @@ private MethodSpec evalVector() {
129142 builder .beginControlFlow ("if (vector.isConstant())" );
130143 {
131144 catchingWarnExceptions (builder , () -> {
132- var constVectType = blockType ( resultType );
145+ var constVectType = processFunction . resultDataType ( true );
133146 builder .addStatement (
134147 "return driverContext.blockFactory().newConstant$TWith($N, positionCount)" ,
135148 constVectType ,
@@ -139,7 +152,7 @@ private MethodSpec evalVector() {
139152 }
140153 builder .endControlFlow ();
141154
142- ClassName resultBuilderType = builderType (blockType ( resultType ));
155+ ClassName resultBuilderType = builderType (processFunction . resultDataType ( true ));
143156 builder .beginControlFlow (
144157 "try ($T builder = driverContext.blockFactory().$L(positionCount))" ,
145158 resultBuilderType ,
@@ -150,7 +163,11 @@ private MethodSpec evalVector() {
150163 {
151164 catchingWarnExceptions (
152165 builder ,
153- () -> builder .addStatement ("builder.$L($N)" , appendMethod (resultType ), evalValueCall ("vector" , "p" , scratchPadName )),
166+ () -> builder .addStatement (
167+ "builder.$L($N)" ,
168+ processFunction .appendMethod (),
169+ evalValueCall ("vector" , "p" , scratchPadName )
170+ ),
154171 () -> builder .addStatement ("builder.appendNull()" )
155172 );
156173 }
@@ -185,7 +202,7 @@ private MethodSpec evalBlock() {
185202 TypeName blockType = blockType (argumentType );
186203 builder .addStatement ("$T block = ($T) b" , blockType , blockType );
187204 builder .addStatement ("int positionCount = block.getPositionCount()" );
188- TypeName resultBuilderType = builderType (blockType ( resultType ));
205+ TypeName resultBuilderType = builderType (processFunction . resultDataType ( true ));
189206 builder .beginControlFlow (
190207 "try ($T builder = driverContext.blockFactory().$L(positionCount))" ,
191208 resultBuilderType ,
@@ -196,19 +213,18 @@ private MethodSpec evalBlock() {
196213 builder .addStatement ("BytesRef $N = new BytesRef()" , scratchPadName );
197214 }
198215
199- String appendMethod = appendMethod (resultType );
216+ String appendMethod = processFunction . appendMethod ();
200217 builder .beginControlFlow ("for (int p = 0; p < positionCount; p++)" );
201218 {
202219 builder .addStatement ("int valueCount = block.getValueCount(p)" );
203220 builder .addStatement ("int start = block.getFirstValueIndex(p)" );
204221 builder .addStatement ("int end = start + valueCount" );
205222 builder .addStatement ("boolean positionOpened = false" );
206223 builder .addStatement ("boolean valuesAppended = false" );
207- // builder.addStatement("builder.beginPositionEntry()");
208224 builder .beginControlFlow ("for (int i = start; i < end; i++)" );
209225 {
210226 catchingWarnExceptions (builder , () -> {
211- builder .addStatement ("$T value = $N" , resultType , evalValueCall ("block" , "i" , scratchPadName ));
227+ builder .addStatement ("$T value = $N" , processFunction . returnType () , evalValueCall ("block" , "i" , scratchPadName ));
212228 builder .beginControlFlow ("if (positionOpened == false && valueCount > 1)" );
213229 {
214230 builder .addStatement ("builder.beginPositionEntry()" );
@@ -253,8 +269,8 @@ private String evalValueCall(String container, String index, String scratchPad)
253269
254270 private MethodSpec evalValue (boolean forVector ) {
255271 MethodSpec .Builder builder = MethodSpec .methodBuilder ("evalValue" )
256- .addModifiers (Modifier .PRIVATE , Modifier . STATIC )
257- .returns (resultType );
272+ .addModifiers (Modifier .PRIVATE )
273+ .returns (processFunction . returnType () );
258274
259275 if (forVector ) {
260276 builder .addParameter (vectorType (argumentType ), "container" );
@@ -269,8 +285,17 @@ private MethodSpec evalValue(boolean forVector) {
269285 builder .addStatement ("$T value = container.$N(index)" , argumentType , getMethod (argumentType ));
270286 }
271287
272- builder .addStatement ("return $T.$N(value)" , declarationType , processFunction .getSimpleName ());
273-
288+ StringBuilder pattern = new StringBuilder ();
289+ List <Object > args = new ArrayList <>();
290+ pattern .append ("return $T.$N(value" );
291+ args .add (declarationType );
292+ args .add (processFunction .function .getSimpleName ());
293+ for (int a = 1 ; a < processFunction .args .size (); a ++) {
294+ pattern .append (", " );
295+ processFunction .args .get (a ).buildInvocation (pattern , args , false /* block style parameter should be unused */ );
296+ }
297+ pattern .append (")" );
298+ builder .addStatement (pattern .toString (), args .toArray ());
274299 return builder .build ();
275300 }
276301
@@ -280,42 +305,11 @@ private TypeSpec factory() {
280305 builder .addModifiers (Modifier .PUBLIC , Modifier .STATIC );
281306
282307 builder .addField (SOURCE , "source" , Modifier .PRIVATE , Modifier .FINAL );
283- builder .addField (EXPRESSION_EVALUATOR_FACTORY , "field" , Modifier .PRIVATE , Modifier .FINAL );
284-
285- builder .addMethod (factoryCtor ());
286- builder .addMethod (factoryGet ());
287- builder .addMethod (factoryToString ());
288- return builder .build ();
289- }
290-
291- private MethodSpec factoryCtor () {
292- MethodSpec .Builder builder = MethodSpec .constructorBuilder ().addModifiers (Modifier .PUBLIC );
293- builder .addParameter (EXPRESSION_EVALUATOR_FACTORY , "field" );
294- builder .addParameter (SOURCE , "source" );
295- builder .addStatement ("this.field = field" );
296- builder .addStatement ("this.source = source" );
297- return builder .build ();
298- }
299-
300- private MethodSpec factoryGet () {
301- MethodSpec .Builder builder = MethodSpec .methodBuilder ("get" ).addAnnotation (Override .class );
302- builder .addModifiers (Modifier .PUBLIC );
303- builder .addParameter (DRIVER_CONTEXT , "context" );
304- builder .returns (implementation );
305-
306- List <String > args = new ArrayList <>();
307- args .add ("field.get(context)" );
308- args .add ("source" );
309- args .add ("context" );
310- builder .addStatement ("return new $T($L)" , implementation , args .stream ().collect (Collectors .joining (", " )));
311- return builder .build ();
312- }
308+ processFunction .args .forEach (a -> a .declareFactoryField (builder ));
313309
314- private MethodSpec factoryToString () {
315- MethodSpec .Builder builder = MethodSpec .methodBuilder ("toString" ).addAnnotation (Override .class );
316- builder .addModifiers (Modifier .PUBLIC );
317- builder .returns (String .class );
318- builder .addStatement ("return $S + field + $S" , declarationType .getSimpleName () + extraName + "Evaluator[field=" , "]" );
310+ builder .addMethod (processFunction .factoryCtor ());
311+ builder .addMethod (processFunction .factoryGet (implementation ));
312+ builder .addMethod (processFunction .toStringMethod (implementation ));
319313 return builder .build ();
320314 }
321315}
0 commit comments