/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.sdk.schemas.utils;

import static org.apache.beam.sdk.util.ByteBuddyUtils.getClassLoadingStrategy;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull;

import com.google.auto.value.AutoValue;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Modifier;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.asm.AsmVisitorWrapper;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.modifier.FieldManifestation;
import net.bytebuddy.description.modifier.Visibility;
import net.bytebuddy.description.type.TypeDescription.ForLoadedType;
import net.bytebuddy.description.type.TypeDescription.Generic;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.dynamic.scaffold.InstrumentedType;
import net.bytebuddy.implementation.Implementation;
import net.bytebuddy.implementation.Implementation.Context;
import net.bytebuddy.implementation.bytecode.ByteCodeAppender;
import net.bytebuddy.implementation.bytecode.ByteCodeAppender.Size;
import net.bytebuddy.implementation.bytecode.Duplication;
import net.bytebuddy.implementation.bytecode.Removal;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.implementation.bytecode.StackSize;
import net.bytebuddy.implementation.bytecode.assign.TypeCasting;
import net.bytebuddy.implementation.bytecode.collection.ArrayAccess;
import net.bytebuddy.implementation.bytecode.constant.IntegerConstant;
import net.bytebuddy.implementation.bytecode.constant.NullConstant;
import net.bytebuddy.implementation.bytecode.member.FieldAccess;
import net.bytebuddy.implementation.bytecode.member.MethodInvocation;
import net.bytebuddy.implementation.bytecode.member.MethodReturn;
import net.bytebuddy.implementation.bytecode.member.MethodVariableAccess;
import net.bytebuddy.jar.asm.ClassWriter;
import net.bytebuddy.jar.asm.Label;
import net.bytebuddy.jar.asm.MethodVisitor;
import net.bytebuddy.jar.asm.Opcodes;
import net.bytebuddy.matcher.ElementMatchers;
import org.apache.beam.sdk.schemas.FieldAccessDescriptor;
import org.apache.beam.sdk.schemas.FieldAccessDescriptor.FieldDescriptor;
import org.apache.beam.sdk.schemas.FieldAccessDescriptor.FieldDescriptor.Qualifier;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.Schema.FieldType;
import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.IfNullElse;
import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.ShortCircuitReturnNull;
import org.apache.beam.sdk.util.common.ReflectHelpers;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;

@SuppressWarnings({
  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
})
class SelectByteBuddyHelpers {
  private static final ByteBuddy BYTE_BUDDY = new ByteBuddy();
  private static final String SELECT_SCHEMA_FIELD_NAME = "OUTPUTSCHEMA";

  private static final ForLoadedType ROW_LOADED_TYPE = new ForLoadedType(Row.class);
  private static final ForLoadedType LIST_LOADED_TYPE = new ForLoadedType(List.class);
  private static final ForLoadedType LISTS_LOADED_TYPE = new ForLoadedType(Lists.class);
  private static final ForLoadedType MAP_LOADED_TYPE = new ForLoadedType(Map.class);
  private static final ForLoadedType MAPS_LOADED_TYPE = new ForLoadedType(Maps.class);
  private static final ForLoadedType MAPENTRY_LOADED_TYPE = new ForLoadedType(Map.Entry.class);
  private static final ForLoadedType ITERABLE_LOADED_TYPE = new ForLoadedType(Iterable.class);
  private static final ForLoadedType ITERATOR_LOADED_TYPE = new ForLoadedType(Iterator.class);

  private static final MethodDescription LIST_ADD =
      LIST_LOADED_TYPE
          .getDeclaredMethods()
          .filter(ElementMatchers.named("add").and(ElementMatchers.takesArguments(1)))
          .getOnly();

  private static final MethodDescription LISTS_NEW_ARRAYLIST =
      LISTS_LOADED_TYPE
          .getDeclaredMethods()
          .filter(ElementMatchers.named("newArrayList").and(ElementMatchers.takesArguments(0)))
          .getOnly();

  private static final MethodDescription MAP_ENTRYSET =
      MAP_LOADED_TYPE.getDeclaredMethods().filter(ElementMatchers.named("entrySet")).getOnly();

  private static final MethodDescription MAP_PUT =
      MAP_LOADED_TYPE
          .getDeclaredMethods()
          .filter(ElementMatchers.named("put").and(ElementMatchers.takesArguments(2)))
          .getOnly();

