dag/
vertex_options.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8use std::collections::HashSet;
9
10use serde::Deserialize;
11
12use crate::Group;
13use crate::Vertex;
14
15/// A list of [`Vertex`]s (usually heads) with options attached to each vertex.
16#[derive(Default, Debug, Clone, Deserialize)]
17#[serde(transparent)]
18pub struct VertexListWithOptions {
19    list: Vec<(Vertex, VertexOptions)>,
20}
21
22/// Options attached to a vertex. Usually the vertex is a head. The head and its
23/// ancestors are going to be inserted to the graph. The options controls some
24/// details about the insertion.
25#[derive(Debug, Clone, Deserialize)]
26#[non_exhaustive]
27pub struct VertexOptions {
28    /// How many ids to reserve for this vertex. Suppose this vertex has id `n`,
29    /// then `n+1..=n+reserve_size` can only be used when inserting this vertex
30    /// and its ancestors in the same batch.
31    ///
32    /// Note: if any id `j` in the `n+1..=n+reserve_size` range were already
33    /// taken, then the reserve range becomes `n+1..j` instead. This avoids
34    /// fragmentation.
35    #[serde(default = "Default::default")]
36    pub reserve_size: u32,
37
38    /// The desired [`Group`] for this vertex. Note a vertex's group can be
39    /// moved down (ex. `NON_MASTER` to `MASTER`) but not moved up
40    /// (ex. `MASTER` to `NON_MASTER`).
41    /// - If the vertex does not exist, it will be inserted into the
42    ///   `desired_group`.
43    /// - If the vertex is already in a lower group than `desired_group`,
44    ///   the vertex will stay in that lower group unchanged.
45    /// - If the vertex is in a higher group than `desired_group`,
46    ///   the implementation (ex. `add_heads_and_flush`) might move the vertex
47    ///   (and its ancestors) to a lower group, or error out.
48    #[serde(default = "default_desired_group")]
49    pub desired_group: Group,
50}
51
52const fn default_desired_group() -> Group {
53    Group::NON_MASTER
54}
55
56impl Default for VertexOptions {
57    fn default() -> Self {
58        Self {
59            reserve_size: 0,
60            desired_group: default_desired_group(),
61        }
62    }
63}
64
65impl<'a> From<&'a [Vertex]> for VertexListWithOptions {
66    fn from(list: &'a [Vertex]) -> Self {
67        // Just use default options.
68        Self {
69            list: list
70                .iter()
71                .map(|v| (v.clone(), VertexOptions::default()))
72                .collect(),
73        }
74    }
75}
76
77impl From<Vec<Vertex>> for VertexListWithOptions {
78    fn from(list: Vec<Vertex>) -> Self {
79        // Just use default options.
80        Self {
81            list: list
82                .into_iter()
83                .map(|v| (v, VertexOptions::default()))
84                .collect(),
85        }
86    }
87}
88
89impl From<Vec<(Vertex, VertexOptions)>> for VertexListWithOptions {
90    fn from(list: Vec<(Vertex, VertexOptions)>) -> Self {
91        Self { list }
92    }
93}
94
95impl VertexListWithOptions {
96    /// Get the vertexes and their options.
97    pub fn vertex_options(&self) -> Vec<(Vertex, VertexOptions)> {
98        self.list.clone()
99    }
100
101    /// Get the vertexes.
102    pub fn vertexes(&self) -> Vec<Vertex> {
103        self.list.iter().map(|i| i.0.clone()).collect()
104    }
105
106    /// Sort the heads. Lower desired group first.
107    pub fn sort_by_group(mut self) -> Self {
108        self.list.sort_by_key(|(_v, opts)| opts.desired_group.0);
109        self
110    }
111
112    /// Get the vertexes, filter by the `desired_group` option.
113    pub fn vertexes_by_group(&self, group: Group) -> Vec<Vertex> {
114        self.list
115            .iter()
116            .filter_map(|(v, o)| {
117                if o.desired_group == group {
118                    Some(v.clone())
119                } else {
120                    None
121                }
122            })
123            .collect()
124    }
125
126    /// Test if this list is empty.
127    pub fn is_empty(&self) -> bool {
128        self.list.is_empty()
129    }
130
131    /// Add a new item to the list.
132    pub fn push(&mut self, head_opts: (Vertex, VertexOptions)) {
133        self.list.push(head_opts);
134    }
135
136    /// Set the `desired_group` option for all vertexes.
137    pub fn with_desired_group(mut self, group: Group) -> Self {
138        for (_v, opts) in self.list.iter_mut() {
139            opts.desired_group = group;
140        }
141        self
142    }
143
144    /// Chain another list. Vertexes that are already in this list are skipped.
145    pub fn chain(mut self, other: impl Into<Self>) -> Self {
146        let other = other.into();
147        let existed: HashSet<_> = self.vertexes().into_iter().collect();
148        for (v, o) in other.list {
149            if !existed.contains(&v) {
150                self.list.push((v, o))
151            }
152        }
153        self
154    }
155
156    /// Length of the `VertexListWithOptions`.
157    pub fn len(&self) -> usize {
158        self.list.len()
159    }
160
161    /// Minimal `desired_group` from all heads. `None` for empty heads list.
162    pub fn min_desired_group(&self) -> Option<Group> {
163        self.list.iter().map(|v| v.1.desired_group).min()
164    }
165
166    /// Maximum `desired_group` from all heads. `None` for empty heads list.
167    pub fn max_desired_group(&self) -> Option<Group> {
168        self.list.iter().map(|v| v.1.desired_group).max()
169    }
170}
dag/vertex_options.rs

dag/
vertex_options.rs
