1use std::collections::BTreeMap;
  2
  3use anyhow::Context;
  4use gix::{ObjectId, bstr::BStr};
  5use itertools::{Either, Itertools};
  6use rkyv::{Archive, Serialize};
  7use rocksdb::{DB, WriteBatch};
  8use yoke::{Yoke, Yokeable};
  9
 10use super::{
 11    Yoked,
 12    prefixes::{TREE_FAMILY, TREE_ITEM_FAMILY},
 13};
 14
 15#[derive(Serialize, Archive, Debug, PartialEq, Eq, Hash)]
 16pub struct Tree {
 17    pub indexed_tree_id: u64,
 18}
 19
 20impl Tree {
 21    pub fn insert(
 22        &self,
 23        database: &DB,
 24        batch: &mut WriteBatch,
 25        tree_oid: ObjectId,
 26    ) -> Result<(), anyhow::Error> {
 27        let cf = database
 28            .cf_handle(TREE_FAMILY)
 29            .context("tree column family missing")?;
 30
 31        batch.put_cf(
 32            cf,
 33            tree_oid.as_slice(),
 34            rkyv::to_bytes::<rkyv::rancor::Error>(self)?,
 35        );
 36
 37        Ok(())
 38    }
 39
 40    pub fn find(database: &DB, tree_oid: ObjectId) -> Result<Option<u64>, anyhow::Error> {
 41        let cf = database
 42            .cf_handle(TREE_FAMILY)
 43            .context("tree column family missing")?;
 44
 45        let Some(data) = database.get_cf(cf, tree_oid.as_slice())? else {
 46            return Ok(None);
 47        };
 48
 49        let data = rkyv::access::<<Self as Archive>::Archived, rkyv::rancor::Error>(data.as_ref())?;
 50
 51        Ok(Some(data.indexed_tree_id.to_native()))
 52    }
 53}
 54
 55#[derive(Serialize, Archive, Debug, PartialEq, Eq, Hash, Ord, PartialOrd)]
 56#[rkyv(derive(Ord, PartialOrd, Eq, PartialEq, Debug))]
 57#[rkyv(compare(PartialOrd, PartialEq))]
 58pub struct TreeKey(pub String);
 59
 60#[derive(Serialize, Archive, Debug, PartialEq, Eq, Default, Yokeable)]
 61pub struct SortedTree(pub BTreeMap<TreeKey, SortedTreeItem>);
 62
 63impl SortedTree {
 64    pub fn insert(
 65        &self,
 66        digest: u64,
 67        database: &DB,
 68        batch: &mut WriteBatch,
 69    ) -> Result<(), anyhow::Error> {
 70        let cf = database
 71            .cf_handle(TREE_ITEM_FAMILY)
 72            .context("tree column family missing")?;
 73
 74        batch.put_cf(
 75            cf,
 76            digest.to_ne_bytes(),
 77            rkyv::to_bytes::<rkyv::rancor::Error>(self)?,
 78        );
 79
 80        Ok(())
 81    }
 82
 83    pub fn get(digest: u64, database: &DB) -> Result<Option<YokedSortedTree>, anyhow::Error> {
 84        let cf = database
 85            .cf_handle(TREE_ITEM_FAMILY)
 86            .expect("tree column family missing");
 87
 88        database
 89            .get_cf(cf, digest.to_ne_bytes())?
 90            .map(|data| {
 91                Yoke::try_attach_to_cart(data.into_boxed_slice(), |data| {
 92                    rkyv::access::<_, rkyv::rancor::Error>(data)
 93                })
 94            })
 95            .transpose()
 96            .context("failed to parse full tree")
 97    }
 98}
 99
