bxl: expose allbuildfiles() and rbuildfiles()

app/buck2_build_api/src/query/bxl.rs

@@ -77,6 +77,17 @@ pub trait BxlCqueryFunctions: Send {
         dice: &mut DiceComputations<'_>,
         targets: &TargetSet<ConfiguredTargetNode>,
     ) -> buck2_error::Result<Vec<MaybeCompatible<ConfiguredTargetNode>>>;
+    async fn allbuildfiles(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        universe: &TargetSet<ConfiguredTargetNode>,
+    ) -> buck2_error::Result<FileSet>;
+    async fn rbuildfiles(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        universe: &FileSet,
+        argset: &FileSet,
+    ) -> buck2_error::Result<FileSet>;
 }
 
 #[async_trait]
@@ -125,6 +136,17 @@ pub trait BxlUqueryFunctions: Send {
         dice: &mut DiceComputations<'_>,
         file_set: &FileSet,
     ) -> buck2_error::Result<TargetSet<TargetNode>>;
+    async fn allbuildfiles(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        universe: &TargetSet<TargetNode>,
+    ) -> buck2_error::Result<FileSet>;
+    async fn rbuildfiles(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        universe: &FileSet,
+        argset: &FileSet,
+    ) -> buck2_error::Result<FileSet>;
 }
 
 #[async_trait]

app/buck2_bxl/src/bxl/starlark_defs/cquery.rs

@@ -607,4 +607,68 @@ fn cquery_methods(builder: &mut MethodsBuilder) {
             })
             .map(StarlarkFileSet::from)?)
     }
+
+    /// The allbuildfiles query for finding the build files of the given targets and their transitive imports.
+    ///
+    /// Sample usage:
+    /// ```python
+    /// def _impl_allbuildfiles(ctx):
+    ///     result = ctx.cquery().allbuildfiles("root//bin:the_binary")
+    ///     ctx.output.print(result)
+    /// ```
+    fn allbuildfiles<'v>(
+        this: &StarlarkCQueryCtx<'v>,
+        universe: ConfiguredTargetListExprArg<'v>,
+        eval: &mut Evaluator<'v, '_, '_>,
+    ) -> starlark::Result<StarlarkFileSet> {
+        Ok(this
+            .ctx
+            .via_dice(eval, |dice| {
+                dice.via(|dice| {
+                    async {
+                        let universe = unpack_targets(this, dice, universe).await?;
+                        get_cquery_env(&this.ctx, &this.global_cfg_options_override)
+                            .await?
+                            .allbuildfiles(dice, &universe)
+                            .await
+                    }
+                    .boxed_local()
+                })
+            })
+            .map(StarlarkFileSet::from)?)
+    }
+
+    /// The rbuildfiles query for finding all build files that transitively depend on the given files.
+    ///
+    /// Sample usage:
+    /// ```python
+    /// def _impl_rbuildfiles(ctx):
+    ///     result = ctx.cquery().rbuildfiles("bin/TARGETS", "bin/defs.bzl")
+    ///     ctx.output.print(result)
+    /// ```
+    fn rbuildfiles<'v>(
+        this: &StarlarkCQueryCtx<'v>,
+        universe: FileSetExpr,
+        argset: FileSetExpr,
+        eval: &mut Evaluator<'v, '_, '_>,
+    ) -> starlark::Result<StarlarkFileSet> {
+        Ok(this
+            .ctx
+            .via_dice(eval, |dice| {
+                dice.via(|dice| {
+                    async {
+                        get_cquery_env(&this.ctx, &this.global_cfg_options_override)
+                            .await?
+                            .rbuildfiles(
+                                dice,
+                                (universe.get(&this.ctx).await?).as_ref(),
+                                (argset.get(&this.ctx).await?).as_ref(),
+                            )
+                            .await
+                    }
+                    .boxed_local()
+                })
+            })
+            .map(StarlarkFileSet::from)?)
+    }
 }

app/buck2_bxl/src/bxl/starlark_defs/lazy_ctx/lazy_uquery_ctx.rs

@@ -263,6 +263,38 @@ fn lazy_uquery_methods(builder: &mut MethodsBuilder) {
         Ok(StarlarkLazy::new_uquery(op))
     }
 