  private static final MethodDescription MAPS_NEW_HASHMAP =
      MAPS_LOADED_TYPE
          .getDeclaredMethods()
          .filter(ElementMatchers.named("newHashMap").and(ElementMatchers.takesArguments(0)))
          .getOnly();

  private static final MethodDescription ITERABLE_ITERATOR =
      ITERABLE_LOADED_TYPE.getDeclaredMethods().filter(ElementMatchers.named("iterator")).getOnly();

  private static final MethodDescription ITERATOR_HASNEXT =
      ITERATOR_LOADED_TYPE.getDeclaredMethods().filter(ElementMatchers.named("hasNext")).getOnly();

  private static final MethodDescription ITERATOR_NEXT =
      ITERATOR_LOADED_TYPE.getDeclaredMethods().filter(ElementMatchers.named("next")).getOnly();

  private static final MethodDescription MAPENTRY_GETKEY =
      MAPENTRY_LOADED_TYPE.getDeclaredMethods().filter(ElementMatchers.named("getKey")).getOnly();

  private static final MethodDescription MAPENTRY_GETVALUE =
      MAPENTRY_LOADED_TYPE.getDeclaredMethods().filter(ElementMatchers.named("getValue")).getOnly();

  @AutoValue
  abstract static class SchemaAndDescriptor {
    abstract Schema getSchema();

    abstract FieldAccessDescriptor getFieldAccessDecriptor();

    static SchemaAndDescriptor of(Schema schema, FieldAccessDescriptor fieldAccessDescriptor) {
      return new AutoValue_SelectByteBuddyHelpers_SchemaAndDescriptor(
          schema, fieldAccessDescriptor);
    }
  }

  private static final Map<SchemaAndDescriptor, RowSelector> CACHED_SELECTORS =
      Maps.newConcurrentMap();

  static RowSelector getRowSelector(
      Schema inputSchema, FieldAccessDescriptor fieldAccessDescriptor) {
    return CACHED_SELECTORS.computeIfAbsent(
        SchemaAndDescriptor.of(inputSchema, fieldAccessDescriptor),
        SelectByteBuddyHelpers::createRowSelector);
  }

  static RowSelector createRowSelector(SchemaAndDescriptor schemaAndDescriptor) {
    Schema outputSchema =
        SelectHelpers.getOutputSchema(
            schemaAndDescriptor.getSchema(), schemaAndDescriptor.getFieldAccessDecriptor());
    try {
      DynamicType.Builder<RowSelector> builder =
          BYTE_BUDDY
              .subclass(RowSelector.class)
              .method(ElementMatchers.named("select"))
              .intercept(
                  new SelectInstruction(
                      schemaAndDescriptor.getFieldAccessDecriptor(),
                      schemaAndDescriptor.getSchema(),
                      outputSchema))
              .defineField(
                  SELECT_SCHEMA_FIELD_NAME,
                  Schema.class,
                  Visibility.PRIVATE,
                  FieldManifestation.FINAL)
              .defineConstructor(Modifier.PUBLIC)
              .withParameters(Schema.class)
              .intercept(new SelectInstructionConstructor());

      return builder
          .visit(new AsmVisitorWrapper.ForDeclaredMethods().writerFlags(ClassWriter.COMPUTE_FRAMES))
          .make()
          .load(
              ReflectHelpers.findClassLoader(Row.class.getClassLoader()),
              getClassLoadingStrategy(RowSelector.class))
          .getLoaded()
          .getDeclaredConstructor(Schema.class)
          .newInstance(outputSchema);
    } catch (InstantiationException
        | IllegalAccessException
        | NoSuchMethodException
        | InvocationTargetException e) {
      throw new RuntimeException("Unable to generate");
    }
  }

  private static class SelectInstructionConstructor implements Implementation {
    @Override
    public InstrumentedType prepare(InstrumentedType instrumentedType) {
      return instrumentedType;
    }

    @Override
    public ByteCodeAppender appender(final Target implementationTarget) {
      return (methodVisitor, implementationContext, instrumentedMethod) -> {
        int numLocals = 1 + instrumentedMethod.getParameters().size();
        StackManipulation stackManipulation =
            new StackManipulation.Compound(
                // Call the base constructor for Object.
                MethodVariableAccess.loadThis(),
                Duplication.SINGLE,
                MethodInvocation.invoke(
                    new ForLoadedType(Object.class)
                        .getDeclaredMethods()
                        .filter(
                            ElementMatchers.isConstructor().and(ElementMatchers.takesArguments(0)))
                        .getOnly()),
                MethodVariableAccess.REFERENCE.loadFrom(1),
                FieldAccess.forField(
                        implementationTarget
                            .getInstrumentedType()
                            .getDeclaredFields()
                            .filter(ElementMatchers.named(SELECT_SCHEMA_FIELD_NAME))
                            .getOnly())
                    .write(),
                MethodReturn.VOID);
        StackManipulation.Size size = stackManipulation.apply(methodVisitor, implementationContext);
        return new ByteCodeAppender.Size(size.getMaximalSize(), numLocals);
      };
    }
  }

