Preserve Starlark field declaration order in StarlarkProvider.getSchema (former getSchemaWithDocumentation)

We want the original field declaration order for starlark_doc_extract, both
to match existing Stardoc behavior and to match existing logic of preserving
the original declaration order of rule attributes etc. in proto output (leaving
sorting up to the renderer).

We want to take this opportunity to clean up StarlarkProvider's field and
method names a little bit:

* rename StarlarkProvider.schema to `fields` to emphasize that these are just
  the fields names, not the full schema given to the provider() call - it has
  a different iteration order and is missing doc strings - and to match the
  getFields() method name, logically enough
* rename StarlarkProvider.getSchemaWithDocumentation to just getSchema - since
  it now returns the original schema from Starlark with the same order, and
  since the method may be used regardless of whether the schema is documented

PiperOrigin-RevId: 528545413
Change-Id: I2756f4520c80acaca17dcd244b61d60aac5539b4
diff --git a/src/main/java/com/google/devtools/build/lib/packages/StarlarkProvider.java b/src/main/java/com/google/devtools/build/lib/packages/StarlarkProvider.java
index 5a82bcd..9654435 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/StarlarkProvider.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/StarlarkProvider.java
@@ -14,13 +14,11 @@
 
 package com.google.devtools.build.lib.packages;
 
-import static com.google.common.base.Verify.verify;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSortedMap;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.Depset;
 import com.google.devtools.build.lib.collect.nestedset.Depset.ElementType;
@@ -70,13 +68,13 @@
 
   // For schemaful providers, the sorted list of allowed field names.
   // The requirement for sortedness comes from StarlarkInfoWithSchema and lets us bisect the fields.
-  @Nullable private final ImmutableList<String> schema;
+  @Nullable private final ImmutableList<String> fields;
 
-  // For schemaful providers, an optional list of field documentation strings, of the same size and
-  // sorted in the same order as the schema. Null for schemaless providers and for providers whose
-  // schema is undocumented. In accordance with the provider() Starlark API, either all schema
-  // fields have documentation strings (possibly empty strings), or none do.
-  @Nullable private final ImmutableList<String> schemaDocumentation;
+  // For schemaful providers, an optional map from field names to documentation strings (if any). In
+  // accordance with the provider() Starlark API, either all schema fields have documentation
+  // strings (possibly empty strings), or none do. The iteration order is the order of fields in the
+  // provider() invocation in Starlark - thus, *not* the order of the `fields` list above.
+  @Nullable private final ImmutableMap<String, Optional<String>> schema;
 
   // Optional custom initializer callback. If present, it is invoked with the same positional and
   // keyword arguments as were passed to the provider constructor. The return value must be a
@@ -134,9 +132,7 @@
 
     @Nullable private String documentation;
 
-    @Nullable private ImmutableList<String> schema;
-
-    @Nullable private ImmutableList<String> schemaDocumentation;
+    @Nullable private ImmutableMap<String, Optional<String>> schema;
 
     @Nullable private StarlarkCallable init;
 
@@ -146,23 +142,31 @@
       this.location = location;
     }
 
-    /** Sets the schema (the list of allowed field names) for the provider built by this builder. */
+    /**
+     * Sets the list of allowed fields for the provider built by this builder, and marks the fields'
+     * documentation as empty.
+     */
     @CanIgnoreReturnValue