100#[derive(Serialize, Archive, Debug, PartialEq, Eq)]
101#[rkyv(
102    bytecheck(bounds(__C: rkyv::validation::ArchiveContext)),
103    serialize_bounds(__S: rkyv::ser::Writer + rkyv::ser::Allocator, __S::Error: rkyv::rancor::Source),
104)]
105pub enum SortedTreeItem {
106    File,
107    Directory(#[rkyv(omit_bounds)] SortedTree),
108}
109
110#[derive(Serialize, Archive, Debug, PartialEq, Eq, Hash)]
111pub struct Submodule {
112    pub url: String,
113    pub oid: [u8; 20],
114}
115
116#[derive(Serialize, Archive, Debug, PartialEq, Eq, Hash)]
117pub enum TreeItemKind {
118    Submodule(Submodule),
119    Tree,
120    File,
121}
122
123#[derive(Serialize, Archive, Debug, PartialEq, Eq, Hash, Yokeable)]
124pub struct TreeItem {
125    pub mode: u16,
126    pub kind: TreeItemKind,
127}
128
129pub type YokedSortedTree = Yoked<&'static <SortedTree as Archive>::Archived>;
130pub type YokedTreeItem = Yoked<&'static <TreeItem as Archive>::Archived>;
131pub type YokedTreeItemKey = Yoked<&'static [u8]>;
132pub type YokedTreeItemKeyUtf8 = Yoked<&'static str>;
133
134impl TreeItem {
135    pub fn insert(
136        &self,
137        buffer: &mut Vec<u8>,
138        digest: u64,
139        path: &BStr,
140        database: &DB,
141        batch: &mut WriteBatch,
142    ) -> Result<(), anyhow::Error> {
143        let cf = database
144            .cf_handle(TREE_ITEM_FAMILY)
145            .context("tree column family missing")?;
146
147        buffer.clear();
148        buffer.reserve(std::mem::size_of::<u64>() + path.len() + std::mem::size_of::<usize>());
149        buffer.extend_from_slice(&digest.to_ne_bytes());
150        buffer.extend_from_slice(&memchr::memchr_iter(b'/', path).count().to_be_bytes());
151        buffer.extend_from_slice(path.as_ref());
152
153        batch.put_cf(cf, &buffer, rkyv::to_bytes::<rkyv::rancor::Error>(self)?);
154
155        Ok(())
156    }
157
158    pub fn find_exact(
159        database: &DB,
160        digest: u64,
161        path: &[u8],
162    ) -> Result<Option<YokedTreeItem>, anyhow::Error> {
163        let cf = database
164            .cf_handle(TREE_ITEM_FAMILY)
165            .expect("tree column family missing");
166
167        let mut buffer = Vec::with_capacity(std::mem::size_of::<u64>() + path.len());
168        buffer.extend_from_slice(&digest.to_ne_bytes());
169        buffer.extend_from_slice(&memchr::memchr_iter(b'/', path).count().to_be_bytes());
170        buffer.extend_from_slice(path);
171
172        database
173            .get_cf(cf, buffer)?
174            .map(|data| {
175                Yoke::try_attach_to_cart(data.into_boxed_slice(), |data| {
176                    rkyv::access::<_, rkyv::rancor::Error>(data)
177                })
178            })
179            .transpose()
180            .context("failed to parse tree item")
181    }
182
183    pub fn find_prefix<'a>(
184        database: &'a DB,
185        digest: u64,
186        prefix: Option<&[u8]>,
187    ) -> impl Iterator<Item = Result<(YokedTreeItemKey, YokedTreeItem), anyhow::Error>> + use<'a>
188    {
189        let cf = database
190            .cf_handle(TREE_ITEM_FAMILY)
191            .expect("tree column family missing");
192
193        let (iterator, key) = match prefix {
194            None => {
195                let iterator = database.prefix_iterator_cf(cf, digest.to_ne_bytes());
196
197                (iterator, Either::Left(Either::Left(digest.to_be_bytes())))
198            }
199            Some([]) => {
200                let mut buffer = [0_u8; std::mem::size_of::<u64>() + std::mem::size_of::<usize>()];
201                buffer[..std::mem::size_of::<u64>()].copy_from_slice(&digest.to_ne_bytes());
202                buffer[std::mem::size_of::<u64>()..].copy_from_slice(&0_usize.to_be_bytes());
203
204                let iterator = database.prefix_iterator_cf(cf, buffer);
205
206                (iterator, Either::Left(Either::Right(buffer)))
207            }
208            Some(prefix) => {
209                let mut buffer = Vec::with_capacity(
210                    std::mem::size_of::<u64>() + prefix.len() + std::mem::size_of::<usize>(),
211                );
212                buffer.extend_from_slice(&digest.to_ne_bytes());
213                buffer.extend_from_slice(
214                    &(memchr::memchr_iter(b'/', prefix).count() + 1).to_be_bytes(),
215                );
216                buffer.extend_from_slice(prefix);
217                buffer.push(b'/');
218
219                let iterator = database.prefix_iterator_cf(cf, &buffer);
220
221                (iterator, Either::Right(buffer))
222            }
223        };
224
225        iterator
226            .take_while(move |v| {
227                v.as_ref().is_ok_and(|(k, _)| {
228                    k.starts_with(match key.as_ref() {
229                        Either::Left(Either::Right(v)) => v.as_ref(),
230                        Either::Left(Either::Left(v)) => v.as_ref(),
231                        Either::Right(v) => v.as_ref(),
232                    })
233                })
234            })
235            .map_ok(|(key, value)| {
236                let key = Yoke::attach_to_cart(key, |data| {
237                    &data[std::mem::size_of::<u64>() + std::mem::size_of::<usize>()..]
238                });
239                let value = Yoke::try_attach_to_cart(value, |data| {
240                    rkyv::access::<_, rkyv::rancor::Error>(data)
241                })
242                .context("Failed to open repository")?;
243                Ok((key, value))
244            })
245            .flatten()
246    }
247
248    pub fn contains(database: &DB, digest: u64) -> Result<bool, anyhow::Error> {
249        let cf = database
250            .cf_handle(TREE_ITEM_FAMILY)
251            .context("tree column family missing")?;
252
253        Ok(database
254            .prefix_iterator_cf(cf, digest.to_ne_bytes())
255            .next()
256            .transpose()?
257            .is_some())
258    }
259}