  // Manage array creation and appending.
  private static class ArrayManager {
    private final StackManipulation.Size sizeDecreaseArrayStore;
    private final int arraySize;
    int currentArrayField = 0;

    ArrayManager(int arraySize) {
      // Size decreases by index and array reference (2) and array element (1, 2) after each element
      // storage.
      this.sizeDecreaseArrayStore =
          StackSize.DOUBLE
              .toDecreasingSize()
              .aggregate(Generic.OBJECT.getStackSize().toDecreasingSize());
      this.arraySize = arraySize;
    }

    StackManipulation createArray() {
      return new StackManipulation() {
        @Override
        public boolean isValid() {
          return true;
        }

        @Override
        public Size apply(MethodVisitor methodVisitor, Context context) {
          StackManipulation.Size size =
              IntegerConstant.forValue(arraySize).apply(methodVisitor, context);
          methodVisitor.visitTypeInsn(
              Opcodes.ANEWARRAY, Generic.OBJECT.asErasure().getInternalName());
          size = size.aggregate(StackSize.ZERO.toDecreasingSize());
          return size;
        }
      };
    }

    StackManipulation append(StackManipulation valueToWrite) {
      return store(currentArrayField++, valueToWrite);
    }

    int reserveSlot() {
      return currentArrayField++;
    }

    StackManipulation store(int arrayIndexToWrite, StackManipulation valueToWrite) {
      Preconditions.checkArgument(arrayIndexToWrite < arraySize);
      return new StackManipulation() {
        @Override
        public boolean isValid() {
          return true;
        }

        @Override
        public StackManipulation.Size apply(MethodVisitor methodVisitor, Context context) {
          StackManipulation stackManipulation =
              new StackManipulation.Compound(
                  Duplication.SINGLE, // Duplicate the array reference
                  IntegerConstant.forValue(arrayIndexToWrite),
                  valueToWrite);
          StackManipulation.Size size = stackManipulation.apply(methodVisitor, context);
          methodVisitor.visitInsn(Opcodes.AASTORE);
          size = size.aggregate(sizeDecreaseArrayStore);
          return size;
        }
      };
    }
  }

  private static class SelectInstruction implements Implementation {
    private final FieldAccessDescriptor fieldAccessDescriptor;
    private final Schema inputSchema;
    private final Schema outputSchema;
    private ByteBuddyLocalVariableManager localVariables;
    private static final int INPUT_ROW_ARG = 1;
    private final int currentSelectRowArg;
    private final int fieldValueArg;

    private static final int NUM_FUNCTION_ARGS = 2; // Args are this and the row to select.

    public SelectInstruction(
        FieldAccessDescriptor fieldAccessDescriptor, Schema inputSchema, Schema outputSchema) {
      this.fieldAccessDescriptor = fieldAccessDescriptor;
      this.inputSchema = inputSchema;
      this.outputSchema = outputSchema;
      this.localVariables = new ByteBuddyLocalVariableManager(NUM_FUNCTION_ARGS);
      currentSelectRowArg = localVariables.createVariable();
      fieldValueArg = localVariables.createVariable();
    }

    @Override
    public InstrumentedType prepare(InstrumentedType instrumentedType) {
      return instrumentedType;
    }