-    public Builder setSchema(Collection<String> schema) {
-      this.schema = ImmutableList.sortedCopyOf(schema);
+    public Builder setSchema(Collection<String> fields) {
+      ImmutableMap.Builder<String, Optional<String>> builder = ImmutableMap.builder();
+      for (String field : fields) {
+        builder.put(field, Optional.empty());
+      }
+      this.schema = builder.buildOrThrow();
       return this;
     }
 
     /**
-     * Sets the schema and its documentation (meaning, the list of allowed field names and their
-     * corresponding documentation strings) for the provider built by this builder.
+     * Sets the list of allowed field names and their corresponding documentation strings for the
+     * provider built by this builder.
      */
     @CanIgnoreReturnValue
     public Builder setSchema(Map<String, String> schemaWithDocumentation) {
-      ImmutableSortedMap<String, String> sortedSchemaWithDocumentation =
-          ImmutableSortedMap.copyOf(schemaWithDocumentation);
-      this.schema = sortedSchemaWithDocumentation.keySet().asList();
-      this.schemaDocumentation = sortedSchemaWithDocumentation.values().asList();
+      ImmutableMap.Builder<String, Optional<String>> builder = ImmutableMap.builder();
+      for (Map.Entry<String, String> entry : schemaWithDocumentation.entrySet()) {
+        builder.put(entry.getKey(), Optional.of(entry.getValue()));
+      }
+      this.schema = builder.buildOrThrow();
       return this;
     }
 
@@ -201,7 +205,7 @@
 
     /** Builds a StarlarkProvider. */
     public StarlarkProvider build() {
-      return new StarlarkProvider(location, documentation, schema, schemaDocumentation, init, key);
+      return new StarlarkProvider(location, documentation, schema, init, key);
     }
   }
 
