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