    @Override
    public ByteCodeAppender appender(final Target implementationTarget) {
      return (methodVisitor, implementationContext, instrumentedMethod) -> {
        int numLocals = 1 + instrumentedMethod.getParameters().size();
        StackManipulation.Size size = new StackManipulation.Size(0, numLocals);

        // Create the Row builder.
        StackManipulation createRowBuilder =
            new StackManipulation.Compound(
                MethodVariableAccess.loadThis(),
                FieldAccess.forField(
                        implementationTarget
                            .getInstrumentedType()
                            .getDeclaredFields()
                            .filter(ElementMatchers.named(SELECT_SCHEMA_FIELD_NAME))
                            .getOnly())
                    .read(),
                MethodInvocation.invoke(
                    new ForLoadedType(Row.class)
                        .getDeclaredMethods()
                        .filter(ElementMatchers.named("withSchema"))
                        .getOnly()));
        size = size.aggregate(createRowBuilder.apply(methodVisitor, implementationContext));

        // Create a new object array with one entry for each field in the output schema.
        ArrayManager arrayManager = new ArrayManager(outputSchema.getFieldCount());
        size =
            size.aggregate(arrayManager.createArray().apply(methodVisitor, implementationContext));

        // Store the current input row into a local variable.
        StackManipulation storeRowInLocalVariable =
            localVariables.copy(INPUT_ROW_ARG, currentSelectRowArg);
        size = size.aggregate(storeRowInLocalVariable.apply(methodVisitor, implementationContext));

        // Fill the array values with those selected from the row.
        size =
            size.aggregate(
                selectIntoArray(
                    inputSchema,
                    fieldAccessDescriptor,
                    arrayManager,
                    methodVisitor,
                    implementationContext));

        // Return the actual row.
        StackManipulation attachToRow =
            new StackManipulation.Compound(
                MethodInvocation.invoke(
                    new ForLoadedType(Row.Builder.class)
                        .getDeclaredMethods()
                        .filter(
                            ElementMatchers.named("attachValues")
                                .and(ElementMatchers.takesArguments(Object[].class)))
                        .getOnly()),
                MethodReturn.REFERENCE);
        size = size.aggregate(attachToRow.apply(methodVisitor, implementationContext));
        return new ByteCodeAppender.Size(
            size.getMaximalSize(), localVariables.getTotalNumVariables());
      };
    }

    // Selects a field from the current row being selected (the one stored in
    // currentSelectRowArg).
    private StackManipulation getCurrentRowFieldValue(int i) {
      StackManipulation readRow = localVariables.readVariable(currentSelectRowArg, Row.class);
      StackManipulation getValue =
          new StackManipulation.Compound(
              localVariables.readVariable(currentSelectRowArg, Row.class),
              IntegerConstant.forValue(i),
              MethodInvocation.invoke(
                  ROW_LOADED_TYPE
                      .getDeclaredMethods()
                      .filter(
                          ElementMatchers.named("getValue")
                              .and(ElementMatchers.takesArguments(int.class)))
                      .getOnly()));

      return new ShortCircuitReturnNull(readRow, getValue);
    }

    // Generate bytecode to select all specified fields from the Row. The current row being selected
    // is stored
    // in the local variable at position 2. The current array being written to is at the top of the
    // stack.
    StackManipulation.Size selectIntoArray(
        Schema inputSchema,
        FieldAccessDescriptor fieldAccessDescriptor,
        ArrayManager arrayManager,
        MethodVisitor methodVisitor,
        Context implementationContext) {
      StackManipulation.Size size = new StackManipulation.Size(0, 0);
      if (fieldAccessDescriptor.getAllFields()) {
        StackManipulation storeAllValues =
            new StackManipulation.Compound(
                IntStream.range(0, inputSchema.getFieldCount())
                    .mapToObj(i -> arrayManager.append(getCurrentRowFieldValue(i)))
                    .collect(Collectors.toList()));
        return size.aggregate(storeAllValues.apply(methodVisitor, implementationContext));
      }

      StackManipulation storeAllValues =
          new StackManipulation.Compound(
              fieldAccessDescriptor.fieldIdsAccessed().stream()
                  .map(i -> arrayManager.append(getCurrentRowFieldValue(i)))
                  .collect(Collectors.toList()));
      size = size.aggregate(storeAllValues.apply(methodVisitor, implementationContext));

      for (Map.Entry<FieldDescriptor, FieldAccessDescriptor> nested :
          fieldAccessDescriptor.getNestedFieldsAccessed().entrySet()) {
        FieldDescriptor field = nested.getKey();
        FieldAccessDescriptor nestedAccess = nested.getValue();
        FieldType nestedInputType = inputSchema.getField(field.getFieldId()).getType();

        StackManipulation.Size subSelectSize =
            selectIntoArrayHelper(
                field.getQualifiers(),
                0,
                field.getFieldId(),
                nestedAccess,
                nestedInputType,
                arrayManager,
                methodVisitor,
                implementationContext);
        size = size.aggregate(subSelectSize);
      }
      return size;
    }

