Add CodecGenerator and FieldGenerator abstractions.

These are intended to be sufficiently general for any AutoCodec implementations.

PiperOrigin-RevId: 586436258
Change-Id: If6dfcc8a7dd9b64c546d17f84f6c4a3833c2b425

src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecProcessor.java

@@ -321,7 +321,7 @@ private MethodSpec buildSerializeMethodWithInstantiator(
         }
         TypeKind typeKind = parameter.asType().getKind();
         serializeBuilder.addStatement(
-            "$T unsafe_$L = ($T) $T.unsafe().get$L(input, $L_offset)",
+            "$T unsafe_$L = ($T) $T.unsafe().get$L(obj, $L_offset)",
             sanitizeTypeParameter(parameter.asType(), env),
             parameter.getSimpleName(),
             sanitizeTypeParameter(parameter.asType(), env),
@@ -390,7 +390,7 @@ private void addSerializeParameterWithGetter(
   }
 
   private static String turnGetterIntoExpression(String getterName) {
-    return "input." + getterName + "()";
+    return "obj." + getterName + "()";
   }
 
   private MethodSpec buildSerializeMethodWithInstantiatorForAutoValue(

src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/BUILD

@@ -107,6 +107,8 @@ java_library(
     name = "autocodec-processor",
     srcs = [
         "AutoCodecProcessor.java",
+        "CodecGenerator.java",
+        "FieldGenerator.java",
         "Initializers.java",
         "Marshallers.java",
         "SerializationCodeGenerator.java",

src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/CodecGenerator.java

@@ -0,0 +1,101 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed 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 com.google.devtools.build.lib.skyframe.serialization.autocodec;
+
+import static com.google.devtools.build.lib.skyframe.serialization.autocodec.Initializers.initializeCodecClassBuilder;
+import static com.google.devtools.build.lib.skyframe.serialization.autocodec.Initializers.initializeSerializeMethodBuilder;
+
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
+import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.TypeSpec;
+import java.util.List;
+import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.element.ExecutableElement;
+import javax.lang.model.element.TypeElement;
+
+/** Defines an abstract strategy for generating {@link ObjectCodec} implementations. */
+// TODO(b/297857068): migrate other types to this class.
+abstract class CodecGenerator {
+  final ProcessingEnvironment env;
+
+  CodecGenerator(ProcessingEnvironment env) {
+    this.env = env;
+  }
+
+  /** Creates the codec by delegating lower level implementation methods. */
+  final TypeSpec defineCodec(
+      TypeElement encodedType, AutoCodec annotation, ExecutableElement instantiator)
+      throws SerializationProcessingException {
+    List<? extends FieldGenerator> fieldGenerators = getFieldGenerators(encodedType);
+
+    TypeSpec.Builder classBuilder = initializeCodecClassBuilder(encodedType, env);
+    performAdditionalCodecInitialization(classBuilder, encodedType, instantiator);
+
+    MethodSpec.Builder constructor = initializeConstructor(!fieldGenerators.isEmpty());
+    MethodSpec.Builder serialize = initializeSerializeMethodBuilder(encodedType, annotation, env);
+    MethodSpec.Builder deserialize = initializeDeserializeMethod(encodedType);
+
+    for (FieldGenerator generator : fieldGenerators) {
+      generator.generateOffsetMember(classBuilder, constructor);
+      generator.generateAdditionalMemberVariables(classBuilder);
+      generator.generateConstructorCode(constructor);
+      generator.generateSerializeCode(serialize);
+      generator.generateDeserializeCode(deserialize);
+    }
+
+    addImplementationToEndOfMethods(
+        instantiator, constructor, deserialize, !fieldGenerators.isEmpty());
+
+    return classBuilder
+        .addMethod(constructor.build())
+        .addMethod(serialize.build())
+        .addMethod(deserialize.build())
+        .build();
+  }
+
+  /**
+   * Performs additional initialization steps on the codec being created.
+   *
+   * <p>Adds the correct superclass. May define additional field-independent methods.
+   */
+  abstract void performAdditionalCodecInitialization(
+      TypeSpec.Builder classBuilder, TypeElement encodedType, ExecutableElement instantiator);
+
+  /** Creates {@link FieldGenerator} instances that generate code for serialized fields. */
+  abstract List<? extends FieldGenerator> getFieldGenerators(TypeElement type)
+      throws SerializationProcessingException;
+
+  /**
+   * Initializes the field-independent parts of the constructor.
+   *
+   * @param hasFields true if there are any fields to serialize. Exception handling logic may depend
+   *     on the presence of fields.
+   */
+  abstract MethodSpec.Builder initializeConstructor(boolean hasFields);
+
+  /** Initializes the method that performs deserialization work. */
+  abstract MethodSpec.Builder initializeDeserializeMethod(TypeElement encodedType);
+
+  /**
+   * Adds field-independent code at the end of methods after per-field code is added.
+   *
+   * @param hasFields true if there are any fields to serialize, based on the result of {@link
+   *     #getFieldGenerators}. Exception handling logic may depend on the presence of fields.
+   */
+  abstract void addImplementationToEndOfMethods(
+      ExecutableElement instantiator,
+      MethodSpec.Builder constructor,
+      MethodSpec.Builder deserialize,
+      boolean hasFields);
+}

src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/FieldGenerator.java

@@ -0,0 +1,130 @@
+// Copyright 2023 The Bazel Authors. All rights reserved.
+//
+// Licensed 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 com.google.devtools.build.lib.skyframe.serialization.autocodec;
+
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
+import com.google.devtools.build.lib.unsafe.UnsafeProvider;
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.TypeSpec;
+import javax.lang.model.element.Modifier;
+import javax.lang.model.element.Name;
+import javax.lang.model.element.TypeElement;
+import javax.lang.model.element.VariableElement;
+
+/**
+ * Generates code for a specific field.
+ *
+ * <p>Always stores the offset of the field in a variable named {@link #getOffsetName}.
+ */
+abstract class FieldGenerator {
+  /**
+   * A string that should never occur in user code identifiers.
+   *
+   * <p>This is used to create unique identifiers that won't collide.
+   */
+  private static final String GENERATED_TAG = "$AutoCodec$";
+
+  private static final String OFFSET_SUFFIX = "_offset";
+
+  private final VariableElement variable;
+  private final ClassName parentName;
+  private final String namePrefix;
+
+  /**
+   * Constructor.
+   *
+   * @param variable the field being serialized. Note that {@link VariableElement} contains a
+   *     reference to the enclosing class.
+   * @param hierarchyLevel a variable could occur in either the class being serialized or in one of
+   *     its ancestor classes. This is 0 for the class itself, 1 for its superclass, and so on. It
+   *     is used to avoid naming collisions, particularly in the case of shadowed variables.
+   */
+  FieldGenerator(VariableElement variable, int hierarchyLevel) {
+    this.variable = variable;
+    this.parentName = ClassName.get((TypeElement) variable.getEnclosingElement());
+    this.namePrefix = variable.getSimpleName() + GENERATED_TAG + hierarchyLevel;
+  }
+
+  /** Any created member variables should start with this prefix. */
+  final String getNamePrefix() {
+    return namePrefix;
+  }
+
+  /** The name of the enclosing class, used in generated code. */
+  final ClassName getParentName() {
+    return parentName;
+  }
+
+  /**
+   * Generated codecs store the offset of every field of the serialized class in a member variable.
+   *
+   * @return name of the offset member variable
+   */
+  final String getOffsetName() {
+    return namePrefix + OFFSET_SUFFIX;
+  }
+
+  /**
+   * Name of the field being serialized.
+   *
+   * <p>Implementations may refer to this for reflection.
+   */
+  final Name getParameterName() {
+    return variable.getSimpleName();
+  }
+
+  final void generateOffsetMember(TypeSpec.Builder classBuilder, MethodSpec.Builder constructor) {
+    classBuilder.addField(long.class, getOffsetName(), Modifier.PRIVATE, Modifier.FINAL);
+    constructor.addStatement(
+        "this.$L = $T.unsafe().objectFieldOffset($T.class.getDeclaredField(\"$L\"))",
+        getOffsetName(),
+        UnsafeProvider.class,
+        getParentName(),
+        variable.getSimpleName());
+  }
+
+  /**
+   * Generates any additional member variables needed for this field.
+   *
+   * <p>To avoid collisions, field specific field names should be prefixed with {@link #namePrefix}.
+   *
+   * <p>The *offset* field is already generated by {@link #generateOffsetMember}.
+   */
+  void generateAdditionalMemberVariables(TypeSpec.Builder classBuilder) {}
+
+  /**
+   * Adds field specific code to the constructor.
+   *
+   * <p>Many implementations don't need to do anything here given that the offset is already
+   * initialized by {@link #generateOffsetMember}.
+   */
+  void generateConstructorCode(MethodSpec.Builder constructor) {}
+
+  /**
+   * Adds {@link ObjectCodec#serialize} code to serialize this field.
+   *
+   * <p>Implementations may assume that the parameters of {@link ObjectCodec#serialize} are present
+   * in the generated method's scope and may be referenced.
+   */
+  abstract void generateSerializeCode(MethodSpec.Builder serialize);
+
+  /**
+   * Adds field specific deserialize code.
+   *
+   * <p>Implementations may assume the presence of parameters of {@link
+   * AsyncObjectCodec#deserializeAsync}.
+   */
+  abstract void generateDeserializeCode(MethodSpec.Builder deserialize);
+}

src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Initializers.java

@@ -75,10 +75,10 @@ static MethodSpec.Builder initializeSerializeMethodBuilder(
             .addException(SerializationException.class)
             .addException(IOException.class)
             .addParameter(SerializationContext.class, "context")
-            .addParameter(getErasure(encodedType, env), "input")
+            .addParameter(getErasure(encodedType, env), "obj")
             .addParameter(CodedOutputStream.class, "codedOut");
     if (annotation.checkClassExplicitlyAllowed()) {
-      builder.addStatement("context.checkClassExplicitlyAllowed(getEncodedClass(), input)");
+      builder.addStatement("context.checkClassExplicitlyAllowed(getEncodedClass(), obj)");
     }
     List<? extends TypeMirror> explicitlyAllowedClasses;
     try {