@@ -215,14 +219,13 @@
   private StarlarkProvider(
       Location location,
       @Nullable String documentation,
-      @Nullable ImmutableList<String> schema,
-      @Nullable ImmutableList<String> schemaDocumentation,
+      @Nullable ImmutableMap<String, Optional<String>> schema,
       @Nullable StarlarkCallable init,
       @Nullable Key key) {
     this.location = location;
     this.documentation = documentation;
+    this.fields = schema != null ? ImmutableList.sortedCopyOf(schema.keySet()) : null;
     this.schema = schema;
-    this.schemaDocumentation = schemaDocumentation;
     this.init = init;
     this.key = key;
     if (schema != null) {
@@ -349,32 +352,20 @@
    * schemaless.
    */
   @Nullable
-  // TODO(adonovan): rename getSchema.
   public ImmutableList<String> getFields() {
-    return schema;
+    return fields;
   }
 
   /**
    * Returns the map of fields allowed by this provider mapping to their corresponding documentation
    * strings (if any), or null if this provider is schemaless.
    *
-   * <p>The returned map is guaranteed to have a stable iteration order.
+   * <p>The returned map's iteration order matches the order of fields in the {@code provider()}
+   * invocation in Starlark - thus, different from the order of fields in {@link #getFields}.
    */
   @Nullable
-  public ImmutableMap<String, Optional<String>> getSchemaWithDocumentation() {
-    if (schema == null) {
-      verify(schemaDocumentation == null);
-      return null;
-    }
-    verify(schemaDocumentation == null || schemaDocumentation.size() == schema.size());
-    ImmutableMap.Builder<String, Optional<String>> builder =
-        ImmutableMap.builderWithExpectedSize(schema.size());
-    for (int i = 0; i < schema.size(); i++) {
-      builder.put(
-          schema.get(i),
-          schemaDocumentation == null ? Optional.empty() : Optional.of(schemaDocumentation.get(i)));
-    }
-    return builder.buildOrThrow();
+  public ImmutableMap<String, Optional<String>> getSchema() {
+    return schema;
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/rules/starlarkdocextract/ModuleInfoExtractor.java b/src/main/java/com/google/devtools/build/lib/rules/starlarkdocextract/ModuleInfoExtractor.java
index 8d0ee26..303d045 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/starlarkdocextract/ModuleInfoExtractor.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/starlarkdocextract/ModuleInfoExtractor.java
@@ -294,10 +294,9 @@
     // Allow providers to be exported under a different name (e.g. in a struct).
     providerInfoBuilder.setProviderName(exportedName);
     provider.getDocumentation().ifPresent(providerInfoBuilder::setDocString);
-    ImmutableMap<String, Optional<String>> schemaWithDocumentation =
-        provider.getSchemaWithDocumentation();
-    if (schemaWithDocumentation != null) {
-      for (Map.Entry<String, Optional<String>> entry : schemaWithDocumentation.entrySet()) {
+    ImmutableMap<String, Optional<String>> schema = provider.getSchema();
+    if (schema != null) {
+      for (Map.Entry<String, Optional<String>> entry : schema.entrySet()) {
         if (isExportableName(entry.getKey())) {
           ProviderFieldInfo.Builder fieldInfoBuilder = ProviderFieldInfo.newBuilder();
           fieldInfoBuilder.setName(entry.getKey());
diff --git a/src/test/java/com/google/devtools/build/lib/packages/StarlarkProviderTest.java b/src/test/java/com/google/devtools/build/lib/packages/StarlarkProviderTest.java
index 4f7412d..37ceb83 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/StarlarkProviderTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/StarlarkProviderTest.java
@@ -274,54 +274,62 @@
   }
 
   @Test
-  public void schemalessProvider_getSchemaWithDocumentation() throws Exception {
+  public void schemalessProvider_getSchema() throws Exception {
     StarlarkProvider provider = StarlarkProvider.builder(Location.BUILTIN).build();
-    assertThat(provider.getSchemaWithDocumentation()).isNull();
+    assertThat(provider.getSchema()).isNull();
   }
 
   @Test
   public void providerWithUndocumentedSchema_getFields() throws Exception {
     StarlarkProvider provider =
         StarlarkProvider.builder(Location.BUILTIN)
-            .setSchema(ImmutableList.of("a", "b", "c"))
+            // Note fields in setSchema() call below are not alphabetized to simulate
+            // non-alphabetized field order in a provider declaration in Starlark code.
+            .setSchema(ImmutableList.of("c", "a", "b"))
             .build();
     assertThat(provider.getFields()).containsExactly("a", "b", "c").inOrder();
   }
 
   @Test
-  public void providerWithUndocumentedSchema_getSchemaWithDocumentation() throws Exception {
+  public void providerWithUndocumentedSchema_getSchema() throws Exception {
     StarlarkProvider provider =
         StarlarkProvider.builder(Location.BUILTIN)
-            .setSchema(ImmutableList.of("a", "b", "c"))
+            // Note fields in setSchema() call below are not alphabetized to simulate
+            // non-alphabetized field order in a provider declaration in Starlark code.
+            .setSchema(ImmutableList.of("c", "a", "b"))
             .build();
-    assertThat(provider.getSchemaWithDocumentation().keySet())
-        .containsExactly("a", "b", "c")
+    assertThat(provider.getSchema())
+        .containsExactly("c", Optional.empty(), "a", Optional.empty(), "b", Optional.empty())
         .inOrder();
-    assertThat(provider.getSchemaWithDocumentation().values())
-        .containsExactly(Optional.empty(), Optional.empty(), Optional.empty());
   }
 
   @Test
   public void providerWithDocumentedSchema_getFields() throws Exception {
     StarlarkProvider provider =
         StarlarkProvider.builder(Location.BUILTIN)
-            .setSchema(ImmutableMap.of("a", "Parameter a", "b", "Parameter b", "c", "Parameter c"))
+            // Note fields in setSchema() call below are not alphabetized to simulate
+            // non-alphabetized field order in a provider declaration in Starlark code.
+            .setSchema(ImmutableMap.of("c", "Parameter c", "a", "Parameter a", "b", "Parameter b"))
             .build();
     assertThat(provider.getFields()).containsExactly("a", "b", "c").inOrder();
   }
 
   @Test
-  public void providerWithDocumentedSchema_getSchemaWithDocumentation() throws Exception {
+  public void providerWithDocumentedSchema_getSchema() throws Exception {
     StarlarkProvider provider =
         StarlarkProvider.builder(Location.BUILTIN)
-            .setSchema(ImmutableMap.of("a", "Parameter a", "b", "Parameter b", "c", "Parameter c"))
+            // Note fields in setSchema() call below are not alphabetized to simulate
+            // non-alphabetized field order in a provider declaration in Starlark code.
+            .setSchema(ImmutableMap.of("c", "Parameter c", "a", "Parameter a", "b", "Parameter b"))
             .build();
-    assertThat(provider.getSchemaWithDocumentation().keySet())
-        .containsExactly("a", "b", "c")
-        .inOrder();
-    assertThat(provider.getSchemaWithDocumentation().values())
+    assertThat(provider.getSchema())
         .containsExactly(
-            Optional.of("Parameter a"), Optional.of("Parameter b"), Optional.of("Parameter c"))
+            "c",
+            Optional.of("Parameter c"),
+            "a",
+            Optional.of("Parameter a"),
+            "b",
+            Optional.of("Parameter b"))
         .inOrder();
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/rules/starlarkdocextract/ModuleInfoExtractorTest.java b/src/test/java/com/google/devtools/build/lib/rules/starlarkdocextract/ModuleInfoExtractorTest.java
index 4eb09e7..dcfeecf 100644
--- a/src/test/java/com/google/devtools/build/lib/rules/starlarkdocextract/ModuleInfoExtractorTest.java
+++ b/src/test/java/com/google/devtools/build/lib/rules/starlarkdocextract/ModuleInfoExtractorTest.java
@@ -261,18 +261,22 @@
   public void providerFields() throws Exception {
     Module module =
         exec(
-            "DocumentedInfo = provider(fields = {'a': 'A', 'b': 'B', '_hidden': 'Hidden'})",
-            "UndocumentedInfo = provider(fields = ['a', 'b', '_hidden'])");
+            // Note fields below are not alphabetized
+            "DocumentedInfo = provider(fields = {'c': 'C', 'a': 'A', 'b': 'B', '_hidden':"
+                + " 'Hidden'})",
+            "UndocumentedInfo = provider(fields = ['c', 'a', 'b', '_hidden'])");
     ModuleInfo moduleInfo = getExtractor().extractFrom(module);
     assertThat(moduleInfo.getProviderInfoList())
         .containsExactly(
             ProviderInfo.newBuilder()
                 .setProviderName("DocumentedInfo")
+                .addFieldInfo(ProviderFieldInfo.newBuilder().setName("c").setDocString("C"))
                 .addFieldInfo(ProviderFieldInfo.newBuilder().setName("a").setDocString("A"))
                 .addFieldInfo(ProviderFieldInfo.newBuilder().setName("b").setDocString("B"))
                 .build(),
             ProviderInfo.newBuilder()
                 .setProviderName("UndocumentedInfo")
+                .addFieldInfo(ProviderFieldInfo.newBuilder().setName("c"))
                 .addFieldInfo(ProviderFieldInfo.newBuilder().setName("a"))
                 .addFieldInfo(ProviderFieldInfo.newBuilder().setName("b"))
                 .build());
diff --git a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java
index b07b842..cfef607 100644
--- a/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/starlark/StarlarkRuleClassFunctionsTest.java
@@ -1788,8 +1788,9 @@
         ev,
         "UndocumentedInfo = provider()",
         "DocumentedInfo = provider(doc = 'My documented provider')",
-        "SchemafulWithoutDocsInfo = provider(fields = ['a', 'b'])",
-        "SchemafulWithDocsInfo = provider(fields = {'a': 'Field a', 'b': 'Field b'})");
+        // Note fields below are not alphabetized
+        "SchemafulWithoutDocsInfo = provider(fields = ['b', 'a'])",
+        "SchemafulWithDocsInfo = provider(fields = {'b': 'Field b', 'a': 'Field a'})");
 
     StarlarkProvider undocumentedInfo = (StarlarkProvider) ev.lookup("UndocumentedInfo");
     StarlarkProvider documentedInfo = (StarlarkProvider) ev.lookup("DocumentedInfo");
@@ -1799,10 +1800,10 @@
 
     assertThat(undocumentedInfo.getDocumentation()).isEmpty();
     assertThat(documentedInfo.getDocumentation()).hasValue("My documented provider");
-    assertThat(schemafulWithoutDocsInfo.getSchemaWithDocumentation())
-        .containsExactly("a", Optional.empty(), "b", Optional.empty());
-    assertThat(schemafulWithDocsInfo.getSchemaWithDocumentation())
-        .containsExactly("a", Optional.of("Field a"), "b", Optional.of("Field b"));
+    assertThat(schemafulWithoutDocsInfo.getSchema())
+        .containsExactly("b", Optional.empty(), "a", Optional.empty());
+    assertThat(schemafulWithDocsInfo.getSchema())
+        .containsExactly("b", Optional.of("Field b"), "a", Optional.of("Field a"));
   }
 
   @Test