+    /// Finds the build files of the given targets and their transitive imports.
+    ///
+    /// Example:
+    /// ```python
+    /// res = ctx.lazy.uquery().allbuildfiles("//:foo").catch().resolve()
+    /// ```
+    fn allbuildfiles<'v>(
+        #[starlark(this)] _this: &'v StarlarkLazyUqueryCtx,
+        #[starlark(require = pos)] universe: TargetListExprArg<'v>,
+    ) -> starlark::Result<StarlarkLazy> {
+        let universe = OwnedTargetListExprArg::from_ref(&universe);
+        let op = LazyUqueryOperation::AllBuildfiles(universe);
+        Ok(StarlarkLazy::new_uquery(op))
+    }
+
+    /// Finds all build files that transitively depend on the given files.
+    ///
+    /// Example:
+    /// ```python
+    /// res = ctx.lazy.uquery().rbuildfiles("bin/TARGETS", "bin/defs.bzl").catch().resolve()
+    /// ```
+    fn rbuildfiles<'v>(
+        #[starlark(this)] _this: &'v StarlarkLazyUqueryCtx,
+        #[starlark(require = pos)] universe: FileSetExpr<'v>,
+        #[starlark(require = pos)] argset: FileSetExpr<'v>,
+    ) -> starlark::Result<StarlarkLazy> {
+        let universe = OwnedFileSetExpr::from_ref(&universe);
+        let argset = OwnedFileSetExpr::from_ref(&argset);
+        let op = LazyUqueryOperation::Rbuildfiles { universe, argset };
+        Ok(StarlarkLazy::new_uquery(op))
+    }
+
     /// Finds targets that own the specified files.
     ///
     /// Example:

app/buck2_bxl/src/bxl/starlark_defs/lazy_ctx/operation/uquery.rs

@@ -71,6 +71,11 @@ pub(crate) enum LazyUqueryOperation {
         targets: OwnedTargetListExprArg,
     },
     Buildfile(OwnedTargetListExprArg),
+    AllBuildfiles(OwnedTargetListExprArg),
+    Rbuildfiles {
+        universe: OwnedFileSetExpr,
+        argset: OwnedFileSetExpr,
+    },
     Owner {
         files: OwnedFileSetExpr,
     },
@@ -95,6 +100,8 @@ pub(crate) enum LazyUqueryResult {
     Rdeps(StarlarkTargetSet<TargetNode>),
     Filter(StarlarkTargetSet<TargetNode>),
     Buildfile(StarlarkFileSet),
+    AllBuildfiles(StarlarkFileSet),
+    Rbuildfiles(StarlarkFileSet),
     Owner(StarlarkTargetSet<TargetNode>),
     TargetsInBuildfile(StarlarkTargetSet<TargetNode>),
     Eval(QueryEvaluationResult<TargetNode>),
@@ -114,6 +121,8 @@ impl LazyUqueryResult {
             LazyUqueryResult::Rdeps(target_set) => Ok(heap.alloc(target_set)),
             LazyUqueryResult::Filter(target_set) => Ok(heap.alloc(target_set)),
             LazyUqueryResult::Buildfile(file_set) => Ok(heap.alloc(file_set)),
+            LazyUqueryResult::AllBuildfiles(file_set) => Ok(heap.alloc(file_set)),
+            LazyUqueryResult::Rbuildfiles(file_set) => Ok(heap.alloc(file_set)),
             LazyUqueryResult::Owner(target_set) => Ok(heap.alloc(target_set)),
             LazyUqueryResult::TargetsInBuildfile(target_set) => Ok(heap.alloc(target_set)),
             LazyUqueryResult::Eval(result) => parse_query_evaluation_result(result, heap),
@@ -258,6 +267,27 @@ impl LazyUqueryOperation {
 
                 Ok(LazyUqueryResult::Buildfile(StarlarkFileSet::from(res)))
             }
+            LazyUqueryOperation::AllBuildfiles(expr) => {
+                let target_set = expr.to_unconfigured_target_set(core_data, dice).await?;
+
+                let res = get_uquery_env(core_data)
+                    .await?
+                    .allbuildfiles(dice, &target_set)
+                    .await?;
+
+                Ok(LazyUqueryResult::AllBuildfiles(StarlarkFileSet::from(res)))
+            }
+            LazyUqueryOperation::Rbuildfiles { universe, argset } => {
+                let universe = universe.get(core_data)?;
+                let argset = argset.get(core_data)?;
+
+                let res = get_uquery_env(core_data)
+                    .await?
+                    .rbuildfiles(dice, &universe, &argset)
+                    .await?;
+
+                Ok(LazyUqueryResult::Rbuildfiles(StarlarkFileSet::from(res)))
+            }
             LazyUqueryOperation::Owner { files } => {
                 let file_set = files.get(core_data)?;
 

app/buck2_bxl/src/bxl/starlark_defs/uquery.rs

@@ -490,6 +490,70 @@ fn uquery_methods(builder: &mut MethodsBuilder) {
             .map(StarlarkTargetSet::from)?)
     }
 
+    /// The allbuildfiles query for finding the build files of the given targets and their transitive imports.
+    ///
+    /// Sample usage:
+    /// ```python
+    /// def _impl_allbuildfiles(ctx):
+    ///     result = ctx.uquery().allbuildfiles("root//bin:the_binary")
+    ///     ctx.output.print(result)
+    /// ```
+    fn allbuildfiles<'v>(
+        this: &StarlarkUQueryCtx<'v>,
+        universe: TargetListExprArg<'v>,
+        eval: &mut Evaluator<'v, '_, '_>,
+    ) -> starlark::Result<StarlarkFileSet> {
+        Ok(this
+            .ctx
+            .via_dice(eval, |dice| {
+                dice.via(|dice| {
+                    async {
+                        let universe = unpack_targets(this, dice, universe).await?;
+                        get_uquery_env(&this.ctx)
+                            .await?
+                            .allbuildfiles(dice, &universe)
+                            .await
+                    }
+                    .boxed_local()
+                })
+            })
+            .map(StarlarkFileSet::from)?)
+    }
+
+    /// The rbuildfiles query for finding all build files that transitively depend on the given files.
+    ///
+    /// Sample usage:
+    /// ```python
+    /// def _impl_rbuildfiles(ctx):
+    ///     result = ctx.uquery().rbuildfiles("bin/TARGETS", "bin/defs.bzl")
+    ///     ctx.output.print(result)
+    /// ```
+    fn rbuildfiles<'v>(
+        this: &StarlarkUQueryCtx<'v>,
+        universe: FileSetExpr,
+        argset: FileSetExpr,
+        eval: &mut Evaluator<'v, '_, '_>,
+    ) -> starlark::Result<StarlarkFileSet> {
+        Ok(this
+            .ctx
+            .via_dice(eval, |dice| {
+                dice.via(|dice| {
+                    async {
+                        get_uquery_env(&this.ctx)
+                            .await?
+                            .rbuildfiles(
+                                dice,
+                                (universe.get(&this.ctx).await?).as_ref(),
+                                (argset.get(&this.ctx).await?).as_ref(),
+                            )
+                            .await
+                    }
+                    .boxed_local()
+                })
+            })
+            .map(StarlarkFileSet::from)?)
+    }
+
     /// The attrregexfilter query for rule attribute filtering with regex.
     ///
     /// Sample usage:

