Inspect post-evaluation exported variables to obtain rule names.

This is a much cleaner, more elegant approach than previous regex matching.
This still leaves room for unknown-name rule definitions, in case, for example, a user namespaces their rule definition not at the top level.
For example: "foo.bar = rule(...)"

RELNOTES: None.
PiperOrigin-RevId: 202380975

10 files changed, 195 insertions, 45 deletions

diff --git a/src/main/java/com/google/devtools/build/skydoc/SkydocMain.java b/src/main/java/com/google/devtools/build/skydoc/SkydocMain.java
index 1e14cc1a16..1d9043f8b7 100644
--- a/src/main/java/com/google/devtools/build/skydoc/SkydocMain.java
+++ b/src/main/java/com/google/devtools/build/skydoc/SkydocMain.java
@@ -14,12 +14,12 @@
 
 package com.google.devtools.build.skydoc;
 
-import static java.nio.charset.StandardCharsets.UTF_8;
-
+import com.google.common.base.Functions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.events.EventHandler;
-import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.skylarkbuildapi.TopLevelBootstrap;
+import com.google.devtools.build.lib.syntax.BaseFunction;
 import com.google.devtools.build.lib.syntax.BuildFileAST;
 import com.google.devtools.build.lib.syntax.Environment;
 import com.google.devtools.build.lib.syntax.Environment.Extension;