    // Load the current field value to process. If a fieldId is specified, then we assume that we
    // need to select
    // a field from the Row stored in currentSelectRowArg. Otherwise we assume that the value has
    // already been stored
    // in fieldValueArg, and we read from there.
    private StackManipulation loadFieldValue(int fieldId) {
      if (fieldId != -1) {
        // If a field id was specified, then load it.
        return getCurrentRowFieldValue(fieldId);
      } else {
        // If not specified, then assume it was stored in this member variable.
        return localVariables.readVariable(fieldValueArg);
      }
    }

    private StackManipulation.Size selectIntoArrayHelper(
        List<Qualifier> qualifiers,
        int qualifierPosition,
        int fieldId,
        FieldAccessDescriptor fieldAccessDescriptor,
        FieldType inputType,
        ArrayManager arrayManager,
        MethodVisitor methodVisitor,
        Context implementationContext) {
      StackManipulation.Size size = new StackManipulation.Size(0, 0);

      if (qualifierPosition >= qualifiers.size()) {
        // We have already constructed all arrays and maps. What remains must be a Row.
        ByteBuddyLocalVariableManager.BackupLocalVariable backup =
            localVariables.backupVariable(currentSelectRowArg);
        StackManipulation updateLocalVariable =
            new StackManipulation.Compound(
                loadFieldValue(fieldId),
                // Backup the current value of the currentSelectRowArg variable.
                backup.backup(),
                // Update the row local with the current stack value.
                localVariables.writeVariable(currentSelectRowArg));
        size = size.aggregate(updateLocalVariable.apply(methodVisitor, implementationContext));

        size =
            size.aggregate(
                selectIntoArray(
                    inputType.getRowSchema(),
                    fieldAccessDescriptor,
                    arrayManager,
                    methodVisitor,
                    implementationContext));
        // Restore the value of currentSelectRowArg from the temp variable.
        size = size.aggregate(backup.restore().apply(methodVisitor, implementationContext));
        return size;
      }

      Qualifier qualifier = qualifiers.get(qualifierPosition);
      switch (qualifier.getKind()) {
        case LIST:
          return size.aggregate(
              processList(
                  inputType,
                  fieldAccessDescriptor,
                  qualifiers,
                  qualifierPosition,
                  fieldId,
                  arrayManager,
                  methodVisitor,
                  implementationContext));

        case MAP:
          return size.aggregate(
              processMap(
                  inputType,
                  fieldAccessDescriptor,
                  qualifiers,
                  qualifierPosition,
                  fieldId,
                  arrayManager,
                  methodVisitor,
                  implementationContext));

        default:
          throw new RuntimeException("Unexpected type " + qualifier.getKind());
      }
    }

