vim: Add global marks (#25702)

Closes https://github.com/zed-industries/zed/issues/13111

Release Notes:

- vim: Added global marks `'[A-Z]`
- vim: Added persistence for global (and local) marks. When re-opening
the same workspace your previous marks will be available.

---------

Co-authored-by: Conrad Irwin <conrad.irwin@gmail.com>

crates/vim/src/state.rs

@@ -3,27 +3,34 @@ use crate::normal::repeat::Replayer;
 use crate::surrounds::SurroundsType;
 use crate::{motion::Motion, object::Object};
 use crate::{ToggleRegistersView, UseSystemClipboard, Vim, VimSettings};
+use anyhow::Result;
 use collections::HashMap;
 use command_palette_hooks::{CommandPaletteFilter, CommandPaletteInterceptor};
+use db::define_connection;
+use db::sqlez_macros::sql;
 use editor::display_map::{is_invisible, replacement};
-use editor::{Anchor, ClipboardSelection, Editor};
+use editor::{Anchor, ClipboardSelection, Editor, MultiBuffer};
 use gpui::{
-    Action, App, BorrowAppContext, ClipboardEntry, ClipboardItem, Entity, Global, HighlightStyle,
-    StyledText, Task, TextStyle, WeakEntity,
+    Action, App, AppContext, BorrowAppContext, ClipboardEntry, ClipboardItem, Entity, EntityId,
+    Global, HighlightStyle, StyledText, Subscription, Task, TextStyle, WeakEntity,
 };
-use language::Point;
+use language::{Buffer, BufferEvent, BufferId, Point};
 use picker::{Picker, PickerDelegate};
+use project::{Project, ProjectItem, ProjectPath};
 use serde::{Deserialize, Serialize};
 use settings::{Settings, SettingsStore};
 use std::borrow::BorrowMut;
+use std::path::Path;
 use std::{fmt::Display, ops::Range, sync::Arc};
+use text::Bias;
 use theme::ThemeSettings;
 use ui::{
     h_flex, rems, ActiveTheme, Context, Div, FluentBuilder, KeyBinding, ParentElement,
     SharedString, Styled, StyledTypography, Window,
 };
+use util::ResultExt;
 use workspace::searchable::Direction;
-use workspace::Workspace;
+use workspace::{Workspace, WorkspaceDb, WorkspaceId};
 
 #[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
 pub enum Mode {
@@ -179,7 +186,7 @@ impl From<String> for Register {
     }
 }
 
-#[derive(Default, Clone)]
+#[derive(Default)]
 pub struct VimGlobals {
     pub last_find: Option<Motion>,
 
@@ -208,7 +215,399 @@ pub struct VimGlobals {
     pub recordings: HashMap<char, Vec<ReplayableAction>>,
 
     pub focused_vim: Option<WeakEntity<Vim>>,
+
+    pub marks: HashMap<EntityId, Entity<MarksState>>,
 }
+
+pub struct MarksState {
+    workspace: WeakEntity<Workspace>,
+
+    multibuffer_marks: HashMap<EntityId, HashMap<String, Vec<Anchor>>>,
+    buffer_marks: HashMap<BufferId, HashMap<String, Vec<text::Anchor>>>,
+    watched_buffers: HashMap<BufferId, (MarkLocation, Subscription, Subscription)>,
+
+    serialized_marks: HashMap<Arc<Path>, HashMap<String, Vec<Point>>>,
+    global_marks: HashMap<String, MarkLocation>,
+
+    _subscription: Subscription,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub enum MarkLocation {
+    Buffer(EntityId),
+    Path(Arc<Path>),
+}
+
+pub enum Mark {
+    Local(Vec<Anchor>),
+    Buffer(EntityId, Vec<Anchor>),
+    Path(Arc<Path>, Vec<Point>),
+}
+
+impl MarksState {
+    pub fn new(workspace: &Workspace, cx: &mut App) -> Entity<MarksState> {
+        cx.new(|cx| {
+            let buffer_store = workspace.project().read(cx).buffer_store().clone();
+            let subscription =
+                cx.subscribe(
+                    &buffer_store,
+                    move |this: &mut Self, _, event, cx| match event {
+                        project::buffer_store::BufferStoreEvent::BufferAdded(buffer) => {
+                            this.on_buffer_loaded(buffer, cx);
+                        }
+                        _ => {}
+                    },
+                );
+
+            let mut this = Self {
+                workspace: workspace.weak_handle(),
+                multibuffer_marks: HashMap::default(),
+                buffer_marks: HashMap::default(),
+                watched_buffers: HashMap::default(),
+                serialized_marks: HashMap::default(),
+                global_marks: HashMap::default(),
+                _subscription: subscription,
+            };
+
+            this.load(cx);
+            this
+        })
+    }
+
+    fn workspace_id(&self, cx: &App) -> Option<WorkspaceId> {
+        self.workspace
+            .read_with(cx, |workspace, _| workspace.database_id())
+            .ok()
+            .flatten()
+    }
+
+    fn project(&self, cx: &App) -> Option<Entity<Project>> {
+        self.workspace
+            .read_with(cx, |workspace, _| workspace.project().clone())
+            .ok()
+    }
+
+    fn load(&mut self, cx: &mut Context<Self>) {
+        cx.spawn(|this, mut cx| async move {
+            let Some(workspace_id) = this.update(&mut cx, |this, cx| this.workspace_id(cx))? else {
+                return Ok(());
+            };
+            let (marks, paths) = cx
+                .background_spawn(async move {
+                    let marks = DB.get_marks(workspace_id)?;
+                    let paths = DB.get_global_marks_paths(workspace_id)?;
+                    anyhow::Ok((marks, paths))
+                })
+                .await?;
+            this.update(&mut cx, |this, cx| this.loaded(marks, paths, cx))
+        })
+        .detach_and_log_err(cx);
+    }
+
+    fn loaded(
+        &mut self,
+        marks: Vec<SerializedMark>,
+        global_mark_paths: Vec<(String, Arc<Path>)>,
+        cx: &mut Context<Self>,
+    ) {
+        let Some(project) = self.project(cx) else {
+            return;
+        };
+
+        for mark in marks {
+            self.serialized_marks
+                .entry(mark.path)
+                .or_default()
+                .insert(mark.name, mark.points);
+        }
+
+        for (name, path) in global_mark_paths {
+            self.global_marks
+                .insert(name, MarkLocation::Path(path.clone()));
+
+            let project_path = project
+                .read(cx)
+                .worktrees(cx)
+                .filter_map(|worktree| {
+                    let relative = path.strip_prefix(worktree.read(cx).abs_path()).ok()?;
+                    Some(ProjectPath {
+                        worktree_id: worktree.read(cx).id(),
+                        path: relative.into(),
+                    })
+                })
+                .next();
+            if let Some(buffer) = project_path
+                .and_then(|project_path| project.read(cx).get_open_buffer(&project_path, cx))
+            {
+                self.on_buffer_loaded(&buffer, cx)
+            }
+        }
+    }
+
+    pub fn on_buffer_loaded(&mut self, buffer_handle: &Entity<Buffer>, cx: &mut Context<Self>) {
+        let Some(project) = self.project(cx) else {
+            return;
+        };
+        let Some(project_path) = buffer_handle.read(cx).project_path(cx) else {
+            return;
+        };
+        let Some(abs_path) = project.read(cx).absolute_path(&project_path, cx) else {
+            return;
+        };
+        let abs_path: Arc<Path> = abs_path.into();
+
+        let Some(serialized_marks) = self.serialized_marks.get(&abs_path) else {
+            return;
+        };
+
+        let mut loaded_marks = HashMap::default();
+        let buffer = buffer_handle.read(cx);
+        for (name, points) in serialized_marks.iter() {
+            loaded_marks.insert(
+                name.clone(),
+                points
+                    .iter()
+                    .map(|point| buffer.anchor_before(buffer.clip_point(*point, Bias::Left)))
+                    .collect(),
+            );
+        }
+        self.buffer_marks.insert(buffer.remote_id(), loaded_marks);
+        self.watch_buffer(MarkLocation::Path(abs_path), buffer_handle, cx)
+    }
+
+    fn serialize_buffer_marks(
+        &mut self,
+        path: Arc<Path>,
+        buffer: &Entity<Buffer>,
+        cx: &mut Context<Self>,
+    ) {
+        let new_points: HashMap<String, Vec<Point>> =
+            if let Some(anchors) = self.buffer_marks.get(&buffer.read(cx).remote_id()) {
+                anchors
+                    .iter()
+                    .map(|(name, anchors)| {
+                        (
+                            name.clone(),
+                            buffer
+                                .read(cx)
+                                .summaries_for_anchors::<Point, _>(anchors)
+                                .collect(),
+                        )
+                    })
+                    .collect()
+            } else {
+                HashMap::default()
+            };
+        let old_points = self.serialized_marks.get(&path.clone());
+        if old_points == Some(&new_points) {
+            return;
+        }
+        let mut to_write = HashMap::default();
+
+        for (key, value) in &new_points {
+            if self.is_global_mark(key) {
+                if self.global_marks.get(key) != Some(&MarkLocation::Path(path.clone())) {
+                    if let Some(workspace_id) = self.workspace_id(cx) {
+                        let path = path.clone();
+                        let key = key.clone();
+                        cx.background_spawn(async move {
+                            DB.set_global_mark_path(workspace_id, key, path).await
+                        })
+                        .detach_and_log_err(cx);
+                    }
+
+                    self.global_marks
+                        .insert(key.clone(), MarkLocation::Path(path.clone()));
+                }
+            }
+            if old_points.and_then(|o| o.get(key)) != Some(value) {
+                to_write.insert(key.clone(), value.clone());
+            }
+        }
+
+        self.serialized_marks.insert(path.clone(), new_points);
+
+        if let Some(workspace_id) = self.workspace_id(cx) {
+            cx.background_spawn(async move {
+                DB.set_marks(workspace_id, path.clone(), to_write).await?;
+                anyhow::Ok(())
+            })
+            .detach_and_log_err(cx);
+        }
+    }
+
+    fn is_global_mark(&self, key: &str) -> bool {
+        key.chars()
+            .next()
+            .is_some_and(|c| c.is_uppercase() || c.is_digit(10))
+    }
+
+    fn rename_buffer(
+        &mut self,
+        old_path: MarkLocation,
+        new_path: Arc<Path>,
+        buffer: &Entity<Buffer>,
+        cx: &mut Context<Self>,
+    ) {
+        if let MarkLocation::Buffer(entity_id) = old_path {
+            if let Some(old_marks) = self.multibuffer_marks.remove(&entity_id) {
+                let buffer_marks = old_marks
+                    .into_iter()
+                    .map(|(k, v)| (k, v.into_iter().map(|anchor| anchor.text_anchor).collect()))
+                    .collect();
+                self.buffer_marks
+                    .insert(buffer.read(cx).remote_id(), buffer_marks);
+            }
+        }
+        self.watch_buffer(MarkLocation::Path(new_path.clone()), buffer, cx);
+        self.serialize_buffer_marks(new_path, buffer, cx);
+    }
+
+    fn path_for_buffer(&self, buffer: &Entity<Buffer>, cx: &App) -> Option<Arc<Path>> {
+        let project_path = buffer.read(cx).project_path(cx)?;
+        let project = self.project(cx)?;
+        let abs_path = project.read(cx).absolute_path(&project_path, cx)?;
+        Some(abs_path.into())
+    }
+
+    fn points_at(
+        &self,
+        location: &MarkLocation,
+        multi_buffer: &Entity<MultiBuffer>,
+        cx: &App,
+    ) -> bool {
+        match location {
+            MarkLocation::Buffer(entity_id) => entity_id == &multi_buffer.entity_id(),
+            MarkLocation::Path(path) => {
+                let Some(singleton) = multi_buffer.read(cx).as_singleton() else {
+                    return false;
+                };
+                self.path_for_buffer(&singleton, cx).as_ref() == Some(path)
+            }
+        }
+    }
+
+    pub fn watch_buffer(
+        &mut self,
+        mark_location: MarkLocation,
+        buffer_handle: &Entity<Buffer>,
+        cx: &mut Context<Self>,
+    ) {
+        let on_change = cx.subscribe(buffer_handle, move |this, buffer, event, cx| match event {
+            BufferEvent::Edited => {
+                if let Some(path) = this.path_for_buffer(&buffer, cx) {
+                    this.serialize_buffer_marks(path, &buffer, cx);
+                }
+            }
+            BufferEvent::FileHandleChanged => {
+                let buffer_id = buffer.read(cx).remote_id();
+                if let Some(old_path) = this
+                    .watched_buffers
+                    .get(&buffer_id.clone())
+                    .map(|(path, _, _)| path.clone())
+                {
+                    if let Some(new_path) = this.path_for_buffer(&buffer, cx) {
+                        this.rename_buffer(old_path, new_path, &buffer, cx)
+                    }
+                }
+            }
+            _ => {}
+        });
+
+        let on_release = cx.observe_release(buffer_handle, |this, buffer, _| {
+            this.watched_buffers.remove(&buffer.remote_id());
+            this.buffer_marks.remove(&buffer.remote_id());
+        });
+
+        self.watched_buffers.insert(
+            buffer_handle.read(cx).remote_id(),
+            (mark_location, on_change, on_release),
+        );
+    }
+
+    pub fn set_mark(
+        &mut self,
+        name: String,
+        multibuffer: &Entity<MultiBuffer>,
+        anchors: Vec<Anchor>,
+        cx: &mut Context<Self>,
+    ) {
+        let buffer = multibuffer.read(cx).as_singleton();
+        let abs_path = buffer.as_ref().and_then(|b| self.path_for_buffer(&b, cx));
+
+        let Some(abs_path) = abs_path else {
+            self.multibuffer_marks
+                .entry(multibuffer.entity_id())
+                .or_default()
+                .insert(name.clone(), anchors);
+            if self.is_global_mark(&name) {
+                self.global_marks
+                    .insert(name.clone(), MarkLocation::Buffer(multibuffer.entity_id()));
+            }
+            if let Some(buffer) = buffer {
+                let buffer_id = buffer.read(cx).remote_id();
+                if !self.watched_buffers.contains_key(&buffer_id) {
+                    self.watch_buffer(MarkLocation::Buffer(multibuffer.entity_id()), &buffer, cx)
+                }
+            }
+            return;
+        };
+        let buffer = buffer.unwrap();
+
+        let buffer_id = buffer.read(cx).remote_id();
+        self.buffer_marks.entry(buffer_id).or_default().insert(
+            name.clone(),
+            anchors
+                .into_iter()
+                .map(|anchor| anchor.text_anchor)
+                .collect(),
+        );
+        if !self.watched_buffers.contains_key(&buffer_id) {
+            self.watch_buffer(MarkLocation::Path(abs_path.clone()), &buffer, cx)
+        }
+        self.serialize_buffer_marks(abs_path, &buffer, cx)
+    }
+
+    pub fn get_mark(
+        &self,
+        name: &str,
+        multi_buffer: &Entity<MultiBuffer>,
+        cx: &App,
+    ) -> Option<Mark> {
+        let target = self.global_marks.get(name);
+
+        if !self.is_global_mark(name) || target.is_some_and(|t| self.points_at(t, multi_buffer, cx))
+        {
+            if let Some(anchors) = self.multibuffer_marks.get(&multi_buffer.entity_id()) {
+                return Some(Mark::Local(anchors.get(name)?.clone()));
+            }
+
+            let singleton = multi_buffer.read(cx).as_singleton()?;
+            let excerpt_id = *multi_buffer.read(cx).excerpt_ids().first().unwrap();
+            let buffer_id = singleton.read(cx).remote_id();
+            if let Some(anchors) = self.buffer_marks.get(&buffer_id) {
+                let text_anchors = anchors.get(name)?;
+                let anchors = text_anchors
+                    .into_iter()
+                    .map(|anchor| Anchor::in_buffer(excerpt_id, buffer_id, *anchor))
+                    .collect();
+                return Some(Mark::Local(anchors));
+            }
+        }
+
+        match target? {
+            MarkLocation::Buffer(entity_id) => {
+                let anchors = self.multibuffer_marks.get(&entity_id)?;
+                return Some(Mark::Buffer(*entity_id, anchors.get(name)?.clone()));
+            }
+            MarkLocation::Path(path) => {
+                let points = self.serialized_marks.get(path)?;
+                return Some(Mark::Path(path.clone(), points.get(name)?.clone()));
+            }
+        }
+    }
+}
+
 impl Global for VimGlobals {}
 
 impl VimGlobals {
@@ -228,8 +627,15 @@ impl VimGlobals {
         })
         .detach();
 
+        let mut was_enabled = None;
+
         cx.observe_global::<SettingsStore>(move |cx| {
-            if Vim::enabled(cx) {
+            let is_enabled = Vim::enabled(cx);
+            if was_enabled == Some(is_enabled) {
+                return;
+            }
+            was_enabled = Some(is_enabled);
+            if is_enabled {
                 KeyBinding::set_vim_mode(cx, true);
                 CommandPaletteFilter::update_global(cx, |filter, _| {
                     filter.show_namespace(Vim::NAMESPACE);
@@ -237,6 +643,17 @@ impl VimGlobals {
                 CommandPaletteInterceptor::update_global(cx, |interceptor, _| {
                     interceptor.set(Box::new(command_interceptor));
                 });
+                for window in cx.windows() {
+                    if let Some(workspace) = window.downcast::<Workspace>() {
+                        workspace
+                            .update(cx, |workspace, _, cx| {
+                                Vim::update_globals(cx, |globals, cx| {
+                                    globals.register_workspace(workspace, cx)
+                                });
+                            })
+                            .ok();
+                    }
+                }
             } else {
                 KeyBinding::set_vim_mode(cx, false);
                 *Vim::globals(cx) = VimGlobals::default();
@@ -249,6 +666,21 @@ impl VimGlobals {
             }
         })
         .detach();
+        cx.observe_new(|workspace: &mut Workspace, _, cx| {
+            Vim::update_globals(cx, |globals, cx| globals.register_workspace(workspace, cx));
+        })
+        .detach()
+    }
+
+    fn register_workspace(&mut self, workspace: &Workspace, cx: &mut Context<Workspace>) {
+        let entity_id = cx.entity_id();
+        self.marks.insert(entity_id, MarksState::new(workspace, cx));
+        cx.observe_release(&cx.entity(), move |_, _, cx| {
+            Vim::update_globals(cx, |globals, _| {
+                globals.marks.remove(&entity_id);
+            })
+        })
+        .detach();
     }
 
     pub(crate) fn write_registers(
@@ -799,3 +1231,111 @@ impl RegistersView {
             .modal(true)
     }
 }
+
+define_connection! (
+    pub static ref DB: VimDb<WorkspaceDb> = &[
+        sql! (
+            CREATE TABLE vim_marks (
+              workspace_id INTEGER,
+              mark_name TEXT,
+              path BLOB,
+              value TEXT
+            );
+            CREATE UNIQUE INDEX idx_vim_marks ON vim_marks (workspace_id, mark_name, path);
+        ),
+        sql! (
+            CREATE TABLE vim_global_marks_paths(
+                workspace_id INTEGER,
+                mark_name TEXT,
+                path BLOB
+            );
+            CREATE UNIQUE INDEX idx_vim_global_marks_paths
+            ON vim_global_marks_paths(workspace_id, mark_name);
+        ),
+    ];
+);
+
+struct SerializedMark {
+    path: Arc<Path>,
+    name: String,
+    points: Vec<Point>,
+}
+
+impl VimDb {
+    pub(crate) async fn set_marks(
+        &self,
+        workspace_id: WorkspaceId,
+        path: Arc<Path>,
+        marks: HashMap<String, Vec<Point>>,
+    ) -> Result<()> {
+        let result = self
+            .write(move |conn| {
+                let mut query = conn.exec_bound(sql!(
+                    INSERT OR REPLACE INTO vim_marks
+                        (workspace_id, mark_name, path, value)
+                    VALUES
+                        (?, ?, ?, ?)
+                ))?;
+                for (mark_name, value) in marks {
+                    let pairs: Vec<(u32, u32)> = value
+                        .into_iter()
+                        .map(|point| (point.row, point.column))
+                        .collect();
+                    let serialized = serde_json::to_string(&pairs)?;
+                    query((workspace_id, mark_name, path.clone(), serialized))?;
+                }
+                Ok(())
+            })
+            .await;
+        result
+    }
+
+    fn get_marks(&self, workspace_id: WorkspaceId) -> Result<Vec<SerializedMark>> {
+        let result: Vec<(Arc<Path>, String, String)> = self.select_bound(sql!(
+            SELECT path, mark_name, value FROM vim_marks
+                WHERE workspace_id = ?
+        ))?(workspace_id)?;
+
+        Ok(result
+            .into_iter()
+            .filter_map(|(path, name, value)| {
+                let pairs: Vec<(u32, u32)> = serde_json::from_str(&value).log_err()?;
+                Some(SerializedMark {
+                    path,
+                    name,
+                    points: pairs
+                        .into_iter()
+                        .map(|(row, column)| Point { row, column })
+                        .collect(),
+                })
+            })
+            .collect())
+    }
+
+    pub(crate) async fn set_global_mark_path(
+        &self,
+        workspace_id: WorkspaceId,
+        mark_name: String,
+        path: Arc<Path>,
+    ) -> Result<()> {
+        self.write(move |conn| {
+            conn.exec_bound(sql!(
+                INSERT OR REPLACE INTO vim_global_marks_paths
+                    (workspace_id, mark_name, path)
+                VALUES
+                    (?, ?, ?)
+            ))?((workspace_id, mark_name, path))
+        })
+        .await
+    }
+
+    pub fn get_global_marks_paths(
+        &self,
+        workspace_id: WorkspaceId,
+    ) -> Result<Vec<(String, Arc<Path>)>> {
+        self.select_bound(sql!(
+        SELECT mark_name, path FROM vim_global_marks_paths
+            WHERE workspace_id = ?
+        ))?(workspace_id)
+    }
+}