app/buck2_query_impls/src/cquery/bxl.rs

@@ -244,6 +244,46 @@ impl BxlCqueryFunctions for BxlCqueryFunctionsImpl {
             })
             .await?)
     }
+
+    async fn allbuildfiles(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        universe: &TargetSet<ConfiguredTargetNode>,
+    ) -> buck2_error::Result<FileSet> {
+        Ok(dice
+            .with_linear_recompute(|dice| async move {
+                cquery_functions()
+                    .allbuildfiles(
+                        &self
+                            .cquery_env(&self.setup_dice_query_delegate(&dice).await?, None)
+                            .await?,
+                        universe,
+                    )
+                    .await
+            })
+            .await?)
+    }
+
+    async fn rbuildfiles(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        universe: &FileSet,
+        argset: &FileSet,
+    ) -> buck2_error::Result<FileSet> {
+        Ok(dice
+            .with_linear_recompute(|dice| async move {
+                cquery_functions()
+                    .rbuildfiles(
+                        &self
+                            .cquery_env(&self.setup_dice_query_delegate(&dice).await?, None)
+                            .await?,
+                        universe,
+                        argset,
+                    )
+                    .await
+            })
+            .await?)
+    }
 }
 
 pub(crate) fn init_new_bxl_cquery_functions() {

app/buck2_query_impls/src/uquery/bxl.rs

@@ -208,6 +208,40 @@ impl BxlUqueryFunctions for BxlUqueryFunctionsImpl {
             })
             .await?)
     }
+    async fn allbuildfiles(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        universe: &TargetSet<TargetNode>,
+    ) -> buck2_error::Result<FileSet> {
+        Ok(dice
+            .with_linear_recompute(|dice| async move {
+                uquery_functions()
+                    .allbuildfiles(
+                        &self.uquery_env(&self.uquery_delegate(&dice).await?).await?,
+                        universe,
+                    )
+                    .await
+            })
+            .await?)
+    }
+    async fn rbuildfiles(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        universe: &FileSet,
+        argset: &FileSet,
+    ) -> buck2_error::Result<FileSet> {
+        Ok(dice
+            .with_linear_recompute(|dice| async move {
+                uquery_functions()
+                    .rbuildfiles(
+                        &self.uquery_env(&self.uquery_delegate(&dice).await?).await?,
+                        universe,
+                        argset,
+                    )
+                    .await
+            })
+            .await?)
+    }
 }
 
 pub(crate) fn init_new_bxl_uquery_functions() {