1use std::path::PathBuf;
2use std::sync::atomic::AtomicBool;
3use std::sync::Arc;
4
5use anyhow::{anyhow, Result};
6use futures::future::{self, BoxFuture, Shared};
7use futures::FutureExt;
8use fuzzy::StringMatchCandidate;
9use gpui::{AppContext, BackgroundExecutor, Global, ReadGlobal, Task, UpdateGlobal};
10use heed::types::SerdeBincode;
11use heed::Database;
12use serde::{Deserialize, Serialize};
13use util::paths::SUPPORT_DIR;
14use util::ResultExt;
15
16use crate::crawler::{RustdocCrawler, RustdocProvider};
17use crate::{RustdocItem, RustdocItemKind};
18
19struct GlobalRustdocStore(Arc<RustdocStore>);
20
21impl Global for GlobalRustdocStore {}
22
23pub struct RustdocStore {
24 executor: BackgroundExecutor,
25 database_future: Shared<BoxFuture<'static, Result<Arc<RustdocDatabase>, Arc<anyhow::Error>>>>,
26}
27
28impl RustdocStore {
29 pub fn global(cx: &AppContext) -> Arc<Self> {
30 GlobalRustdocStore::global(cx).0.clone()
31 }
32
33 pub fn init_global(cx: &mut AppContext) {
34 GlobalRustdocStore::set_global(
35 cx,
36 GlobalRustdocStore(Arc::new(Self::new(cx.background_executor().clone()))),
37 );
38 }
39
40 pub fn new(executor: BackgroundExecutor) -> Self {
41 let database_future = executor
42 .spawn({
43 let executor = executor.clone();
44 async move {
45 RustdocDatabase::new(SUPPORT_DIR.join("docs/rust/rustdoc-db.0.mdb"), executor)
46 }
47 })
48 .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new)))
49 .boxed()
50 .shared();
51
52 Self {
53 executor,
54 database_future,
55 }
56 }
57
58 pub async fn load(
59 &self,
60 crate_name: String,
61 item_path: Option<String>,
62 ) -> Result<RustdocDatabaseEntry> {
63 self.database_future
64 .clone()
65 .await
66 .map_err(|err| anyhow!(err))?
67 .load(crate_name, item_path)
68 .await
69 }
70
71 pub fn index(
72 &self,
73 crate_name: String,
74 provider: Box<dyn RustdocProvider + Send + Sync + 'static>,
75 ) -> Task<Result<()>> {
76 let database_future = self.database_future.clone();
77 self.executor.spawn(async move {
78 let crawler = RustdocCrawler::new(provider);
79
80 let Some(crate_docs) = crawler.crawl(crate_name.clone()).await? else {
81 return Ok(());
82 };
83
84 let database = database_future.await.map_err(|err| anyhow!(err))?;
85
86 database
87 .insert(crate_name.clone(), None, crate_docs.crate_root_markdown)
88 .await?;
89
90 for (item, item_docs) in crate_docs.items {
91 database
92 .insert(crate_name.clone(), Some(&item), item_docs)
93 .await?;
94 }
95
96 Ok(())
97 })
98 }
99
100 pub fn search(&self, query: String) -> Task<Vec<String>> {
101 let executor = self.executor.clone();
102 let database_future = self.database_future.clone();
103 self.executor.spawn(async move {
104 if query.is_empty() {
105 return Vec::new();
106 }
107
108 let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
109 return Vec::new();
110 };
111
112 let Some(items) = database.keys().await.log_err() else {
113 return Vec::new();
114 };
115
116 let candidates = items
117 .iter()
118 .enumerate()
119 .map(|(ix, item_path)| StringMatchCandidate::new(ix, item_path.clone()))
120 .collect::<Vec<_>>();
121
122 let matches = fuzzy::match_strings(
123 &candidates,
124 &query,
125 false,
126 100,
127 &AtomicBool::default(),
128 executor,
129 )
130 .await;
131
132 matches
133 .into_iter()
134 .map(|mat| items[mat.candidate_id].clone())
135 .collect()
136 })
137 }
138}
139
140#[derive(Serialize, Deserialize)]
141pub enum RustdocDatabaseEntry {
142 Crate { docs: String },
143 Item { kind: RustdocItemKind, docs: String },
144}
145
146impl RustdocDatabaseEntry {
147 pub fn docs(&self) -> &str {
148 match self {
149 Self::Crate { docs } | Self::Item { docs, .. } => &docs,
150 }
151 }
152}
153
154struct RustdocDatabase {
155 executor: BackgroundExecutor,
156 env: heed::Env,
157 entries: Database<SerdeBincode<String>, SerdeBincode<RustdocDatabaseEntry>>,
158}
159
160impl RustdocDatabase {
161 pub fn new(path: PathBuf, executor: BackgroundExecutor) -> Result<Self> {
162 std::fs::create_dir_all(&path)?;
163
164 const ONE_GB_IN_BYTES: usize = 1024 * 1024 * 1024;
165 let env = unsafe {
166 heed::EnvOpenOptions::new()
167 .map_size(ONE_GB_IN_BYTES)
168 .max_dbs(1)
169 .open(path)?
170 };
171
172 let mut txn = env.write_txn()?;
173 let entries = env.create_database(&mut txn, Some("rustdoc_entries"))?;
174 txn.commit()?;
175
176 Ok(Self {
177 executor,
178 env,
179 entries,
180 })
181 }
182
183 pub fn keys(&self) -> Task<Result<Vec<String>>> {
184 let env = self.env.clone();
185 let entries = self.entries;
186
187 self.executor.spawn(async move {
188 let txn = env.read_txn()?;
189 let mut iter = entries.iter(&txn)?;
190 let mut keys = Vec::new();
191 while let Some((key, _value)) = iter.next().transpose()? {
192 keys.push(key);
193 }
194
195 Ok(keys)
196 })
197 }
198
199 pub fn load(
200 &self,
201 crate_name: String,
202 item_path: Option<String>,
203 ) -> Task<Result<RustdocDatabaseEntry>> {
204 let env = self.env.clone();
205 let entries = self.entries;
206 let item_path = if let Some(item_path) = item_path {
207 format!("{crate_name}::{item_path}")
208 } else {
209 crate_name
210 };
211
212 self.executor.spawn(async move {
213 let txn = env.read_txn()?;
214 entries
215 .get(&txn, &item_path)?
216 .ok_or_else(|| anyhow!("no docs found for {item_path}"))
217 })
218 }
219
220 pub fn insert(
221 &self,
222 crate_name: String,
223 item: Option<&RustdocItem>,
224 docs: String,
225 ) -> Task<Result<()>> {
226 let env = self.env.clone();
227 let entries = self.entries;
228 let (item_path, entry) = if let Some(item) = item {
229 (
230 format!("{crate_name}::{}", item.display()),
231 RustdocDatabaseEntry::Item {
232 kind: item.kind,
233 docs,
234 },
235 )
236 } else {
237 (crate_name, RustdocDatabaseEntry::Crate { docs })
238 };
239
240 self.executor.spawn(async move {
241 let mut txn = env.write_txn()?;
242 entries.put(&mut txn, &item_path, &entry)?;
243 txn.commit()?;
244 Ok(())
245 })
246 }
247}