Simplify schema provider definition (#8829)

ethereum/spec/src/main/java/tech/pegasys/teku/spec/SpecMilestone.java

@@ -13,6 +13,7 @@
 
 package tech.pegasys.teku.spec;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import java.util.Arrays;
@@ -55,12 +56,21 @@ public boolean isLessThanOrEqualTo(final SpecMilestone other) {
   }
 
   /** Returns the milestone prior to this milestone */
+  @SuppressWarnings("EnumOrdinal")
   public SpecMilestone getPreviousMilestone() {
-    if (equals(PHASE0)) {
-      throw new IllegalArgumentException("There is no milestone prior to Phase0");
+    checkArgument(!equals(PHASE0), "There is no milestone prior to Phase0");
+    final SpecMilestone[] values = SpecMilestone.values();
+    return values[ordinal() - 1];
+  }
+
+  /** Returns the milestone prior to this milestone */
+  @SuppressWarnings("EnumOrdinal")
+  public Optional<SpecMilestone> getPreviousMilestoneIfExists() {
+    if (this.equals(PHASE0)) {
+      return Optional.empty();
     }
-    final List<SpecMilestone> priorMilestones = getAllPriorMilestones(this);
-    return priorMilestones.getLast();
+    final SpecMilestone[] values = SpecMilestone.values();
+    return Optional.of(values[ordinal() - 1]);
   }
 
   /**

ethereum/spec/src/main/java/tech/pegasys/teku/spec/schemas/registry/BaseSchemaProvider.java

@@ -29,32 +29,38 @@ class BaseSchemaProvider<T> implements SchemaProvider<T> {
   private final TreeMap<SpecMilestone, SchemaProviderCreator<T>> milestoneToSchemaCreator =
       new TreeMap<>();
   private final SchemaId<T> schemaId;
+  private final boolean alwaysCreateNewSchema;
 
   private BaseSchemaProvider(
       final SchemaId<T> schemaId,
       final List<SchemaProviderCreator<T>> schemaProviderCreators,
       final SpecMilestone untilMilestone,
-      final boolean isConstant) {
+      final boolean alwaysCreateNewSchema) {
     this.schemaId = schemaId;
+    this.alwaysCreateNewSchema = alwaysCreateNewSchema;
     final List<SchemaProviderCreator<T>> creatorsList = new ArrayList<>(schemaProviderCreators);
 
     SchemaProviderCreator<T> lastCreator = null;
 
     for (final SpecMilestone milestone : SpecMilestone.getMilestonesUpTo(untilMilestone)) {
-      if (!creatorsList.isEmpty() && creatorsList.getFirst().milestone == milestone) {
+      if (!creatorsList.isEmpty() && creatorsList.getFirst().baseMilestone == milestone) {
         lastCreator = creatorsList.removeFirst();
       }
 
       if (lastCreator != null) {
-        milestoneToSchemaCreator.put(
-            milestone, isConstant ? lastCreator : lastCreator.withMilestone(milestone));
+        milestoneToSchemaCreator.put(milestone, lastCreator);
       }
     }
   }
 
   @Override
-  public SpecMilestone getEffectiveMilestone(final SpecMilestone milestone) {
-    return getSchemaCreator(milestone).milestone;
+  public SpecMilestone getBaseMilestone(final SpecMilestone milestone) {
+    return getSchemaCreator(milestone).baseMilestone;
+  }
+
+  @Override
+  public boolean alwaysCreateNewSchema() {
+    return alwaysCreateNewSchema;
   }
 
   @Override
@@ -90,100 +96,107 @@ public Set<SpecMilestone> getSupportedMilestones() {
   }
 
   protected record SchemaProviderCreator<T>(
-      SpecMilestone milestone, BiFunction<SchemaRegistry, SpecConfig, T> creator) {
-
-    private SchemaProviderCreator<T> withMilestone(final SpecMilestone milestone) {
-      return new SchemaProviderCreator<>(milestone, creator);
-    }
+      SpecMilestone baseMilestone, BiFunction<SchemaRegistry, SpecConfig, T> creator) {
 
     @Override
     public String toString() {
-      return MoreObjects.toStringHelper(this).add("milestone", milestone).toString();
+      return MoreObjects.toStringHelper(this).add("baseMilestone", baseMilestone).toString();
     }
   }
 
   /**
-   * Creates a builder for a constant schema provider.<br>
-   * This can be used when schema remains the same across multiple milestones. <br>
+   * Creates a builder for a schema provider.<br>
    * Example usage:
    *
    * <pre>{@code
-   * constantProviderBuilder(EXAMPLE_SCHEMA)
-   *    .withCreator(ALTAIR, (registry, config) -> new ExampleSchemaAltair())
-   *    .withCreator(ELECTRA, (registry, config) -> new ExampleSchemaElectra())
+   * providerBuilder(EXAMPLE_SCHEMA)
+   *    .withCreator(ALTAIR, (registry, config) -> ExampleSchema1.create(registry, config))
+   *    .withCreator(ELECTRA, (registry, config) -> ExampleSchema2.create(registry, config))
    *    .build();
    *
    * }</pre>
    *
    * this will create a schema provider that will generate: <br>
-   * - only one ExampleSchemaAltair instance which will be reused from ALTAIR to BELLATRIX <br>
-   * - only one ExampleSchemaElectra instance which will be reused from ELECTRA to last known
-   * milestone <br>
-   */
-  static <T> Builder<T> constantProviderBuilder(final SchemaId<T> schemaId) {
-    return new Builder<>(schemaId, true);
-  }
-
-  /**
-   * Creates a builder for a variable schema provider.<br>
-   * This can be used when schema changes across multiple milestones (i.e. depends on changing
-   * schemas) <br>
-   * Example usage:
+   * - a new ExampleSchema1 for each milestone from ALTAIR to CAPELLA <br>
+   * - a new ExampleSchema2 for each milestone from ELECTRA to last known milestone <br>
+   *
+   * <p>By default, the schema provider will check for schema equality when a creator is used
+   * multiple times. In the previous example, if ExampleSchema1.create generates schemas that are
+   * <b>equals</b> in both ALTAIR and BELLATRIX context, the ALTAIR instance will be used for
+   * BELLATRIX too.<br>
+   * Since the equality check does not consider names, semantically equivalent schemas with
+   * different fields or container names will be considered equal.<br>
+   *
+   * <p>If the equality check is relevant, this behavior can be avoided in two ways:<br>
+   * - specifying a new creator like: <br>
    *
    * <pre>{@code
    * variableProviderBuilder(EXAMPLE_SCHEMA)
-   *    .withCreator(ALTAIR, (registry, config) -> new ExampleSchema1(registry, config))
-   *    .withCreator(ELECTRA, (registry, config) -> new ExampleSchema2(registry, config))
+   *    .withCreator(ALTAIR, (registry, config) -> ExampleSchema1.create(registry, config))
+   *    .withCreator(BELLATRIX, (registry, config) -> ExampleSchema1.create(registry, config))
+   *    .withCreator(ELECTRA, (registry, config) -> ExampleSchema2.create(registry, config))
    *    .build();
    *
    * }</pre>
    *
-   * this will create a schema provider that will generate: <br>
-   * - a new instance of ExampleSchema1 for each milestone from ALTAIR to ELECTRA <br>
-   * - a new instance of ExampleSchema2 for each milestone from ELECTRA to last known milestone <br>
+   * - using {@link Builder#alwaysCreateNewSchema()}
    */
-  static <T> Builder<T> variableProviderBuilder(final SchemaId<T> schemaId) {
-    return new Builder<>(schemaId, false);
+  static <T> Builder<T> providerBuilder(final SchemaId<T> schemaId) {
+    return new Builder<>(schemaId);
   }
 
   static class Builder<T> {
     private final SchemaId<T> schemaId;
-    private final boolean isConstant;
     final List<SchemaProviderCreator<T>> schemaProviderCreators = new ArrayList<>();
     private SpecMilestone untilMilestone = SpecMilestone.getHighestMilestone();
+    private boolean alwaysCreateNewSchema = false;
 
-    private Builder(final SchemaId<T> schemaId, final boolean isConstant) {
+    private Builder(final SchemaId<T> schemaId) {
       this.schemaId = schemaId;
-      this.isConstant = isConstant;
     }
 
     public Builder<T> withCreator(
         final SpecMilestone milestone,
         final BiFunction<SchemaRegistry, SpecConfig, T> creationSchema) {
       checkArgument(
           schemaProviderCreators.isEmpty()
-              || milestone.isGreaterThan(schemaProviderCreators.getLast().milestone),
+              || milestone.isGreaterThan(schemaProviderCreators.getLast().baseMilestone),
           "Creator's milestones must added in strict ascending order for %s",
           schemaId);
 
       schemaProviderCreators.add(new SchemaProviderCreator<>(milestone, creationSchema));
       return this;
     }
 
+    /**
+     * This can be used when a schema is deprecated and should not be used for newer milestones.
+     *
+     * @param untilMilestone the last milestone for which the schema will be created
+     */
     public Builder<T> until(final SpecMilestone untilMilestone) {
       this.untilMilestone = untilMilestone;
       return this;
     }
 
+    /**
+     * Forces schema provider to create a new schema on each milestone, disabling schema equality
+     * check with previous milestone. Refer to {@link BaseSchemaProvider} for more information.
+     */
+    public Builder<T> alwaysCreateNewSchema() {
+      this.alwaysCreateNewSchema = true;
+      return this;
+    }
+
     public BaseSchemaProvider<T> build() {
       checkArgument(
           !schemaProviderCreators.isEmpty(), "There should be at least 1 creator for %s", schemaId);
 
       checkArgument(
-          untilMilestone.isGreaterThanOrEqualTo(schemaProviderCreators.getLast().milestone),
+          untilMilestone.isGreaterThanOrEqualTo(schemaProviderCreators.getLast().baseMilestone),
           "until must be greater or equal than last creator milestone in %s",
           schemaId);
-      return new BaseSchemaProvider<>(schemaId, schemaProviderCreators, untilMilestone, isConstant);
+      return new BaseSchemaProvider<>(
+          schemaId, schemaProviderCreators, untilMilestone, alwaysCreateNewSchema);
     }
   }
 }

ethereum/spec/src/main/java/tech/pegasys/teku/spec/schemas/registry/SchemaProvider.java

@@ -22,7 +22,9 @@ interface SchemaProvider<T> {
 
   Set<SpecMilestone> getSupportedMilestones();
 
-  SpecMilestone getEffectiveMilestone(SpecMilestone version);
+  SpecMilestone getBaseMilestone(SpecMilestone version);
+
+  boolean alwaysCreateNewSchema();
 
   SchemaId<T> getSchemaId();
 }

ethereum/spec/src/main/java/tech/pegasys/teku/spec/schemas/registry/SchemaRegistry.java

@@ -75,22 +75,11 @@ public <T> T get(final SchemaId<T> schemaId) {
               + " or it does not support milestone "
               + milestone);
     }
-    T schema = cache.get(milestone, schemaId);
+    final T schema = cache.get(milestone, schemaId);
     if (schema != null) {
       return schema;
     }
 
-    // let's check if the schema is stored associated to the effective milestone
-    final SpecMilestone effectiveMilestone = provider.getEffectiveMilestone(milestone);
-    if (effectiveMilestone != milestone) {
-      schema = cache.get(effectiveMilestone, schemaId);
-      if (schema != null) {
-        // let's cache the schema for current milestone as well
-        cache.put(milestone, schemaId, schema);
-        return schema;
-      }
-    }
-
     // The schema was not found.
     // we reach this point only during priming when we actually ask providers to generate schemas
     checkState(!primed, "Registry is primed but schema not found for %s", schemaId);
@@ -101,17 +90,31 @@ public <T> T get(final SchemaId<T> schemaId) {
     }
 
     // actual schema creation (may trigger recursive registry lookups)
-    schema = provider.getSchema(this);
+    final T createdSchema = provider.getSchema(this);
 
     // release the provider
     INFLIGHT_PROVIDERS.remove(provider);
 
-    // cache the schema
-    cache.put(effectiveMilestone, schemaId, schema);
-    if (effectiveMilestone != milestone) {
-      cache.put(milestone, schemaId, schema);
+    // let's check if the created schema is equal to the one from the previous milestone
+    final SpecMilestone effectiveMilestone = provider.getBaseMilestone(milestone);
+    final T resolvedSchema;
+    if (provider.alwaysCreateNewSchema()) {
+      resolvedSchema = createdSchema;
+    } else {
+      resolvedSchema =
+          milestone
+              .getPreviousMilestoneIfExists()
+              .filter(
+                  previousMilestone -> previousMilestone.isGreaterThanOrEqualTo(effectiveMilestone))
+              .map(previousMilestone -> cache.get(previousMilestone, schemaId))
+              .filter(previousSchema -> previousSchema.equals(createdSchema))
+              .orElse(createdSchema);
     }
-    return schema;
+
+    // cache the schema
+    cache.put(milestone, schemaId, resolvedSchema);
+
+    return resolvedSchema;
   }
 
   public SpecMilestone getMilestone() {