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::VertexName;
14
15/// A list of [`VertexName`]s (usually heads) with options attached to each vertex.
16#[derive(Default, Debug, Clone, Deserialize)]
17#[serde(transparent)]
18pub struct VertexListWithOptions {
19    list: Vec<(VertexName, 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 highest [`Group`] for this vertex. If set to `NON_MASTER` then
39    /// this vertex could end up in `MASTER` or `NON_MASTER`. If set to
40    /// `MASTER` then this vertex will end up in `MASTER` group.
41    #[serde(default = "default_highest_group")]
42    pub highest_group: Group,
43}
44
45const fn default_highest_group() -> Group {
46    Group::NON_MASTER
47}
48
49impl Default for VertexOptions {
50    fn default() -> Self {
51        Self {
52            reserve_size: 0,
53            highest_group: default_highest_group(),
54        }
55    }
56}
57
58impl<'a> From<&'a [VertexName]> for VertexListWithOptions {
59    fn from(list: &'a [VertexName]) -> Self {
60        // Just use default options.
61        Self {
62            list: list
63                .iter()
64                .map(|v| (v.clone(), VertexOptions::default()))
65                .collect(),
66        }
67    }
68}
69
70impl From<Vec<VertexName>> for VertexListWithOptions {
71    fn from(list: Vec<VertexName>) -> Self {
72        // Just use default options.
73        Self {
74            list: list
75                .into_iter()
76                .map(|v| (v, VertexOptions::default()))
77                .collect(),
78        }
79    }
80}
81
82impl From<Vec<(VertexName, VertexOptions)>> for VertexListWithOptions {
83    fn from(list: Vec<(VertexName, VertexOptions)>) -> Self {
84        Self { list }
85    }
86}
87
88impl VertexListWithOptions {
89    /// Get the vertexes and their options.
90    pub fn vertex_options(&self) -> Vec<(VertexName, VertexOptions)> {
91        self.list.clone()
92    }
93
94    /// Get the vertexes.
95    pub fn vertexes(&self) -> Vec<VertexName> {
96        self.list.iter().map(|i| i.0.clone()).collect()
97    }
98
99    /// Get the vertexes, filter by the `highest_group` option.
100    pub fn vertexes_by_group(&self, group: Group) -> Vec<VertexName> {
101        self.list
102            .iter()
103            .filter_map(|(v, o)| {
104                if o.highest_group == group {
105                    Some(v.clone())
106                } else {
107                    None
108                }
109            })
110            .collect()
111    }
112
113    /// Test if this list is empty.
114    pub fn is_empty(&self) -> bool {
115        self.list.is_empty()
116    }
117
118    /// Add a new item to the list.
119    pub fn push(&mut self, head_opts: (VertexName, VertexOptions)) {
120        self.list.push(head_opts);
121    }
122
123    /// Set the `highest_group` option for all vertexes.
124    pub fn with_highest_group(mut self, group: Group) -> Self {
125        for (_v, opts) in self.list.iter_mut() {
126            opts.highest_group = group;
127        }
128        self
129    }
130
131    /// Chain another list. Vertexes that are already in this list are skipped.
132    pub fn chain(mut self, other: impl Into<Self>) -> Self {
133        let other = other.into();
134        let existed: HashSet<_> = self.vertexes().into_iter().collect();
135        for (v, o) in other.list {
136            if !existed.contains(&v) {
137                self.list.push((v, o))
138            }
139        }
140        self
141    }
142}