commit_graph: return boundary changesets in ancestors_within_distance

Summary: Currently there's an `ancestors_within_distance` method that returns a stream of all ancestors of a set of changesets that are within a given distance `max_distance`. This diff renames this method into `ancestors_within_distance_stream` and makes it return for each changeset the minimum distance it takes to reach it alongside its changeset id. Another non-stream version is added `ancestors_within_distance` that returns the ancestors as well as all boundary changesets, which are changesets for which the minimum distance to reach them is exactly the given `max_distance`.

Reviewed By: mitrandir77

Differential Revision: D57280689

fbshipit-source-id: 3ed53049ab3b24e44482160417222f70449b93e8

eden/mononoke/repo_attributes/commit_graph/commit_graph/src/lib.rs

@@ -20,6 +20,7 @@ use anyhow::Result;
 use borrowed::borrowed;
 use commit_graph_types::edges::ChangesetNode;
 pub use commit_graph_types::edges::ChangesetParents;
+use commit_graph_types::frontier::AncestorsWithinDistance;
 use commit_graph_types::frontier::ChangesetFrontierWithinDistance;
 use commit_graph_types::segments::BoundaryChangesets;
 use commit_graph_types::segments::ChangesetSegment;
@@ -251,15 +252,51 @@ impl CommitGraph {
             .await
     }
 