@@ -47,8 +47,8 @@ import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
+import java.util.Map.Entry;
+import java.util.stream.Collectors;
 
 /**
  * Main entry point for the Skydoc binary.
@@ -66,11 +66,6 @@ import java.util.regex.Pattern;
  */
 public class SkydocMain {
 
-  // Pattern to match the assignment of a variable to a rule definition
-  // For example, 'my_rule = rule(' will match and have 'my_rule' available as group(1).
-  private static final Pattern ruleDefinitionLinePattern =
-      Pattern.compile("([^\\s]+) = rule\\(");
-
   private final EventHandler eventHandler = new SystemOutEventHandler();
 
   public static void main(String[] args) throws IOException, InterruptedException {
@@ -88,47 +83,52 @@ public class SkydocMain {
     ParserInputSource parserInputSource =
         ParserInputSource.create(content, PathFragment.create(path.toString()));
 
-    List<RuleInfo> ruleInfoList = new SkydocMain().eval(parserInputSource);
+    ImmutableMap.Builder<String, RuleInfo> ruleInfoMap = ImmutableMap.builder();
+    ImmutableList.Builder<RuleInfo> unexportedRuleInfos = ImmutableList.builder();
+
+    new SkydocMain().eval(parserInputSource, ruleInfoMap, unexportedRuleInfos);
 
     try (PrintWriter printWriter = new PrintWriter(outputPath, "UTF-8")) {
-      printRuleInfos(printWriter, ruleInfoList);
+      printRuleInfos(printWriter, ruleInfoMap.build(), unexportedRuleInfos.build());
     }
   }
 
   // TODO(cparsons): Improve output (markdown or HTML).
   private static void printRuleInfos(
-      PrintWriter printWriter, List<RuleInfo> ruleInfos) throws IOException {
-    for (RuleInfo ruleInfo : ruleInfos) {
-      Location location = ruleInfo.getLocation();
-      Path filePath = Paths.get(location.getPath().getPathString());
-      List<String> lines = Files.readAllLines(filePath, UTF_8);
-      String definingString = lines.get(location.getStartLine() - 1);
-      // Rule definitions don't specify their own visible name directly. Instead, the name of
-      // a rule is dependent on the name of the variable assigend to the return value of rule().
-      // This attempts to find a line of the form 'foo = rule(' and thus label the rule as
-      // named 'foo'.
-      // TODO(cparsons): Inspect the global bindings of the environment instead of using string
-      // matching.
-      Matcher matcher = ruleDefinitionLinePattern.matcher(definingString);
-      if (matcher.matches()) {
-        printWriter.println(matcher.group(1));
-      } else {
-        printWriter.println("<unknown name>");
-      }
-      printWriter.println(ruleInfo.getDescription());
+      PrintWriter printWriter,
+      Map<String, RuleInfo> ruleInfos,
+      List<RuleInfo> unexportedRuleInfos) throws IOException {
+    for (Entry<String, RuleInfo> ruleInfoEntry : ruleInfos.entrySet()) {
+      printRuleInfo(printWriter, ruleInfoEntry.getKey(), ruleInfoEntry.getValue());
+    }
+    for (RuleInfo unexportedRuleInfo : unexportedRuleInfos) {
+      printRuleInfo(printWriter, "<unknown name>", unexportedRuleInfo);
     }
   }
 
+  private static void printRuleInfo(
+      PrintWriter printWriter, String exportedName, RuleInfo ruleInfo) {
+    printWriter.println(exportedName);
+    printWriter.println(ruleInfo.getDescription());
+    printWriter.println();
+  }
+
   /**
    * Evaluates/interprets the skylark file at the given input source using a fake build API and
    * collects information about all rule definitions made in that file.
    *
    * @param parserInputSource the input source representing the input skylark file
-   * @return a list of {@link RuleInfo} objects describing the rule definitions
+   * @param ruleInfoMap a map builder to be populated with rule definition information for
+   *     named rules. Keys are exported names of rules, and values are their {@link RuleInfo}
+   *     rule descriptions. For example, 'my_rule = rule(...)' has key 'my_rule'
+   * @param unexportedRuleInfos a list builder to be populated with rule definition information
+   *     for rules which were not exported as top level symbols
    * @throws InterruptedException if evaluation is interrupted
    */
   // TODO(cparsons): Evaluate load statements recursively.
-  public List<RuleInfo> eval(ParserInputSource parserInputSource)
+  public void eval(ParserInputSource parserInputSource,
+      ImmutableMap.Builder<String, RuleInfo> ruleInfoMap,
+      ImmutableList.Builder<RuleInfo> unexportedRuleInfos)
       throws InterruptedException {
     List<RuleInfo> ruleInfoList = new ArrayList<>();
 
@@ -146,7 +146,19 @@ public class SkydocMain {
 
     env.mutability().freeze();
 
-    return ruleInfoList;
+    Map<BaseFunction, RuleInfo> ruleFunctions = ruleInfoList.stream()
+        .collect(Collectors.toMap(
+            RuleInfo::getIdentifierFunction,
+            Functions.identity()));
+
+    for (Entry<String, Object> envEntry : env.getGlobals().getBindings().entrySet()) {
+      if (ruleFunctions.containsKey(envEntry.getValue())) {
+        ruleInfoMap.put(envEntry.getKey(), ruleFunctions.get(envEntry.getValue()));
+        ruleFunctions.remove(envEntry.getValue());
+      }
+    }
+
+    unexportedRuleInfos.addAll(ruleFunctions.values());
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeSkylarkRuleFunctionsApi.java b/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeSkylarkRuleFunctionsApi.java
index dcae202abc..16635fcbca 100644
--- a/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeSkylarkRuleFunctionsApi.java
+++ b/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/FakeSkylarkRuleFunctionsApi.java
@@ -27,18 +27,20 @@ import com.google.devtools.build.lib.syntax.BaseFunction;
 import com.google.devtools.build.lib.syntax.Environment;
 import com.google.devtools.build.lib.syntax.EvalException;
 import com.google.devtools.build.lib.syntax.FuncallExpression;
+import com.google.devtools.build.lib.syntax.Runtime;
 import com.google.devtools.build.lib.syntax.SkylarkDict;
 import com.google.devtools.build.lib.syntax.SkylarkList;
 import com.google.devtools.build.skydoc.rendering.RuleInfo;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import javax.annotation.Nullable;
 
 /**
  * Fake implementation of {@link SkylarkRuleFunctionsApi}.
  *
- * <p>This fake hooks into the global {@code rule()} function, noting calls of that function
- * with a {@link RuleInfoCollector} given in the class constructor.</p>
+ * <p>This fake hooks into the global {@code rule()} function, adding descriptors of calls of that
+ * function to a {@link List} given in the class constructor.</p>
  */
 public class FakeSkylarkRuleFunctionsApi implements SkylarkRuleFunctionsApi<FileApi> {
 
@@ -67,7 +69,7 @@ public class FakeSkylarkRuleFunctionsApi implements SkylarkRuleFunctionsApi<File
       Boolean executionPlatformConstraintsAllowed, SkylarkList<?> execCompatibleWith,
       FuncallExpression ast, Environment funcallEnv) throws EvalException {
     Set<String> attrNames;
-    if (attrs != null) {
+    if (attrs != null && attrs != Runtime.NONE) {
       SkylarkDict<?, ?> attrsDict = (SkylarkDict<?, ?>) attrs;
       Map<String, Descriptor> attrsMap =
           attrsDict.getContents(String.class, Descriptor.class, "attrs");
@@ -76,9 +78,10 @@ public class FakeSkylarkRuleFunctionsApi implements SkylarkRuleFunctionsApi<File
       attrNames = ImmutableSet.of();
     }
 
+    RuleDefinitionIdentifier functionIdentifier = new RuleDefinitionIdentifier();
     // TODO(cparsons): Improve details given to RuleInfo (for example, attribute types).
-    ruleInfoList.add(new RuleInfo(ast.getLocation(), doc, attrNames));
-    return implementation;
+    ruleInfoList.add(new RuleInfo(functionIdentifier, ast.getLocation(), doc, attrNames));
+    return functionIdentifier;
   }
 
   @Override
@@ -100,4 +103,24 @@ public class FakeSkylarkRuleFunctionsApi implements SkylarkRuleFunctionsApi<File
       FuncallExpression ast, Environment funcallEnv) throws EvalException {
     return null;
   }
+
+  /**
+   * A fake {@link BaseFunction} implementation which serves as an identifier for a rule definition.
+   * A skylark invocation of 'rule()' should spawn a unique instance of this class and return it.
+   * Thus, skylark code such as 'foo = rule()' will result in 'foo' being assigned to a unique
+   * identifier, which can later be matched to a registered rule() invocation saved by the fake
+   * build API implementation.
+   */
+  private static class RuleDefinitionIdentifier extends BaseFunction {
+
+    public RuleDefinitionIdentifier() {
+      super("RuleDefinitionIdentifier");
+    }
+
+    @Override
+    public boolean equals(@Nullable Object other) {
+      // Use exact object matching.
+      return this == other;
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/skydoc/rendering/BUILD b/src/main/java/com/google/devtools/build/skydoc/rendering/BUILD
index 13f2cc9edd..0dc5f0827c 100644
--- a/src/main/java/com/google/devtools/build/skydoc/rendering/BUILD
+++ b/src/main/java/com/google/devtools/build/skydoc/rendering/BUILD
@@ -15,6 +15,7 @@ java_library(
     deps = [
         "//src/main/java/com/google/devtools/build/lib:events",
         "//src/main/java/com/google/devtools/build/lib:skylarkinterface",
+        "//src/main/java/com/google/devtools/build/lib:syntax",
         "//third_party:guava",
         "//third_party:jsr305",
     ],
diff --git a/src/main/java/com/google/devtools/build/skydoc/rendering/RuleInfo.java b/src/main/java/com/google/devtools/build/skydoc/rendering/RuleInfo.java
index 0d8a150dfc..dee5e0f94b 100644
--- a/src/main/java/com/google/devtools/build/skydoc/rendering/RuleInfo.java
+++ b/src/main/java/com/google/devtools/build/skydoc/rendering/RuleInfo.java
@@ -17,6 +17,7 @@ package com.google.devtools.build.skydoc.rendering;
 import com.google.common.base.Joiner;
 import com.google.common.base.Strings;
 import com.google.devtools.build.lib.events.Location;
+import com.google.devtools.build.lib.syntax.BaseFunction;
 import java.util.Collection;
 
 /**
@@ -24,16 +25,25 @@ import java.util.Collection;
  */
 public class RuleInfo {
 
+  private final BaseFunction identifierFunction;
   private final Location location;
   private final String docString;
   private final Collection<String> attrNames;
 
-  public RuleInfo(Location location, String docString, Collection<String> attrNames) {
+  public RuleInfo(BaseFunction identifierFunction,
+      Location location,
+      String docString,
+      Collection<String> attrNames) {
+    this.identifierFunction = identifierFunction;
     this.location = location;
     this.docString = docString;
     this.attrNames = attrNames;
   }
 
+  public BaseFunction getIdentifierFunction() {
+    return identifierFunction;
+  }
+
   public Location getLocation() {
     return location;
   }
diff --git a/src/test/java/com/google/devtools/build/skydoc/BUILD b/src/test/java/com/google/devtools/build/skydoc/BUILD
index d6daaa3398..e5a80fd54d 100644
--- a/src/test/java/com/google/devtools/build/skydoc/BUILD
+++ b/src/test/java/com/google/devtools/build/skydoc/BUILD
@@ -44,3 +44,10 @@ skydoc_test(
     input_file = "unknown_name_test/input.txt",
     skydoc = "//src/main/java/com/google/devtools/build/skydoc",
 )
+
+skydoc_test(
+    name = "multiple_rules_test",
+    golden_file = "multiple_rules_test/golden.txt",
+    input_file = "multiple_rules_test/input.txt",
+    skydoc = "//src/main/java/com/google/devtools/build/skydoc",
+)
diff --git a/src/test/java/com/google/devtools/build/skydoc/SkydocTest.java b/src/test/java/com/google/devtools/build/skydoc/SkydocTest.java
index d80fe4ed35..ad524c6df5 100644
--- a/src/test/java/com/google/devtools/build/skydoc/SkydocTest.java
+++ b/src/test/java/com/google/devtools/build/skydoc/SkydocTest.java
@@ -16,13 +16,17 @@ package com.google.devtools.build.skydoc;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.skylark.util.SkylarkTestCase;
 import com.google.devtools.build.lib.syntax.ParserInputSource;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.skydoc.rendering.RuleInfo;
-import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.stream.Collectors;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -56,12 +60,63 @@ public final class SkydocTest extends SkylarkTestCase {
     ParserInputSource parserInputSource =
         ParserInputSource.create(bytes, file.asFragment());
 
-    List<RuleInfo> ruleInfos = new SkydocMain().eval(parserInputSource);
+    ImmutableMap.Builder<String, RuleInfo> ruleInfoMap = ImmutableMap.builder();
+    ImmutableList.Builder<RuleInfo> unexportedRuleInfos = ImmutableList.builder();
+
+    new SkydocMain().eval(parserInputSource, ruleInfoMap, unexportedRuleInfos);
+    Map<String, RuleInfo> ruleInfos = ruleInfoMap.build();
     assertThat(ruleInfos).hasSize(1);
 
-    RuleInfo ruleInfo = Iterables.getOnlyElement(ruleInfos);
-    assertThat(ruleInfo.getDocString()).isEqualTo("This is my rule. It does stuff.");
-    assertThat(ruleInfo.getAttrNames()).containsExactly(
+    Entry<String, RuleInfo> ruleInfo = Iterables.getOnlyElement(ruleInfos.entrySet());
+    assertThat(ruleInfo.getKey()).isEqualTo("my_rule");
+    assertThat(ruleInfo.getValue().getDocString()).isEqualTo("This is my rule. It does stuff.");
+    assertThat(ruleInfo.getValue().getAttrNames()).containsExactly(
         "first", "second", "third", "fourth");
+    assertThat(unexportedRuleInfos.build()).isEmpty();
+  }
+
+  @Test
+  public void testMultipleRuleNames() throws Exception {
+    Path file =
+        scratch.file(
+            "/test/test.bzl",
+            "def rule_impl(ctx):",
+            "  return struct()",
+            "",
+            "rule_one = rule(",
+            "    doc = 'Rule one',",
+            "    implementation = rule_impl,",
+            ")",
+            "",
+            "rule(",
+            "    doc = 'This rule is not named',",
+            "    implementation = rule_impl,",
+            ")",
+            "",
+            "rule(",
+            "    doc = 'This rule also is not named',",
+            "    implementation = rule_impl,",
+            ")",
+            "",
+            "rule_two = rule(",
+            "    doc = 'Rule two',",
+            "    implementation = rule_impl,",
+            ")");
+    byte[] bytes = FileSystemUtils.readWithKnownFileSize(file, file.getFileSize());
+
+    ParserInputSource parserInputSource =
+        ParserInputSource.create(bytes, file.asFragment());
+
+    ImmutableMap.Builder<String, RuleInfo> ruleInfoMap = ImmutableMap.builder();
+    ImmutableList.Builder<RuleInfo> unexportedRuleInfos = ImmutableList.builder();
+
+    new SkydocMain().eval(parserInputSource, ruleInfoMap, unexportedRuleInfos);
+
+    assertThat(ruleInfoMap.build().keySet()).containsExactly("rule_one", "rule_two");
+
+    assertThat(unexportedRuleInfos.build().stream()
+            .map(ruleInfo -> ruleInfo.getDocString())
+            .collect(Collectors.toList()))
+        .containsExactly("This rule is not named", "This rule also is not named");
   }
 }
diff --git a/src/test/java/com/google/devtools/build/skydoc/multiple_rules_test/golden.txt b/src/test/java/com/google/devtools/build/skydoc/multiple_rules_test/golden.txt
new file mode 100644
index 0000000000..b1a4a66325
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/skydoc/multiple_rules_test/golden.txt
@@ -0,0 +1,12 @@
+my_rule
+This is my rule. It does stuff.
+first,second
+
+other_rule
+This is another rule.
+third,fourth
+
+yet_another_rule
+This is yet another rule
+fifth
+
diff --git a/src/test/java/com/google/devtools/build/skydoc/multiple_rules_test/input.txt b/src/test/java/com/google/devtools/build/skydoc/multiple_rules_test/input.txt
new file mode 100644
index 0000000000..b1c0538231
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/skydoc/multiple_rules_test/input.txt
@@ -0,0 +1,28 @@
+def my_rule_impl(ctx):
+    return struct()
+
+my_rule = rule(
+    implementation = my_rule_impl,
+    doc = "This is my rule. It does stuff.",
+    attrs = {
+        "first": attr.label(mandatory = True, allow_files = True, single_file = True),
+        "second": attr.string_dict(mandatory = True),
+    },
+)
+
+other_rule = rule(
+    implementation = my_rule_impl,
+    doc = "This is another rule.",
+    attrs = {
+        "third": attr.label(mandatory = True, allow_files = True, single_file = True),
+        "fourth": attr.string_dict(mandatory = True),
+    },
+)
+
+yet_another_rule = rule(
+    implementation = my_rule_impl,
+    doc = "This is yet another rule",
+    attrs = {
+        "fifth": attr.label(mandatory = True, allow_files = True, single_file = True),
+    },
+)
diff --git a/src/test/java/com/google/devtools/build/skydoc/simple_test/golden.txt b/src/test/java/com/google/devtools/build/skydoc/simple_test/golden.txt
index 0bf3a8c96a..50b71f076b 100644
--- a/src/test/java/com/google/devtools/build/skydoc/simple_test/golden.txt
+++ b/src/test/java/com/google/devtools/build/skydoc/simple_test/golden.txt
@@ -1,3 +1,4 @@
 my_rule
 This is my rule. It does stuff.
 first,second,third,fourth
+
diff --git a/src/test/java/com/google/devtools/build/skydoc/unknown_name_test/golden.txt b/src/test/java/com/google/devtools/build/skydoc/unknown_name_test/golden.txt
index 9b49c77015..ab79bcf5d6 100644
--- a/src/test/java/com/google/devtools/build/skydoc/unknown_name_test/golden.txt
+++ b/src/test/java/com/google/devtools/build/skydoc/unknown_name_test/golden.txt
@@ -1,2 +1,3 @@
 <unknown name>
 first,second,third,fourth
+