WorkspaceASTFunction returns a list of ASTs so we can split the AST before load statements

Issue #824 Step 2.

--
MOS_MIGRATED_REVID=113986176

1 files changed, 32 insertions, 1 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTFunction.java
index d1077d2260..46f5d0d9d1 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/WorkspaceASTFunction.java
@@ -14,9 +14,12 @@
 
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.packages.RuleClassProvider;
 import com.google.devtools.build.lib.syntax.BuildFileAST;
+import com.google.devtools.build.lib.syntax.LoadStatement;
 import com.google.devtools.build.lib.syntax.ParserInputSource;
+import com.google.devtools.build.lib.syntax.Statement;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.RootedPath;
@@ -27,6 +30,7 @@ import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 
 import java.io.IOException;
+import java.util.List;
 
 /**
  * A SkyFunction to parse WORKSPACE files into a BuildFileAST.
@@ -65,12 +69,39 @@ public class WorkspaceASTFunction implements SkyFunction {
               new IOException("Failed to parse WORKSPACE file"), Transience.PERSISTENT);
         }
       }
-      return new WorkspaceASTValue(ast);
+      return new WorkspaceASTValue(splitAST(ast));
     } catch (IOException ex) {
       throw new WorkspaceASTFunctionException(ex, Transience.TRANSIENT);
     }
   }
 
+  /**
+   * Cut {@code ast} into a list of AST separated by load statements. We cut right before each load
+   * statement series.
+   */
+  private static ImmutableList<BuildFileAST> splitAST(BuildFileAST ast) {
+    ImmutableList.Builder<BuildFileAST> asts = ImmutableList.builder();
+    int prevIdx = 0;
+    boolean lastIsLoad = true; // don't cut if the first statement is a load.
+    List<Statement> statements = ast.getStatements();
+    for (int idx = 0; idx < statements.size(); idx++) {
+      Statement st = statements.get(idx);
+      if (st instanceof LoadStatement) {
+        if (!lastIsLoad) {
+          asts.add(ast.subTree(prevIdx, idx));
+          prevIdx = idx;
+        }
+        lastIsLoad = true;
+      } else {
+        lastIsLoad = false;
+      }
+    }
+    if (!statements.isEmpty()) {
+      asts.add(ast.subTree(prevIdx, statements.size()));
+    }
+    return asts.build();
+  }
+
   private static final class WorkspaceASTFunctionException extends SkyFunctionException {
     public WorkspaceASTFunctionException(Exception e, Transience transience) {
       super(e, transience);