-    /// Returns all ancestors of any changeset in `heads` that's reachable
-    /// by taking no more than `distance` edges from some changeset in `heads`.
+    /// Returns all ancestors of any changeset in `heads` that are reachable
+    /// by taking no more than `max_distance` edges from some changeset in `heads`,
+    /// as well as the boundary changesets which are the changesets for which
+    /// the shortest distance to them is exactly `max_distance`.
     pub async fn ancestors_within_distance(
         &self,
         ctx: &CoreContext,
         heads: Vec<ChangesetId>,
-        distance: u64,
-    ) -> Result<BoxStream<'static, Result<ChangesetId>>> {
-        let frontier = self.frontier_within_distance(ctx, heads, distance).await?;
+        max_distance: u64,
+    ) -> Result<AncestorsWithinDistance> {
+        let cs_id_and_distance = self
+            .ancestors_within_distance_stream(ctx, heads, max_distance)
+            .await?
+            .try_collect::<Vec<_>>()
+            .await?;
+
+        let mut boundaries = vec![];
+        for (cs_id, distance) in &cs_id_and_distance {
+            if *distance == max_distance {
+                boundaries.push(*cs_id);
+            }
+        }
+
+        Ok(AncestorsWithinDistance {
+            ancestors: cs_id_and_distance
+                .into_iter()
+                .map(|(cs_id, _)| cs_id)
+                .collect(),
+            boundaries,
+        })
+    }
+
+    /// Returns a stream of all ancestors of any changeset in `heads` that are
+    /// reachable by taking no more than `max_distance` edges from some changeset
+    /// in `heads`. For each changeset we returns its changeset id alongside
+    /// the minimum distance it takes to reach the changeset.
+    pub async fn ancestors_within_distance_stream(
+        &self,
+        ctx: &CoreContext,
+        heads: Vec<ChangesetId>,
+        max_distance: u64,
+    ) -> Result<BoxStream<'static, Result<(ChangesetId, u64)>>> {
+        let frontier = self
+            .frontier_within_distance(ctx, heads, max_distance)
+            .await?;
 
         struct AncestorsWithinDistanceState {
             commit_graph: CommitGraph,
@@ -282,8 +319,8 @@ impl CommitGraph {
 
                 if let Some((_generation, cs_ids_and_remaining_distance)) = frontier.pop_last() {
                     let output_cs_ids = cs_ids_and_remaining_distance
-                        .keys()
-                        .copied()
+                        .iter()
+                        .map(|(cs_id, remaining_distance)| (*cs_id, max_distance - *remaining_distance))
                         .collect::<Vec<_>>();
 
                     let max_remaining_distance = cs_ids_and_remaining_distance.values().copied()

eden/mononoke/repo_attributes/commit_graph/commit_graph_testlib/src/lib.rs

@@ -1477,58 +1477,130 @@ pub async fn test_ancestors_within_distance(
     .await?;
     storage.flush();
 
-    assert_ancestors_within_distance(&ctx, &graph, vec!["N"], 0, vec!["N"]).await?;
-    assert_ancestors_within_distance(&ctx, &graph, vec!["N"], 1, vec!["N", "L", "K"]).await?;
-    assert_ancestors_within_distance(&ctx, &graph, vec!["N"], 2, vec!["N", "L", "K", "E", "J"])
-        .await?;
+    assert_ancestors_within_distance(&ctx, &graph, vec!["N"], 0, vec![("N", 0)]).await?;
+    assert_ancestors_within_distance(
+        &ctx,
+        &graph,
+        vec!["N"],
+        1,
+        vec![("N", 0), ("L", 1), ("K", 1)],
+    )
+    .await?;
+    assert_ancestors_within_distance(
+        &ctx,
+        &graph,
+        vec!["N"],
+        2,
+        vec![("N", 0), ("L", 1), ("K", 1), ("E", 2), ("J", 2)],
+    )
+    .await?;
     assert_ancestors_within_distance(
         &ctx,
         &graph,
         vec!["N"],
         3,
-        vec!["N", "L", "K", "E", "J", "I", "D"],
+        vec![
+            ("N", 0),
+            ("L", 1),
+            ("K", 1),
+            ("E", 2),
+            ("J", 2),
+            ("I", 3),
+            ("D", 3),
+        ],
     )
     .await?;
     assert_ancestors_within_distance(
         &ctx,
         &graph,
         vec!["N"],
         4,
-        vec!["N", "L", "K", "E", "J", "I", "D", "F", "C"],
+        vec![
+            ("N", 0),
+            ("L", 1),
+            ("K", 1),
+            ("E", 2),
+            ("J", 2),
+            ("I", 3),
+            ("D", 3),
+            ("F", 4),
+            ("C", 4),
+        ],
     )
     .await?;
     assert_ancestors_within_distance(
         &ctx,
         &graph,
         vec!["N"],
         5,
-        vec!["N", "L", "K", "E", "J", "I", "D", "F", "C", "B"],
+        vec![
+            ("N", 0),
+            ("L", 1),
+            ("K", 1),
+            ("E", 2),
+            ("J", 2),
+            ("I", 3),
+            ("D", 3),
+            ("F", 4),
+            ("C", 4),
+            ("B", 5),
+        ],
     )
     .await?;
     assert_ancestors_within_distance(
         &ctx,
         &graph,
         vec!["N"],
         6,
-        vec!["N", "L", "K", "E", "J", "I", "D", "F", "C", "B", "A"],
+        vec![
+            ("N", 0),
+            ("L", 1),
+            ("K", 1),
+            ("E", 2),
+            ("J", 2),
+            ("I", 3),
+            ("D", 3),
+            ("F", 4),
+            ("C", 4),
+            ("B", 5),
+            ("A", 6),
+        ],
     )
     .await?;
     assert_ancestors_within_distance(
         &ctx,
         &graph,
         vec!["N"],
         100,
-        vec!["N", "L", "K", "E", "J", "I", "D", "F", "C", "B", "A"],
+        vec![
+            ("N", 0),
+            ("L", 1),
+            ("K", 1),
+            ("E", 2),
+            ("J", 2),
+            ("I", 3),
+            ("D", 3),
+            ("F", 4),
+            ("C", 4),
+            ("B", 5),
+            ("A", 6),
+        ],
+    )
+    .await?;
+    assert_ancestors_within_distance(
+        &ctx,
+        &graph,
+        vec!["N", "M", "H"],
+        0,
+        vec![("N", 0), ("M", 0), ("H", 0)],
     )
     .await?;
-    assert_ancestors_within_distance(&ctx, &graph, vec!["N", "M", "H"], 0, vec!["N", "M", "H"])
-        .await?;
     assert_ancestors_within_distance(
         &ctx,
         &graph,
         vec!["N", "M", "H"],
         1,
-        vec!["N", "M", "H", "L", "K", "G"],
+        vec![("N", 0), ("M", 0), ("H", 0), ("L", 1), ("K", 1), ("G", 1)],
     )
     .await?;
     assert_ancestors_within_distance(&ctx, &graph, vec![], 100, vec![]).await?;

eden/mononoke/repo_attributes/commit_graph/commit_graph_testlib/src/utils.rs

@@ -653,19 +653,50 @@ pub async fn assert_ancestors_within_distance(
     ctx: &CoreContext,
     graph: &CommitGraph,
     cs_ids: Vec<&str>,
-    distance: u64,
-    expected_ancestors: Vec<&str>,
+    max_distance: u64,
+    expected_ancestors_and_distances: Vec<(&str, u64)>,
 ) -> Result<()> {
     let cs_ids: Vec<_> = cs_ids.into_iter().map(name_cs_id).collect();
-    let ancestors = graph
-        .ancestors_within_distance(ctx, cs_ids, distance)
+    let ancestors_and_distances = graph
+        .ancestors_within_distance_stream(ctx, cs_ids.clone(), max_distance)
         .await?
         .try_collect::<HashSet<_>>()
         .await?;
 
-    let expected_ancestors: HashSet<_> = expected_ancestors.into_iter().map(name_cs_id).collect();
+    let expected_ancestors_and_distances: HashSet<_> = expected_ancestors_and_distances
+        .into_iter()
+        .map(|(name, distance)| (name_cs_id(name), distance))
+        .collect();
+
+    assert_eq!(ancestors_and_distances, expected_ancestors_and_distances);
 
-    assert_eq!(ancestors, expected_ancestors);
+    let ancestors_and_boundaries = graph
+        .ancestors_within_distance(ctx, cs_ids, max_distance)
+        .await?;
+
+    assert_eq!(
+        ancestors_and_boundaries
+            .ancestors
+            .into_iter()
+            .collect::<HashSet<_>>(),
+        expected_ancestors_and_distances
+            .iter()
+            .map(|(cs_id, _)| cs_id)
+            .copied()
+            .collect::<HashSet<_>>()
+    );
+    assert_eq!(
+        ancestors_and_boundaries
+            .boundaries
+            .into_iter()
+            .collect::<HashSet<_>>(),
+        expected_ancestors_and_distances
+            .iter()
+            .filter(|(_, distance)| *distance == max_distance)
+            .map(|(cs_id, _)| cs_id)
+            .copied()
+            .collect::<HashSet<_>>()
+    );
 
     Ok(())
 }

eden/mononoke/repo_attributes/commit_graph/commit_graph_types/src/frontier.rs

@@ -136,6 +136,12 @@ impl Extend<(ChangesetId, Generation)> for ChangesetFrontier {
 #[derive(Clone, Debug)]
 pub struct ChangesetFrontierWithinDistance(BTreeMap<Generation, HashMap<ChangesetId, u64>>);
 
+#[derive(Clone, Debug)]
+pub struct AncestorsWithinDistance {
+    pub ancestors: Vec<ChangesetId>,
+    pub boundaries: Vec<ChangesetId>,
+}
+
 impl ChangesetFrontierWithinDistance {
     pub fn new() -> Self {
         Self(Default::default())