    private StackManipulation.Size processList(
        FieldType inputType,
        FieldAccessDescriptor fieldAccessDescriptor,
        List<Qualifier> qualifiers,
        int qualifierPosition,
        int fieldId,
        ArrayManager arrayManager,
        MethodVisitor methodVisitor,
        Context implementationContext) {
      StackManipulation.Size size = new StackManipulation.Size(0, 0);
      FieldType nestedInputType = checkNotNull(inputType.getCollectionElementType());
      Schema nestedSchema = getNestedSchema(nestedInputType, fieldAccessDescriptor);

      // We create temp local variables to store all the arrays we create. Each field in
      // nestedSchema corresponds to a separate array in the output.
      int[] localVariablesForArrays =
          IntStream.range(0, nestedSchema.getFieldCount())
              .map(i -> localVariables.createVariable())
              .toArray();
      // Each field returned in nestedSchema will become it's own list in the output. So let's
      // iterate and create arrays and store each one in the output.
      StackManipulation createAllArrayLists =
          new StackManipulation.Compound(
              Arrays.stream(localVariablesForArrays)
                  .mapToObj(
                      v -> {
                        StackManipulation createArrayList =
                            new StackManipulation.Compound(
                                MethodInvocation.invoke(LISTS_NEW_ARRAYLIST),
                                // Store the ArrayList in a local variable.
                                Duplication.SINGLE,
                                localVariables.writeVariable(v));
                        StackManipulation storeNull =
                            new StackManipulation.Compound(
                                NullConstant.INSTANCE,
                                localVariables.writeVariable(v),
                                NullConstant.INSTANCE);

                        // Create the array only if the input isn't null. Otherwise store a null
                        // value into the output
                        // array.
                        int arraySlot = arrayManager.reserveSlot();
                        return new IfNullElse(
                            loadFieldValue(fieldId),
                            arrayManager.store(arraySlot, storeNull),
                            arrayManager.store(arraySlot, createArrayList));
                      })
                  .collect(Collectors.toList()));
      size = size.aggregate(createAllArrayLists.apply(methodVisitor, implementationContext));

      // If the input variable is null, then don't try and iterate over it.
      Label onNullLabel = new Label();
      size = size.aggregate(loadFieldValue(fieldId).apply(methodVisitor, implementationContext));
      methodVisitor.visitJumpInsn(Opcodes.IFNULL, onNullLabel);
      size = size.aggregate(StackSize.SINGLE.toDecreasingSize());

      // Now iterate over the value, selecting from each element.
      StackManipulation readListIterator =
          new StackManipulation.Compound(
              loadFieldValue(fieldId),
              TypeCasting.to(new ForLoadedType(Iterable.class)),
              MethodInvocation.invoke(ITERABLE_ITERATOR));
      size = size.aggregate(readListIterator.apply(methodVisitor, implementationContext));

      Label startLoopLabel = new Label();
      Label exitLoopLabel = new Label();
      // Loop over the entire iterable.
      methodVisitor.visitLabel(startLoopLabel);

      StackManipulation checkTerminationCondition =
          new StackManipulation.Compound(
              Duplication.SINGLE, MethodInvocation.invoke(ITERATOR_HASNEXT));
      size = size.aggregate(checkTerminationCondition.apply(methodVisitor, implementationContext));
      methodVisitor.visitJumpInsn(Opcodes.IFEQ, exitLoopLabel); // Exit the loop if !hasNext().
      size = size.aggregate(StackSize.SINGLE.toDecreasingSize());

      // Read the next value in the iterator.
      StackManipulation getNext =
          new StackManipulation.Compound(
              Duplication.SINGLE, MethodInvocation.invoke(ITERATOR_NEXT));
      size = size.aggregate(getNext.apply(methodVisitor, implementationContext));

      ByteBuddyLocalVariableManager.BackupLocalVariable backupFieldValue = null;
      if (fieldId == -1) {
        // Save the field value arg before overwriting it, as we need it in subsequent
        // iterations of the
        // loop.
        backupFieldValue = localVariables.backupVariable(fieldValueArg);
        size =
            size.aggregate(backupFieldValue.backup().apply(methodVisitor, implementationContext));
      }

      // Recursively generate select with one qualifier consumed. Since we pass in -1 as the
      // field id, the iterator.next() value will be consumed as the value to select instead
      // of accessing the fieldId.
      size =
          size.aggregate(
              localVariables
                  .writeVariable(fieldValueArg)
                  .apply(methodVisitor, implementationContext));

      // Select the fields of interest from this row.
      // Create a new object array with one entry for each field selected from this row.
      ArrayManager nestedArrayManager = new ArrayManager(nestedSchema.getFieldCount());
      size =
          size.aggregate(
              nestedArrayManager.createArray().apply(methodVisitor, implementationContext));

      size =
          size.aggregate(
              selectIntoArrayHelper(
                  qualifiers,
                  qualifierPosition + 1,
                  -1,
                  fieldAccessDescriptor,
                  nestedInputType,
                  nestedArrayManager,
                  methodVisitor,
                  implementationContext));

      if (backupFieldValue != null) {
        // Restore the field value.
        size =
            size.aggregate(backupFieldValue.restore().apply(methodVisitor, implementationContext));
      }

      // Now the top of the stack holds an array containing all the fields selected from the
      // row.
      // Now we need to distribute these fields into the separate arrays we created in the
      // result.
      // That is: if this select returned {a, b}, our final resulting schema will contain two
      // lists,
      // so we must add a to the first and b to the second list.

      int tempVariableForField = localVariables.createVariable();
      for (int i = 0; i < nestedSchema.getFieldCount(); ++i) {
        // Extract the field and store it in a temp variable.
        StackManipulation extractField =
            new StackManipulation.Compound(
                Duplication.SINGLE,
                IntegerConstant.forValue(i),
                ArrayAccess.REFERENCE.load(),
                localVariables.writeVariable(tempVariableForField));
        StackManipulation addItemToList =
            new StackManipulation.Compound(
                localVariables.readVariable(localVariablesForArrays[i]),
                localVariables.readVariable(tempVariableForField),
                MethodInvocation.invoke(LIST_ADD),
                // Ignore return value from add().
                Removal.SINGLE);
        size =
            size.aggregate(
                new StackManipulation.Compound(extractField, addItemToList)
                    .apply(methodVisitor, implementationContext));
      }

      // Pop the created array from the top of the stack.
      size = size.aggregate(Removal.SINGLE.apply(methodVisitor, implementationContext));

      // Go back to the beginning of the loop.
      methodVisitor.visitJumpInsn(Opcodes.GOTO, startLoopLabel);
      methodVisitor.visitLabel(exitLoopLabel);
      // Remove the iterator from the top of the stack.
      size = size.aggregate(Removal.SINGLE.apply(methodVisitor, implementationContext));
      methodVisitor.visitLabel(onNullLabel);
      return size;
    }

