1 files changed, 4 insertions, 2 deletions
diff --git a/src/main/java/com/google/devtools/build/docgen/BuildDocCollector.java b/src/main/java/com/google/devtools/build/docgen/BuildDocCollector.java
index 7bb4f71993..870bf9558b 100644
--- a/src/main/java/com/google/devtools/build/docgen/BuildDocCollector.java
+++ b/src/main/java/com/google/devtools/build/docgen/BuildDocCollector.java
@@ -154,7 +154,7 @@ public class BuildDocCollector {
       if (ruleClass != null) {
         if (ruleClass.isDocumented()) {
           Class<? extends RuleDefinition> ruleDefinition =
-              ruleClassProvider.getRuleClassDefinition(ruleDoc.getRuleName());
+              ruleClassProvider.getRuleClassDefinition(ruleDoc.getRuleName()).getClass();
           for (Attribute attribute : ruleClass.getAttributes()) {
             String attrName = attribute.getName();
             List<RuleDocumentationAttribute> attributeDocList =
@@ -167,7 +167,9 @@ public class BuildDocCollector {
               int minLevel = Integer.MAX_VALUE;
               RuleDocumentationAttribute bestAttributeDoc = null;
               for (RuleDocumentationAttribute attributeDoc : attributeDocList) {
-                int level = attributeDoc.getDefinitionClassAncestryLevel(ruleDefinition);
+                int level = attributeDoc.getDefinitionClassAncestryLevel(
+                    ruleDefinition,
+                    ruleClassProvider);
                 if (level >= 0 && level < minLevel) {
                   bestAttributeDoc = attributeDoc;
                   minLevel = level;