neovim_connection.rs

  1use std::path::PathBuf;
  2#[cfg(feature = "neovim")]
  3use std::{
  4    cmp,
  5    ops::{Deref, DerefMut, Range},
  6};
  7
  8#[cfg(feature = "neovim")]
  9use async_compat::Compat;
 10#[cfg(feature = "neovim")]
 11use async_trait::async_trait;
 12#[cfg(feature = "neovim")]
 13use gpui::Keystroke;
 14
 15#[cfg(feature = "neovim")]
 16use language::Point;
 17
 18#[cfg(feature = "neovim")]
 19use nvim_rs::{
 20    Handler, Neovim, UiAttachOptions, Value, create::tokio::new_child_cmd, error::LoopError,
 21};
 22#[cfg(feature = "neovim")]
 23use parking_lot::ReentrantMutex;
 24use serde::{Deserialize, Serialize};
 25#[cfg(feature = "neovim")]
 26use tokio::{
 27    process::{Child, ChildStdin, Command},
 28    task::JoinHandle,
 29};
 30
 31use crate::state::Mode;
 32use collections::VecDeque;
 33
 34// Neovim doesn't like to be started simultaneously from multiple threads. We use this lock
 35// to ensure we are only constructing one neovim connection at a time.
 36#[cfg(feature = "neovim")]
 37static NEOVIM_LOCK: ReentrantMutex<()> = ReentrantMutex::new(());
 38
 39#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
 40pub enum NeovimData {
 41    Put { state: String },
 42    Key(String),
 43    Get { state: String, mode: Mode },
 44    ReadRegister { name: char, value: String },
 45    Exec { command: String },
 46    SetOption { value: String },
 47}
 48
 49pub struct NeovimConnection {
 50    data: VecDeque<NeovimData>,
 51    #[cfg(feature = "neovim")]
 52    test_case_id: String,
 53    #[cfg(feature = "neovim")]
 54    nvim: Neovim<nvim_rs::compat::tokio::Compat<ChildStdin>>,
 55    #[cfg(feature = "neovim")]
 56    _join_handle: JoinHandle<Result<(), Box<LoopError>>>,
 57    #[cfg(feature = "neovim")]
 58    _child: Child,
 59}
 60
 61impl NeovimConnection {
 62    pub async fn new(mut test_case_id: String) -> Self {
 63        // When running under perf, don't create duplicate files.
 64        if cfg!(perf_enabled) {
 65            if test_case_id.ends_with(perf::consts::SUF_NORMAL) {
 66                test_case_id.truncate(test_case_id.len() - perf::consts::SUF_NORMAL.len());
 67            }
 68        }
 69        #[cfg(feature = "neovim")]
 70        let handler = NvimHandler {};
 71        #[cfg(feature = "neovim")]
 72        let (nvim, join_handle, child) = Compat::new(async {
 73            // Ensure we don't create neovim connections in parallel
 74            let _lock = NEOVIM_LOCK.lock();
 75            let (nvim, join_handle, child) = new_child_cmd(
 76                Command::new("nvim")
 77                    .arg("--embed")
 78                    .arg("--clean")
 79                    // disable swap (otherwise after about 1000 test runs you run out of swap file names)
 80                    .arg("-n")
 81                    // disable writing files (just in case)
 82                    .arg("-m"),
 83                handler,
 84            )
 85            .await
 86            .expect("Could not connect to neovim process");
 87
 88            nvim.ui_attach(100, 100, &UiAttachOptions::default())
 89                .await
 90                .expect("Could not attach to ui");
 91
 92            // Makes system act a little more like zed in terms of indentation
 93            nvim.set_option("smartindent", nvim_rs::Value::Boolean(true))
 94                .await
 95                .expect("Could not set smartindent on startup");
 96
 97            (nvim, join_handle, child)
 98        })
 99        .await;
100
101        Self {
102            #[cfg(feature = "neovim")]
103            data: Default::default(),
104            #[cfg(not(feature = "neovim"))]
105            data: Self::read_test_data(&test_case_id),
106            #[cfg(feature = "neovim")]
107            test_case_id,
108            #[cfg(feature = "neovim")]
109            nvim,
110            #[cfg(feature = "neovim")]
111            _join_handle: join_handle,
112            #[cfg(feature = "neovim")]
113            _child: child,
114        }
115    }
116
117    // Sends a keystroke to the neovim process.
118    #[cfg(feature = "neovim")]
119    pub async fn send_keystroke(&mut self, keystroke_text: &str) {
120        let mut keystroke = Keystroke::parse(keystroke_text).unwrap();
121
122        if keystroke.key == "<" {
123            keystroke.key = "lt".to_string()
124        }
125
126        let special = keystroke.modifiers.shift
127            || keystroke.modifiers.control
128            || keystroke.modifiers.alt
129            || keystroke.modifiers.platform
130            || keystroke.key.len() > 1;
131        let start = if special { "<" } else { "" };
132        let shift = if keystroke.modifiers.shift { "S-" } else { "" };
133        let ctrl = if keystroke.modifiers.control {
134            "C-"
135        } else {
136            ""
137        };
138        let alt = if keystroke.modifiers.alt { "M-" } else { "" };
139        let cmd = if keystroke.modifiers.platform {
140            "D-"
141        } else {
142            ""
143        };
144        let end = if special { ">" } else { "" };
145
146        let key = format!("{start}{shift}{ctrl}{alt}{cmd}{}{end}", keystroke.key);
147
148        self.data
149            .push_back(NeovimData::Key(keystroke_text.to_string()));
150        self.nvim
151            .input(&key)
152            .await
153            .expect("Could not input keystroke");
154    }
155
156    #[cfg(not(feature = "neovim"))]
157    pub async fn send_keystroke(&mut self, keystroke_text: &str) {
158        if matches!(self.data.front(), Some(NeovimData::Get { .. })) {
159            self.data.pop_front();
160        }
161        assert_eq!(
162            self.data.pop_front(),
163            Some(NeovimData::Key(keystroke_text.to_string())),
164            "operation does not match recorded script. re-record with --features=neovim"
165        );
166    }
167
168    #[cfg(feature = "neovim")]
169    pub async fn set_state(&mut self, marked_text: &str) {
170        let (text, selections) = parse_state(marked_text);
171
172        let nvim_buffer = self
173            .nvim
174            .get_current_buf()
175            .await
176            .expect("Could not get neovim buffer");
177        let lines = text
178            .split('\n')
179            .map(|line| line.to_string())
180            .collect::<Vec<_>>();
181
182        nvim_buffer
183            .set_lines(0, -1, false, lines)
184            .await
185            .expect("Could not set nvim buffer text");
186
187        self.nvim
188            .input("<escape>")
189            .await
190            .expect("Could not send escape to nvim");
191        self.nvim
192            .input("<escape>")
193            .await
194            .expect("Could not send escape to nvim");
195
196        let nvim_window = self
197            .nvim
198            .get_current_win()
199            .await
200            .expect("Could not get neovim window");
201
202        if selections.len() != 1 {
203            panic!("must have one selection");
204        }
205        let selection = &selections[0];
206
207        let cursor = selection.start;
208        nvim_window
209            .set_cursor((cursor.row as i64 + 1, cursor.column as i64))
210            .await
211            .expect("Could not set nvim cursor position");
212
213        if !selection.is_empty() {
214            self.nvim
215                .input("v")
216                .await
217                .expect("could not enter visual mode");
218
219            let cursor = selection.end;
220            nvim_window
221                .set_cursor((cursor.row as i64 + 1, cursor.column as i64))
222                .await
223                .expect("Could not set nvim cursor position");
224        }
225
226        if let Some(NeovimData::Get { mode, state }) = self.data.back()
227            && *mode == Mode::Normal
228            && *state == marked_text
229        {
230            return;
231        }
232        self.data.push_back(NeovimData::Put {
233            state: marked_text.to_string(),
234        })
235    }
236
237    #[cfg(not(feature = "neovim"))]
238    pub async fn set_state(&mut self, marked_text: &str) {
239        if let Some(NeovimData::Get { mode, state: text }) = self.data.front() {
240            if *mode == Mode::Normal && *text == marked_text {
241                return;
242            }
243            self.data.pop_front();
244        }
245        assert_eq!(
246            self.data.pop_front(),
247            Some(NeovimData::Put {
248                state: marked_text.to_string()
249            }),
250            "operation does not match recorded script. re-record with --features=neovim"
251        );
252    }
253
254    #[cfg(feature = "neovim")]
255    pub async fn set_option(&mut self, value: &str) {
256        self.nvim
257            .command_output(format!("set {}", value).as_str())
258            .await
259            .unwrap();
260
261        self.data.push_back(NeovimData::SetOption {
262            value: value.to_string(),
263        })
264    }
265
266    #[cfg(not(feature = "neovim"))]
267    pub async fn set_option(&mut self, value: &str) {
268        if let Some(NeovimData::Get { .. }) = self.data.front() {
269            self.data.pop_front();
270        };
271        assert_eq!(
272            self.data.pop_front(),
273            Some(NeovimData::SetOption {
274                value: value.to_string(),
275            }),
276            "operation does not match recorded script. re-record with --features=neovim"
277        );
278    }
279
280    #[cfg(feature = "neovim")]
281    pub async fn exec(&mut self, value: &str) {
282        self.nvim.command_output(value).await.unwrap();
283
284        self.data.push_back(NeovimData::Exec {
285            command: value.to_string(),
286        })
287    }
288
289    #[cfg(not(feature = "neovim"))]
290    pub async fn exec(&mut self, value: &str) {
291        if let Some(NeovimData::Get { .. }) = self.data.front() {
292            self.data.pop_front();
293        };
294        assert_eq!(
295            self.data.pop_front(),
296            Some(NeovimData::Exec {
297                command: value.to_string(),
298            }),
299            "operation does not match recorded script. re-record with --features=neovim"
300        );
301    }
302
303    #[cfg(not(feature = "neovim"))]
304    pub async fn read_register(&mut self, register: char) -> String {
305        if let Some(NeovimData::Get { .. }) = self.data.front() {
306            self.data.pop_front();
307        };
308        if let Some(NeovimData::ReadRegister { name, value }) = self.data.pop_front()
309            && name == register
310        {
311            return value;
312        }
313
314        panic!("operation does not match recorded script. re-record with --features=neovim")
315    }
316
317    #[cfg(feature = "neovim")]
318    pub async fn read_register(&mut self, name: char) -> String {
319        let value = self
320            .nvim
321            .command_output(format!("echo getreg('{}')", name).as_str())
322            .await
323            .unwrap();
324
325        self.data.push_back(NeovimData::ReadRegister {
326            name,
327            value: value.clone(),
328        });
329
330        value
331    }
332
333    #[cfg(feature = "neovim")]
334    async fn read_position(&mut self, cmd: &str) -> u32 {
335        self.nvim
336            .command_output(cmd)
337            .await
338            .unwrap()
339            .parse::<u32>()
340            .unwrap()
341    }
342
343    #[cfg(feature = "neovim")]
344    pub async fn state(&mut self) -> (Mode, String) {
345        let nvim_buffer = self
346            .nvim
347            .get_current_buf()
348            .await
349            .expect("Could not get neovim buffer");
350        let text = nvim_buffer
351            .get_lines(0, -1, false)
352            .await
353            .expect("Could not get buffer text")
354            .join("\n");
355
356        // nvim columns are 1-based, so -1.
357        let mut cursor_row = self.read_position("echo line('.')").await - 1;
358        let mut cursor_col = self.read_position("echo col('.')").await - 1;
359        let mut selection_row = self.read_position("echo line('v')").await - 1;
360        let mut selection_col = self.read_position("echo col('v')").await - 1;
361        let total_rows = self.read_position("echo line('$')").await - 1;
362
363        let nvim_mode_text = self
364            .nvim
365            .get_mode()
366            .await
367            .expect("Could not get mode")
368            .into_iter()
369            .find_map(|(key, value)| {
370                if key.as_str() == Some("mode") {
371                    Some(value.as_str().unwrap().to_owned())
372                } else {
373                    None
374                }
375            })
376            .expect("Could not find mode value");
377
378        let mode = match nvim_mode_text.as_ref() {
379            "i" => Mode::Insert,
380            "n" => Mode::Normal,
381            "v" => Mode::Visual,
382            "V" => Mode::VisualLine,
383            "R" => Mode::Replace,
384            "\x16" => Mode::VisualBlock,
385            _ => panic!("unexpected vim mode: {nvim_mode_text}"),
386        };
387
388        let mut selections = Vec::new();
389        // Vim uses the index of the first and last character in the selection
390        // Zed uses the index of the positions between the characters, so we need
391        // to add one to the end in visual mode.
392        match mode {
393            Mode::VisualBlock if selection_row != cursor_row => {
394                // in zed we fake a block selection by using multiple cursors (one per line)
395                // this code emulates that.
396                // to deal with casees where the selection is not perfectly rectangular we extract
397                // the content of the selection via the "a register to get the shape correctly.
398                self.nvim.input("\"aygv").await.unwrap();
399                let content = self.nvim.command_output("echo getreg('a')").await.unwrap();
400                let lines = content.split('\n').collect::<Vec<_>>();
401                let top = cmp::min(selection_row, cursor_row);
402                let left = cmp::min(selection_col, cursor_col);
403                for row in top..=cmp::max(selection_row, cursor_row) {
404                    let content = if row - top >= lines.len() as u32 {
405                        ""
406                    } else {
407                        lines[(row - top) as usize]
408                    };
409                    let line_len = self
410                        .read_position(format!("echo strlen(getline({}))", row + 1).as_str())
411                        .await;
412
413                    if left > line_len {
414                        continue;
415                    }
416
417                    let start = Point::new(row, left);
418                    let end = Point::new(row, left + content.len() as u32);
419                    if cursor_col >= selection_col {
420                        selections.push(start..end)
421                    } else {
422                        selections.push(end..start)
423                    }
424                }
425            }
426            Mode::Visual | Mode::VisualLine | Mode::VisualBlock => {
427                if (selection_row, selection_col) > (cursor_row, cursor_col) {
428                    let selection_line_length =
429                        self.read_position("echo strlen(getline(line('v')))").await;
430                    if selection_line_length > selection_col {
431                        selection_col += 1;
432                    } else if selection_row < total_rows {
433                        selection_col = 0;
434                        selection_row += 1;
435                    }
436                } else {
437                    let cursor_line_length =
438                        self.read_position("echo strlen(getline(line('.')))").await;
439                    if cursor_line_length > cursor_col {
440                        cursor_col += 1;
441                    } else if cursor_row < total_rows {
442                        cursor_col = 0;
443                        cursor_row += 1;
444                    }
445                }
446                selections.push(
447                    Point::new(selection_row, selection_col)..Point::new(cursor_row, cursor_col),
448                )
449            }
450            Mode::Insert | Mode::Normal | Mode::Replace => selections
451                .push(Point::new(selection_row, selection_col)..Point::new(cursor_row, cursor_col)),
452            Mode::HelixNormal | Mode::HelixSelect => unreachable!(),
453        }
454
455        let ranges = encode_ranges(&text, &selections);
456        let state = NeovimData::Get {
457            mode,
458            state: ranges.clone(),
459        };
460
461        if self.data.back() != Some(&state) {
462            self.data.push_back(state);
463        }
464
465        (mode, ranges)
466    }
467
468    #[cfg(not(feature = "neovim"))]
469    pub async fn state(&mut self) -> (Mode, String) {
470        if let Some(NeovimData::Get { state: raw, mode }) = self.data.front() {
471            (*mode, raw.to_string())
472        } else {
473            panic!("operation does not match recorded script. re-record with --features=neovim");
474        }
475    }
476
477    fn test_data_path(test_case_id: &str) -> PathBuf {
478        let mut data_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
479        data_path.push("test_data");
480        data_path.push(format!("{}.json", test_case_id));
481        data_path
482    }
483
484    #[cfg(not(feature = "neovim"))]
485    fn read_test_data(test_case_id: &str) -> VecDeque<NeovimData> {
486        let path = Self::test_data_path(test_case_id);
487        let json = std::fs::read_to_string(path).expect(
488            "Could not read test data. Is it generated? Try running test with '--features neovim'",
489        );
490
491        let mut result = VecDeque::new();
492        for line in json.lines() {
493            result.push_back(
494                serde_json::from_str(line)
495                    .expect("invalid test data. regenerate it with '--features neovim'"),
496            );
497        }
498        result
499    }
500
501    #[cfg(feature = "neovim")]
502    fn write_test_data(test_case_id: &str, data: &VecDeque<NeovimData>) {
503        let path = Self::test_data_path(test_case_id);
504        let mut json = Vec::new();
505        for entry in data {
506            serde_json::to_writer(&mut json, entry).unwrap();
507            json.push(b'\n');
508        }
509        std::fs::create_dir_all(path.parent().unwrap())
510            .expect("could not create test data directory");
511        std::fs::write(path, json).expect("could not write out test data");
512    }
513}
514
515#[cfg(feature = "neovim")]
516impl Deref for NeovimConnection {
517    type Target = Neovim<nvim_rs::compat::tokio::Compat<ChildStdin>>;
518
519    fn deref(&self) -> &Self::Target {
520        &self.nvim
521    }
522}
523
524#[cfg(feature = "neovim")]
525impl DerefMut for NeovimConnection {
526    fn deref_mut(&mut self) -> &mut Self::Target {
527        &mut self.nvim
528    }
529}
530
531#[cfg(feature = "neovim")]
532impl Drop for NeovimConnection {
533    fn drop(&mut self) {
534        Self::write_test_data(&self.test_case_id, &self.data);
535    }
536}
537
538#[cfg(feature = "neovim")]
539#[derive(Clone)]
540struct NvimHandler {}
541
542#[cfg(feature = "neovim")]
543#[async_trait]
544impl Handler for NvimHandler {
545    type Writer = nvim_rs::compat::tokio::Compat<ChildStdin>;
546
547    async fn handle_request(
548        &self,
549        _event_name: String,
550        _arguments: Vec<Value>,
551        _neovim: Neovim<Self::Writer>,
552    ) -> Result<Value, Value> {
553        unimplemented!();
554    }
555
556    async fn handle_notify(
557        &self,
558        _event_name: String,
559        _arguments: Vec<Value>,
560        _neovim: Neovim<Self::Writer>,
561    ) {
562    }
563}
564
565#[cfg(feature = "neovim")]
566fn parse_state(marked_text: &str) -> (String, Vec<Range<Point>>) {
567    let (text, ranges) = util::test::marked_text_ranges(marked_text, true);
568    let point_ranges = ranges
569        .into_iter()
570        .map(|byte_range| {
571            let mut point_range = Point::zero()..Point::zero();
572            let mut ix = 0;
573            let mut position = Point::zero();
574            for c in text.chars().chain(['\0']) {
575                if ix == byte_range.start {
576                    point_range.start = position;
577                }
578                if ix == byte_range.end {
579                    point_range.end = position;
580                }
581                let len_utf8 = c.len_utf8();
582                ix += len_utf8;
583                if c == '\n' {
584                    position.row += 1;
585                    position.column = 0;
586                } else {
587                    position.column += len_utf8 as u32;
588                }
589            }
590            point_range
591        })
592        .collect::<Vec<_>>();
593    (text, point_ranges)
594}
595
596#[cfg(feature = "neovim")]
597fn encode_ranges(text: &str, point_ranges: &Vec<Range<Point>>) -> String {
598    let byte_ranges = point_ranges
599        .iter()
600        .map(|range| {
601            let mut byte_range = 0..0;
602            let mut ix = 0;
603            let mut position = Point::zero();
604            for c in text.chars().chain(['\0']) {
605                if position == range.start {
606                    byte_range.start = ix;
607                }
608                if position == range.end {
609                    byte_range.end = ix;
610                }
611                let len_utf8 = c.len_utf8();
612                ix += len_utf8;
613                if c == '\n' {
614                    position.row += 1;
615                    position.column = 0;
616                } else {
617                    position.column += len_utf8 as u32;
618                }
619            }
620            byte_range
621        })
622        .collect::<Vec<_>>();
623    util::test::generate_marked_text(text, &byte_ranges[..], true)
624}