Support passing in Buck2 config data

Summary:
Now we have `buck2 audit config --all-cells` (or will once we have D54535728) we can use that data to figure out the build files. That means that we will now support new cells as they get created, once everything gets wired together.

I'll do additional follow up once the Buck2 code has been landed for a while.

Reviewed By: Acesine

Differential Revision: D54542224

fbshipit-source-id: 4ee6dbe493993ef931b434063d2372f8dec84f90

btd/src/buck/cells.rs

@@ -94,6 +94,46 @@ impl CellInfo {
         Ok(Self { cells, paths })
     }
 
+    pub fn load_config_data(&mut self, file: &Path) -> anyhow::Result<()> {
+        let data = fs::read_to_string(file)
+            .with_context(|| format!("When reading `{}`", file.display()))?;
+        self.parse_config_data(&data)
+    }
+
+    pub fn parse_config_data(&mut self, data: &str) -> anyhow::Result<()> {
+        let json: HashMap<String, String> = serde_json::from_str(data)?;
+        // name_v2 needs to take precedence, so evaluate it second
+        for v2 in [false, true] {
+            let want_key = if v2 {
+                "buildfile.name_v2"
+            } else {
+                "buildfile.name"
+            };
+            for (k, v) in json.iter() {
+                // Expecting `cell//buildfile.name` = `BUCK,TARGETS`
+                if let Some((cell, key)) = k.split_once("//") {
+                    if key == want_key {
+                        let mut names = v
+                            .split(',')
+                            .map(|x| x.trim().to_owned())
+                            .collect::<Vec<_>>();
+                        if !v2 {
+                            // For name, we infer the .v2 suffix automatically
+                            names = names
+                                .into_iter()
+                                .flat_map(|x| [format!("{x}.v2"), x])
+                                .collect();
+                        }
+                        if let Some(data) = self.cells.get_mut(&CellName::new(cell)) {
+                            data.build_files = names;
+                        }
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
     pub fn resolve(&self, path: &CellPath) -> anyhow::Result<ProjectRelativePath> {
         match self.cells.get(&path.cell()) {
             Some(data) => Ok(data.path.join(path.path().as_str())),
@@ -170,4 +210,35 @@ mod tests {
 
         assert!(cells.resolve(&CellPath::new("missing//foo.txt")).is_err());
     }
+
+    #[test]
+    fn test_cell_config() {
+        let value = serde_json::json!(
+            {
+                "root": "/Users/ndmitchell/repo",
+                "cell1": "/Users/ndmitchell/repo/cell1",
+                "cell2": "/Users/ndmitchell/repo/cell2",
+              }
+        );
+        let mut cells = CellInfo::parse(&serde_json::to_string(&value).unwrap()).unwrap();
+        let value = serde_json::json!(
+            {
+                "cell1//buildfile.name":"BUCK",
+                "cell1//buildfile.name_v2":"TARGETS",
+                "cell2//buildfile.name":"A1,A2",
+            }
+        );
+        cells
+            .parse_config_data(&serde_json::to_string(&value).unwrap())
+            .unwrap();
+        assert_eq!(cells.build_files(&CellName::new("cell1")), &["TARGETS"]);
+        assert_eq!(
+            cells.build_files(&CellName::new("cell2")),
+            &["A1.v2", "A1", "A2.v2", "A2"]
+        );
+        assert_eq!(
+            cells.build_files(&CellName::new("cell3")),
+            &["BUCK.v2", "BUCK"]
+        );
+    }
 }

btd/src/buck/run.rs

@@ -104,6 +104,28 @@ impl Buck2 {
         Ok(String::from_utf8(res.stdout)?)
     }
 
+    pub fn audit_config(&mut self) -> anyhow::Result<String> {
+        let mut command = self.command();
+        command.args([
+            "audit",
+            "config",
+            "--json",
+            "--all-cells",
+            "buildfile.name",
+            "buildfile.name_v2",
+            "--reuse-current-config",
+        ]);
+        command.current_dir(self.root()?);
+        let res = with_command(command, |mut command| {
+            let res = command.output()?;
+            res.status.exit_ok().with_context(|| {
+                format!("Buck2 stderr: {}", String::from_utf8_lossy(&res.stderr))
+            })?;
+            Ok(res)
+        })?;
+        Ok(String::from_utf8(res.stdout)?)
+    }
+
     /// Does a package exist. Doesn't actually invoke Buck2, but does look at the file system.
     pub fn does_package_exist(&mut self, cells: &CellInfo, x: &Package) -> anyhow::Result<bool> {
         let root = self.root()?;

btd/src/lib.rs

@@ -63,6 +63,9 @@ use crate::output::OutputFormat;
 use crate::rerun::PackageStatus;
 use crate::sapling::status::read_status;
 
+/// Set to `true` once we have `buck2 audit config --all-cells` deployed.
+const BUCK2_ALL_CELLS_DEPLOYED: bool = false;
+
 /// Buck-based target determinator.
 #[derive(Parser)]
 pub struct Args {
@@ -71,6 +74,11 @@ pub struct Args {
     #[arg(long, value_name = "FILE")]
     cells: Option<PathBuf>,
 
+    /// File containing the output of `buck2 audit config --cells --json` in the root of the repo.
+    /// If the `cells` are empty this will run the Buck command to figure it out.
+    #[arg(long, value_name = "FILE")]
+    config: Option<PathBuf>,
+
     /// File containing the output of `hg status` for the relevant diff.
     #[arg(long, value_name = "FILE")]
     changes: PathBuf,
@@ -168,10 +176,19 @@ pub fn main(args: Args) -> anyhow::Result<()> {
     let step = |name| info!("Starting {} at {:.3}s", name, t.elapsed().as_secs_f64());
 
     step("reading cells");
-    let cells = match &args.cells {
+    let mut cells = match &args.cells {
         Some(file) => CellInfo::new(file)?,
         None => CellInfo::parse(&buck2.cells()?)?,
     };
+    step("reading config");
+    match &args.config {
+        Some(file) => cells.load_config_data(file)?,
+        None if args.cells.is_none() && BUCK2_ALL_CELLS_DEPLOYED => {
+            cells.parse_config_data(&buck2.audit_config()?)?
+        }
+        _ => (), // We don't auto fill in config data if the user has explicit cells
+    }
+
     step("reading changes");
     let changes = Changes::new(&cells, read_status(&args.changes)?)?;
     step("reading base");