Allow passing location, ast, and environment to @SkylarkCallable methods

RELNOTES: None.
PiperOrigin-RevId: 188201686

2 files changed, 152 insertions, 17 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skylarkinterface/SkylarkCallable.java b/src/main/java/com/google/devtools/build/lib/skylarkinterface/SkylarkCallable.java
index 1b193fe1f1..3ce0f70d9d 100644
--- a/src/main/java/com/google/devtools/build/lib/skylarkinterface/SkylarkCallable.java
+++ b/src/main/java/com/google/devtools/build/lib/skylarkinterface/SkylarkCallable.java
@@ -20,6 +20,22 @@ import java.lang.annotation.Target;
 
 /**
  * A marker interface for Java methods which can be called from Skylark.
+ *
+ * <p>Methods annotated with this annotation are expected to meet certain requirements which
+ * are enforced by an annotation processor:</p>
+ * <ul>
+ * <li>The method must be public.</li>
+ * <li>If structField=true, there must be zero user-supplied parameters.</li>
+ * <li>Method parameters must be supplied in the following order:
+ *   <pre>method([positionals]*[other user-args](Location)(FuncallExpression)(Envrionment))</pre>
+ *   where Location, FuncallExpression, and Environment are supplied by the interpreter if and
+ *   only if useLocation, useAst, and useEnvironment are specified, respectively.
+*  </li>
+ * <li>
+ *   The number of method parameters much match the number of annotation-declared parameters
+ *   plus the number of interpreter-supplied parameters.
+ * </li>
+ * </ul>
  */
 @Target({ElementType.METHOD})
 @Retention(RetentionPolicy.RUNTIME)
@@ -70,4 +86,25 @@ public @interface SkylarkCallable {
    * <code>None</code>). If not set and the Java method returns null, an error will be raised.
    */
   boolean allowReturnNones() default false;
+
+  /**
+   * If true, the location of the call site will be passed as an argument of the annotated function.
+   * (Thus, the annotated method signature must contain Location as a parameter. See the
+   * interface-level javadoc for details.)
+   */
+  boolean useLocation() default false;
+
+  /**
+   * If true, the AST of the call site will be passed as an argument of the annotated function.
+   * (Thus, the annotated method signature must contain FuncallExpression as a parameter. See the
+   * interface-level javadoc for details.)
+   */
+  boolean useAst() default false;
+
+  /**
+   * If true, the Skylark Environment will be passed as an argument of the annotated function.
+   * (Thus, the annotated method signature must contain Environment as a parameter. See the
+   * interface-level javadoc for details.)
+   */
+  boolean useEnvironment() default false;
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skylarkinterface/processor/SkylarkCallableProcessor.java b/src/main/java/com/google/devtools/build/lib/skylarkinterface/processor/SkylarkCallableProcessor.java
index cc98597932..cd4b288db1 100644
--- a/src/main/java/com/google/devtools/build/lib/skylarkinterface/processor/SkylarkCallableProcessor.java
+++ b/src/main/java/com/google/devtools/build/lib/skylarkinterface/processor/SkylarkCallableProcessor.java
@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.skylarkinterface.processor;
 
 import com.google.devtools.build.lib.skylarkinterface.SkylarkCallable;
+import java.util.List;
 import java.util.Set;
 import javax.annotation.processing.AbstractProcessor;
 import javax.annotation.processing.Messager;
@@ -27,6 +28,7 @@ import javax.lang.model.element.Element;
 import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
