1use anyhow::{anyhow, Result};
  2use gpui::{AppContext, Entity, ModelContext, ModelHandle};
  3use language::LanguageRegistry;
  4use project::{Fs, Project};
  5use smol::channel;
  6use std::{path::PathBuf, sync::Arc};
  7use util::ResultExt;
  8use workspace::WorkspaceCreated;
  9
 10pub fn init(fs: Arc<dyn Fs>, language_registry: Arc<LanguageRegistry>, cx: &mut AppContext) {
 11    let vector_store = cx.add_model(|cx| VectorStore::new(fs, language_registry));
 12
 13    cx.subscribe_global::<WorkspaceCreated, _>({
 14        let vector_store = vector_store.clone();
 15        move |event, cx| {
 16            let workspace = &event.0;
 17            if let Some(workspace) = workspace.upgrade(cx) {
 18                let project = workspace.read(cx).project().clone();
 19                if project.read(cx).is_local() {
 20                    vector_store.update(cx, |store, cx| {
 21                        store.add_project(project, cx);
 22                    });
 23                }
 24            }
 25        }
 26    })
 27    .detach();
 28}
 29
 30struct Document {
 31    offset: usize,
 32    name: String,
 33    embedding: Vec<f32>,
 34}
 35
 36struct IndexedFile {
 37    path: PathBuf,
 38    sha1: String,
 39    documents: Vec<Document>,
 40}
 41
 42struct SearchResult {
 43    path: PathBuf,
 44    offset: usize,
 45    name: String,
 46    distance: f32,
 47}
 48
 49struct VectorStore {
 50    fs: Arc<dyn Fs>,
 51    language_registry: Arc<LanguageRegistry>,
 52}
 53
 54impl VectorStore {
 55    fn new(fs: Arc<dyn Fs>, language_registry: Arc<LanguageRegistry>) -> Self {
 56        Self {
 57            fs,
 58            language_registry,
 59        }
 60    }
 61
 62    async fn index_file(
 63        fs: &Arc<dyn Fs>,
 64        language_registry: &Arc<LanguageRegistry>,
 65        file_path: PathBuf,
 66    ) -> Result<IndexedFile> {
 67        eprintln!("indexing file {file_path:?}");
 68        Err(anyhow!("not implemented"))
 69        // todo!();
 70    }
 71
 72    fn add_project(&mut self, project: ModelHandle<Project>, cx: &mut ModelContext<Self>) {
 73        let worktree_scans_complete = project
 74            .read(cx)
 75            .worktrees(cx)
 76            .map(|worktree| worktree.read(cx).as_local().unwrap().scan_complete())
 77            .collect::<Vec<_>>();
 78
 79        let fs = self.fs.clone();
 80        let language_registry = self.language_registry.clone();
 81
 82        cx.spawn(|this, cx| async move {
 83            futures::future::join_all(worktree_scans_complete).await;
 84
 85            let worktrees = project.read_with(&cx, |project, cx| {
 86                project
 87                    .worktrees(cx)
 88                    .map(|worktree| worktree.read(cx).snapshot())
 89                    .collect::<Vec<_>>()
 90            });
 91
 92            let (paths_tx, paths_rx) = channel::unbounded::<PathBuf>();
 93            let (indexed_files_tx, indexed_files_rx) = channel::unbounded::<IndexedFile>();
 94            cx.background()
 95                .spawn(async move {
 96                    for worktree in worktrees {
 97                        for file in worktree.files(false, 0) {
 98                            paths_tx.try_send(worktree.absolutize(&file.path)).unwrap();
 99                        }
100                    }
101                })
102                .detach();
103            cx.background()
104                .spawn(async move {
105                    while let Ok(indexed_file) = indexed_files_rx.recv().await {
106                        // write document to database
107                    }
108                })
109                .detach();
110            cx.background()
111                .scoped(|scope| {
112                    for _ in 0..cx.background().num_cpus() {
113                        scope.spawn(async {
114                            while let Ok(file_path) = paths_rx.recv().await {
115                                if let Some(indexed_file) =
116                                    Self::index_file(&fs, &language_registry, file_path)
117                                        .await
118                                        .log_err()
119                                {
120                                    indexed_files_tx.try_send(indexed_file).unwrap();
121                                }
122                            }
123                        });
124                    }
125                })
126                .await;
127        })
128        .detach();
129    }
130}
131
132impl Entity for VectorStore {
133    type Event = ();
134}