    private StackManipulation.Size processMap(
        FieldType inputType,
        FieldAccessDescriptor fieldAccessDescriptor,
        List<Qualifier> qualifiers,
        int qualifierPosition,
        int fieldId,
        ArrayManager arrayManager,
        MethodVisitor methodVisitor,
        Context implementationContext) {
      StackManipulation.Size size = new StackManipulation.Size(0, 0);
      FieldType nestedInputType = checkNotNull(inputType.getMapValueType());
      Schema nestedSchema = getNestedSchema(nestedInputType, fieldAccessDescriptor);

      // We create temp local variables to store all the maps we create. Each field in
      // nestedSchema corresponds to a separate array in the output.
      int[] localVariablesForMaps =
          IntStream.range(0, nestedSchema.getFieldCount())
              .map(i -> localVariables.createVariable())
              .toArray();

      // Each field returned in nestedSchema will become it's own map in the output. So let's
      // iterate and create arrays and store each one in the output.
      StackManipulation createAllHashMaps =
          new StackManipulation.Compound(
              Arrays.stream(localVariablesForMaps)
                  .mapToObj(
                      v -> {
                        StackManipulation createHashMap =
                            new StackManipulation.Compound(
                                MethodInvocation.invoke(MAPS_NEW_HASHMAP),
                                // Store the HashMap in a local variable.
                                Duplication.SINGLE,
                                localVariables.writeVariable(v));
                        StackManipulation storeNull =
                            new StackManipulation.Compound(
                                NullConstant.INSTANCE,
                                localVariables.writeVariable(v),
                                NullConstant.INSTANCE);
                        int arraySlot = arrayManager.reserveSlot();
                        return new IfNullElse(
                            loadFieldValue(fieldId),
                            arrayManager.store(arraySlot, storeNull),
                            arrayManager.store(arraySlot, createHashMap));
                      })
                  .collect(Collectors.toList()));
      size = size.aggregate(createAllHashMaps.apply(methodVisitor, implementationContext));

      // If the input variable is null, then don't try and iterate over it.
      Label onNullLabel = new Label();
      size = size.aggregate(loadFieldValue(fieldId).apply(methodVisitor, implementationContext));
      methodVisitor.visitJumpInsn(Opcodes.IFNULL, onNullLabel);
      size = size.aggregate(StackSize.SINGLE.toDecreasingSize());

      // Now iterate over the value, selecting from each element.
      StackManipulation readMapEntriesIterator =
          new StackManipulation.Compound(
              loadFieldValue(fieldId),
              TypeCasting.to(new ForLoadedType(Map.class)),
              MethodInvocation.invoke(MAP_ENTRYSET),
              MethodInvocation.invoke(ITERABLE_ITERATOR));
      size = size.aggregate(readMapEntriesIterator.apply(methodVisitor, implementationContext));

      // Loop over the entire entrySet iterable.
      Label startLoopLabel = new Label();
      Label exitLoopLabel = new Label();
      methodVisitor.visitLabel(startLoopLabel);

      StackManipulation checkTerminationCondition =
          new StackManipulation.Compound(
              Duplication.SINGLE, MethodInvocation.invoke(ITERATOR_HASNEXT));
      size = size.aggregate(checkTerminationCondition.apply(methodVisitor, implementationContext));
      methodVisitor.visitJumpInsn(Opcodes.IFEQ, exitLoopLabel); // Exit the loop if !hasNext().
      size = size.aggregate(StackSize.SINGLE.toDecreasingSize());

      int keyVariable = localVariables.createVariable();
      // Read the next value in the iterator.
      StackManipulation getNext =
          new StackManipulation.Compound(
              Duplication.SINGLE,
              MethodInvocation.invoke(ITERATOR_NEXT),
              // Get the key and store it in the keyVariable.
              Duplication.SINGLE,
              MethodInvocation.invoke(MAPENTRY_GETKEY),
              localVariables.writeVariable(keyVariable),
              // Get the value and leave it on the stack.
              MethodInvocation.invoke(MAPENTRY_GETVALUE));
      size = size.aggregate(getNext.apply(methodVisitor, implementationContext));

      ByteBuddyLocalVariableManager.BackupLocalVariable backupFieldValue = null;
      if (fieldId == -1) {
        // Save the field value arg before overwriting it, as we need it in subsequent
        // iterations of the
        // loop.
        backupFieldValue = localVariables.backupVariable(fieldValueArg);
        size =
            size.aggregate(backupFieldValue.backup().apply(methodVisitor, implementationContext));
      }

      // Recursively generate select with one qualifier consumed. Since we pass in -1 as the
      // field id, the iterator.next() value will be consumed as the value to select instead
      // of accessing the fieldId.
      size =
          size.aggregate(
              localVariables
                  .writeVariable(fieldValueArg)
                  .apply(methodVisitor, implementationContext));

      // Select the fields of interest from this row.
      // Create a new object array with one entry for each field selected from this row.
      ArrayManager nestedArrayManager = new ArrayManager(nestedSchema.getFieldCount());
      size =
          size.aggregate(
              nestedArrayManager.createArray().apply(methodVisitor, implementationContext));

      size =
          size.aggregate(
              selectIntoArrayHelper(
                  qualifiers,
                  qualifierPosition + 1,
                  -1,
                  fieldAccessDescriptor,
                  nestedInputType,
                  nestedArrayManager,
                  methodVisitor,
                  implementationContext));

      if (backupFieldValue != null) {
        // Restore the field value.
        size =
            size.aggregate(backupFieldValue.restore().apply(methodVisitor, implementationContext));
      }

      // Now the top of the stack holds an array containing all the fields selected from the
      // row.
      // Now we need to distribute these fields into the separate arrays we created in the
      // result.
      // That is: if this select returned {a, b}, our final resulting schema will contain two
      // maps,
      // so we must add a to the first and b to the second map.

      int tempVariableForField = localVariables.createVariable();
      for (int i = 0; i < nestedSchema.getFieldCount(); ++i) {
        // Extract the field and store it in a temp variable.
        StackManipulation extractField =
            new StackManipulation.Compound(
                Duplication.SINGLE,
                IntegerConstant.forValue(i),
                ArrayAccess.REFERENCE.load(),
                localVariables.writeVariable(tempVariableForField));
        StackManipulation addItemToMap =
            new StackManipulation.Compound(
                localVariables.readVariable(localVariablesForMaps[i]),
                localVariables.readVariable(keyVariable),
                localVariables.readVariable(tempVariableForField),
                MethodInvocation.invoke(MAP_PUT),
                // Ignore return value from add().
                Removal.SINGLE);
        size =
            size.aggregate(
                new StackManipulation.Compound(extractField, addItemToMap)
                    .apply(methodVisitor, implementationContext));
      }

      // Pop the created array from the top of the stack.
      size = size.aggregate(Removal.SINGLE.apply(methodVisitor, implementationContext));

      // Go back to the beginning of the loop.
      methodVisitor.visitJumpInsn(Opcodes.GOTO, startLoopLabel);
      methodVisitor.visitLabel(exitLoopLabel);
      // Remove the iterator from the top of the stack.
      size = size.aggregate(Removal.SINGLE.apply(methodVisitor, implementationContext));
      methodVisitor.visitLabel(onNullLabel);
      return size;
    }

    private Schema getNestedSchema(
        FieldType nestedInputType, FieldAccessDescriptor fieldAccessDescriptor) {
      // When selecting multiple subelements under a list, we distribute the select
      // resulting in multiple lists. For example, if there is a field "list" with type
      // {a: string, b: int}[], selecting list.a, list.b results in a schema of type
      // {a: string[], b: int[]}. This preserves the invariant that the name selected always
      // appears in the top-level schema.
      Schema tempSchema = Schema.builder().addField("a", nestedInputType).build();
      FieldAccessDescriptor tempAccessDescriptor =
          FieldAccessDescriptor.create()
              .withNestedField("a", fieldAccessDescriptor)
              .resolve(tempSchema);
      return SelectHelpers.getOutputSchema(tempSchema, tempAccessDescriptor);
    }
  }
}