+import javax.lang.model.element.VariableElement;
 import javax.tools.Diagnostic;
 
 /**
@@ -35,8 +37,16 @@ import javax.tools.Diagnostic;
  * <p>Checks the following invariants about {@link SkylarkCallable}-annotated methods:
  * <ul>
  * <li>The method must be public.</li>
- * <li>The number of method parameters much match the number of annotation-declared parameters.</li>
- * <li>If structField=true, there must be zero arguments.</li>
+ * <li>If structField=true, there must be zero user-supplied parameters.</li>
+ * <li>Method parameters must be supplied in the following order:
+ *   <pre>method([positionals]*[other user-args](Location)(FuncallExpression)(Envrionment))</pre>
+ *   where Location, FuncallExpression, and Environment are supplied by the interpreter if and
+ *   only if useLocation, useAst, and useEnvironment are specified, respectively.
+ *  </li>
+ * <li>
+ *   The number of method parameters much match the number of annotation-declared parameters
+ *   plus the number of interpreter-supplied parameters.
+ * </li>
  * </ul>
  *
  * <p>These properties can be relied upon at runtime without additional checks.
@@ -47,6 +57,10 @@ public final class SkylarkCallableProcessor extends AbstractProcessor {
 
   private Messager messager;
 
+  private static final String LOCATION = "com.google.devtools.build.lib.events.Location";
+  private static final String AST = "com.google.devtools.build.lib.syntax.FuncallExpression";
+  private static final String ENVIRONMENT = "com.google.devtools.build.lib.syntax.Environment";
+
   @Override
   public synchronized void init(ProcessingEnvironment processingEnv) {
     super.init(processingEnv);
@@ -56,6 +70,7 @@ public final class SkylarkCallableProcessor extends AbstractProcessor {
   @Override
   public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
     for (Element element : roundEnv.getElementsAnnotatedWith(SkylarkCallable.class)) {
+
       // Only methods are annotated with SkylarkCallable. This is verified by the
       // @Target(ElementType.METHOD) annotation.
       ExecutableElement methodElement = (ExecutableElement) element;
@@ -64,32 +79,115 @@ public final class SkylarkCallableProcessor extends AbstractProcessor {
       if (!methodElement.getModifiers().contains(Modifier.PUBLIC)) {
         error(methodElement, "@SkylarkCallable annotated methods must be public.");
       }
-      if (annotation.parameters().length > 0 || annotation.mandatoryPositionals() >= 0) {
-         int numDeclaredArgs = annotation.parameters().length
-             + Math.max(0, annotation.mandatoryPositionals());
-        if (methodElement.getParameters().size() != numDeclaredArgs) {
-          error(methodElement, String.format(
-              "@SkylarkCallable annotated method has %d parameters, but annotation declared %d.",
-              methodElement.getParameters().size(), numDeclaredArgs));
-        }
-      }
-      if (annotation.structField()) {
-        if (!methodElement.getParameters().isEmpty()) {
-          error(methodElement,
-              "@SkylarkCallable annotated methods with structField=true must have zero arguments.");
-        }
+
+      try {
+        verifyNumberOfParameters(methodElement, annotation);
+        verifyExtraInterpreterParams(methodElement, annotation);
+      } catch (SkylarkCallableProcessorException exception) {
+        error(exception.methodElement, exception.errorMessage);
       }
     }
+
     return true;
   }
 
+  private void verifyNumberOfParameters(ExecutableElement methodElement, SkylarkCallable annotation)
+      throws SkylarkCallableProcessorException {
+    List<? extends VariableElement> methodSignatureParams = methodElement.getParameters();
+    int numExtraInterpreterParams = numExpectedExtraInterpreterParams(annotation);
+
+    if (annotation.parameters().length > 0 || annotation.mandatoryPositionals() >= 0) {
+      int numDeclaredArgs =
+          annotation.parameters().length + Math.max(0, annotation.mandatoryPositionals());
+      if (methodSignatureParams.size() != numDeclaredArgs + numExtraInterpreterParams) {
+        throw new SkylarkCallableProcessorException(
+            methodElement,
+            String.format(
+                "@SkylarkCallable annotated method has %d parameters, but annotation declared "
+                    + "%d user-supplied parameters and %d extra interpreter parameters.",
+                methodSignatureParams.size(), numDeclaredArgs, numExtraInterpreterParams));
+      }
+    }
+    if (annotation.structField()) {
+      if (methodSignatureParams.size() > 0) {
+        // TODO(cparsons): Allow structField methods to accept interpreter-supplied arguments.
+        throw new SkylarkCallableProcessorException(
+            methodElement,
+            "@SkylarkCallable annotated methods with structField=true must have zero arguments.");
+      }
+    }
+  }
+
+  private void verifyExtraInterpreterParams(ExecutableElement methodElement,
+      SkylarkCallable annotation) throws SkylarkCallableProcessorException {
+    List<? extends VariableElement> methodSignatureParams = methodElement.getParameters();
+    int currentIndex = methodSignatureParams.size() - numExpectedExtraInterpreterParams(annotation);
+
+    // TODO(cparsons): Matching by class name alone is somewhat brittle, but due to tangled
+    // dependencies, it is difficult for this processor to depend directy on the expected
+    // classes here.
+    if (annotation.useLocation()) {
+      if (!LOCATION.equals(methodSignatureParams.get(currentIndex).asType().toString())) {
+        throw new SkylarkCallableProcessorException(
+            methodElement,
+            String.format(
+                "Expected parameter index %d to be the %s type, matching useLocation, but was %s",
+                currentIndex,
+                LOCATION,
+                methodSignatureParams.get(currentIndex).asType().toString()));
+      }
+      currentIndex++;
+    }
+    if (annotation.useAst()) {
+      if (!AST.equals(methodSignatureParams.get(currentIndex).asType().toString())) {
+        throw new SkylarkCallableProcessorException(
+            methodElement,
+            String.format(
+                "Expected parameter index %d to be the %s type, matching useAst, but was %s",
+                currentIndex, AST, methodSignatureParams.get(currentIndex).asType().toString()));
+      }
+      currentIndex++;
+    }
+    if (annotation.useEnvironment()) {
+      if (!ENVIRONMENT.equals(methodSignatureParams.get(currentIndex).asType().toString())) {
+        throw new SkylarkCallableProcessorException(
+            methodElement,
+            String.format(
+                "Expected parameter index %d to be the %s type, matching useEnvironment, "
+                    + "but was %s",
+                currentIndex,
+                ENVIRONMENT,
+                methodSignatureParams.get(currentIndex).asType().toString()));
+      }
+    }
+  }
+
+  private int numExpectedExtraInterpreterParams(SkylarkCallable annotation) {
+    int numExtraInterpreterParams = 0;
+    numExtraInterpreterParams += annotation.useLocation() ? 1 : 0;
+    numExtraInterpreterParams += annotation.useAst() ? 1 : 0;
+    numExtraInterpreterParams += annotation.useEnvironment() ? 1 : 0;
+    return numExtraInterpreterParams;
+  }
+
   /**
    * Prints an error message & fails the compilation.
    *
    * @param e The element which has caused the error. Can be null
    * @param msg The error message
    */
-  public void error(Element e, String msg) {
+  private void error(Element e, String msg) {
     messager.printMessage(Diagnostic.Kind.ERROR, msg, e);
   }
+
+  private static class SkylarkCallableProcessorException extends Exception {
+    private final ExecutableElement methodElement;
+    private final String errorMessage;
+
+    private SkylarkCallableProcessorException(
+        ExecutableElement methodElement, String errorMessage) {
+      this.methodElement = methodElement;
+      this.errorMessage = errorMessage;
+    }
+  }
 }