dag/idmap/
mem_idmap.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::BTreeMap;
9use std::sync::atomic;
10use std::sync::atomic::AtomicU64;
11
12use super::IdMapWrite;
13use crate::errors::NotFoundError;
14use crate::id::Group;
15use crate::id::Id;
16use crate::id::VertexName;
17use crate::ops::IdConvert;
18use crate::ops::Persist;
19use crate::ops::PrefixLookup;
20use crate::Result;
21use crate::VerLink;
22
23/// Bi-directional mapping between an integer id and a name (`[u8]`).
24///
25/// Private. Stored in memory.
26pub struct MemIdMap {
27    core: CoreMemIdMap,
28    map_id: String,
29    map_version: VerLink,
30}
31
32/// Subset of the `MemIdMap` interface that does not have "map_version".
33/// or "version" concept.
34#[derive(Default, Clone)]
35pub(crate) struct CoreMemIdMap {
36    id2name: BTreeMap<Id, VertexName>,
37    name2id: BTreeMap<VertexName, Id>,
38}
39
40impl MemIdMap {
41    /// Create an empty [`MemIdMap`].
42    pub fn new() -> Self {
43        Self {
44            core: Default::default(),
45            map_id: format!("mem:{}", next_id()),
46            map_version: VerLink::new(),
47        }
48    }
49}
50
51impl Clone for MemIdMap {
52    fn clone(&self) -> Self {
53        Self {
54            core: self.core.clone(),
55            map_id: self.map_id.clone(),
56            map_version: self.map_version.clone(),
57        }
58    }
59}
60
61impl CoreMemIdMap {
62    pub fn lookup_vertex_id(&self, name: &VertexName) -> Option<Id> {
63        self.name2id.get(&name).copied()
64    }
65
66    pub fn lookup_vertex_name(&self, id: Id) -> Option<VertexName> {
67        self.id2name.get(&id).cloned()
68    }
69
70    pub fn lookup_vertexes_by_hex_prefix(
71        &self,
72        hex_prefix: &[u8],
73        limit: usize,
74    ) -> Result<Vec<VertexName>> {
75        let start = VertexName::from_hex(hex_prefix)?;
76        let mut result = Vec::new();
77        for (vertex, _) in self.name2id.range(start..) {
78            if !vertex.to_hex().as_bytes().starts_with(hex_prefix) {
79                break;
80            }
81            result.push(vertex.clone());
82            if result.len() >= limit {
83                break;
84            }
85        }
86        Ok(result)
87    }
88
89    pub fn has_vertex_name(&self, name: &VertexName) -> bool {
90        self.name2id.contains_key(name)
91    }
92
93    pub fn has_vertex_id(&self, id: Id) -> bool {
94        self.id2name.contains_key(&id)
95    }
96
97    pub fn insert_vertex_id_name(&mut self, id: Id, vertex_name: VertexName) {
98        self.name2id.insert(vertex_name.clone(), id);
99        self.id2name.insert(id, vertex_name);
100    }
101
102    pub fn remove_range(&mut self, low: Id, high: Id) -> Result<Vec<VertexName>> {
103        let to_remove: Vec<(Id, VertexName)> = self
104            .id2name
105            .range(low..=high)
106            .map(|(i, n)| (*i, n.clone()))
107            .collect();
108        for (id, name) in &to_remove {
109            self.id2name.remove(id);
110            self.name2id.remove(name);
111        }
112        Ok(to_remove.into_iter().map(|(_, v)| v).collect())
113    }
114}
115
116#[async_trait::async_trait]
117impl IdConvert for MemIdMap {
118    async fn vertex_id(&self, name: VertexName) -> Result<Id> {
119        self.core
120            .lookup_vertex_id(&name)
121            .ok_or_else(|| name.not_found_error())
122    }
123    async fn vertex_id_with_max_group(
124        &self,
125        name: &VertexName,
126        max_group: Group,
127    ) -> Result<Option<Id>> {
128        let optional_id = self.core.name2id.get(name).and_then(|id| {
129            if id.group() <= max_group {
130                Some(*id)
131            } else {
132                None
133            }
134        });
135        Ok(optional_id)
136    }
137    async fn vertex_name(&self, id: Id) -> Result<VertexName> {
138        self.core
139            .lookup_vertex_name(id)
140            .ok_or_else(|| id.not_found_error())
141    }
142    async fn contains_vertex_name(&self, name: &VertexName) -> Result<bool> {
143        Ok(self.core.has_vertex_name(name))
144    }
145
146    async fn contains_vertex_id_locally(&self, ids: &[Id]) -> Result<Vec<bool>> {
147        Ok(ids
148            .iter()
149            .copied()
150            .map(|id| self.core.has_vertex_id(id))
151            .collect())
152    }
153    async fn contains_vertex_name_locally(&self, names: &[VertexName]) -> Result<Vec<bool>> {
154        Ok(names
155            .iter()
156            .map(|name| self.core.has_vertex_name(name))
157            .collect())
158    }
159
160    fn map_id(&self) -> &str {
161        &self.map_id
162    }
163
164    fn map_version(&self) -> &VerLink {
165        &self.map_version
166    }
167}
168
169// TODO: Reconsider re-assign master cases. Currently they are ignored.
170#[async_trait::async_trait]
171impl IdMapWrite for MemIdMap {
172    async fn insert(&mut self, id: Id, name: &[u8]) -> Result<()> {
173        let vertex_name = VertexName::copy_from(name);
174        self.core.insert_vertex_id_name(id, vertex_name);
175        self.map_version.bump();
176        Ok(())
177    }
178    async fn remove_range(&mut self, low: Id, high: Id) -> Result<Vec<VertexName>> {
179        self.map_version = VerLink::new();
180        self.core.remove_range(low, high)
181    }
182}
183
184impl Persist for MemIdMap {
185    type Lock = ();
186
187    fn lock(&mut self) -> Result<Self::Lock> {
188        Ok(())
189    }
190
191    fn reload(&mut self, _lock: &Self::Lock) -> Result<()> {
192        Ok(())
193    }
194
195    fn persist(&mut self, _lock: &Self::Lock) -> Result<()> {
196        Ok(())
197    }
198}
199
200#[async_trait::async_trait]
201impl PrefixLookup for MemIdMap {
202    async fn vertexes_by_hex_prefix(
203        &self,
204        hex_prefix: &[u8],
205        limit: usize,
206    ) -> Result<Vec<VertexName>> {
207        self.core.lookup_vertexes_by_hex_prefix(hex_prefix, limit)
208    }
209}
210
211fn next_id() -> u64 {
212    static ID: AtomicU64 = AtomicU64::new(0);
213    ID.fetch_add(1, atomic::Ordering::AcqRel)
214}
dag/idmap/mem_idmap.rs

dag/idmap/
mem_idmap.rs
