nu_protocol/engine/
prompt_state.rs

1use std::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard};
2
3/// Wakes the active interactive line editor and asks it to repaint the prompt in
4/// place, without disturbing the line currently being edited.
5type Repainter = Arc<dyn Fn() + Send + Sync>;
6
7/// Which prompt segment an asynchronous [`PromptState::set`] targets.
8#[derive(Debug, Clone, Copy)]
9pub enum PromptSegment {
10    /// The main (left) prompt, `$env.PROMPT_COMMAND`.
11    Left,
12
13    /// The right prompt, `$env.PROMPT_COMMAND_RIGHT`.
14    Right,
15
16    /// The prompt indicator for the default/emacs edit mode, `$env.PROMPT_INDICATOR`.
17    Indicator,
18
19    /// The vi insert-mode indicator, `$env.PROMPT_INDICATOR_VI_INSERT`.
20    ViInsert,
21
22    /// The vi normal-mode indicator, `$env.PROMPT_INDICATOR_VI_NORMAL`.
23    ViNormal,
24
25    /// The multiline continuation indicator, `$env.PROMPT_MULTILINE_INDICATOR`.
26    Multiline,
27}
28
29/// The full set of rendered prompt strings, as the line editor draws them.
30#[derive(Debug, Default, Clone)]
31pub struct PromptContents {
32    pub left: Option<Arc<str>>,
33    pub right: Option<Arc<str>>,
34    pub indicator: Option<Arc<str>>,
35    pub vi_insert: Option<Arc<str>>,
36    pub vi_normal: Option<Arc<str>>,
37    pub multiline: Option<Arc<str>>,
38    pub render_right_on_last_line: bool,
39}
40
41impl PromptContents {
42    /// Applies an overriding string to a specific segment.
43    pub fn apply_segment_override(&mut self, segment: PromptSegment, content: impl Into<Arc<str>>) {
44        let content = content.into();
45        match segment {
46            PromptSegment::Left => self.left = Some(content),
47            PromptSegment::Right => self.right = Some(content),
48            PromptSegment::Indicator => self.indicator = Some(content),
49            PromptSegment::ViInsert => self.vi_insert = Some(content),
50            PromptSegment::ViNormal => self.vi_normal = Some(content),
51            PromptSegment::Multiline => self.multiline = Some(content),
52        }
53    }
54
55    /// Layer `overrides` on top of `self`, preferring the override where set
56    /// and falling back to `self` otherwise.
57    ///
58    /// Used to render the transient prompt from the *live* baseline (so late
59    /// [`PromptState::set`] pushes still show up) with `TRANSIENT_PROMPT_*`
60    /// values layered on top.
61    pub fn overridden_by(&self, overrides: &PromptContents) -> PromptContents {
62        PromptContents {
63            left: overrides.left.clone().or_else(|| self.left.clone()),
64            right: overrides.right.clone().or_else(|| self.right.clone()),
65            indicator: overrides
66                .indicator
67                .clone()
68                .or_else(|| self.indicator.clone()),
69            vi_insert: overrides
70                .vi_insert
71                .clone()
72                .or_else(|| self.vi_insert.clone()),
73            vi_normal: overrides
74                .vi_normal
75                .clone()
76                .or_else(|| self.vi_normal.clone()),
77            multiline: overrides
78                .multiline
79                .clone()
80                .or_else(|| self.multiline.clone()),
81            render_right_on_last_line: self.render_right_on_last_line,
82        }
83    }
84}
85
86/// Shared, thread-safe home for the interactive prompt's rendered content.
87#[derive(derive_more::Debug, Default)]
88pub struct PromptState {
89    /// Upgraded to RwLock: Enables infinite concurrent reads, locking only for mutations.
90    contents: RwLock<PromptContents>,
91
92    /// Kept in its own lock: it is installed/cleared by the REPL and read by a
93    /// background job's `set`, never together with `contents`.
94    #[debug(skip)]
95    repainter: Mutex<Option<Repainter>>,
96}
97
98impl PromptState {
99    pub fn new() -> Self {
100        Self::default()
101    }
102
103    fn acquire_read_lock(&self) -> RwLockReadGuard<'_, PromptContents> {
104        self.contents
105            .read()
106            .unwrap_or_else(|poisoned_error| poisoned_error.into_inner())
107    }
108
109    fn acquire_write_lock(&self) -> RwLockWriteGuard<'_, PromptContents> {
110        self.contents
111            .write()
112            .unwrap_or_else(|poisoned_error| poisoned_error.into_inner())
113    }
114
115    fn acquire_repainter_lock(&self) -> MutexGuard<'_, Option<Repainter>> {
116        self.repainter
117            .lock()
118            .unwrap_or_else(|poisoned_error| poisoned_error.into_inner())
119    }
120
121    /// Run an action with shared, read-only access to the current contents.
122    pub fn with_contents<ReturnType>(
123        &self,
124        action: impl FnOnce(&PromptContents) -> ReturnType,
125    ) -> ReturnType {
126        action(&self.acquire_read_lock())
127    }
128
129    /// Run an action with exclusive, mutable access to the current contents.
130    fn modify_contents<ReturnType>(
131        &self,
132        action: impl FnOnce(&mut PromptContents) -> ReturnType,
133    ) -> ReturnType {
134        action(&mut self.acquire_write_lock())
135    }
136
137    /// A snapshot of the current contents.
138    pub fn contents(&self) -> PromptContents {
139        self.with_contents(PromptContents::clone)
140    }
141
142    /// Replace all prompt content (the baseline).
143    pub fn set_contents(&self, new_contents: PromptContents) {
144        self.modify_contents(|contents| *contents = new_contents);
145    }
146
147    /// Apply a batch of overrides under a single write lock, then request one
148    /// in-place repaint. Callers touching several segments at once should use this as to not flash-update
149    pub fn apply(&self, overrides: impl FnOnce(&mut PromptContents)) {
150        self.modify_contents(overrides);
151        self.request_repaint();
152    }
153
154    /// Push an override for a single segment and request an in-place repaint.
155    pub fn set(&self, segment: PromptSegment, content: impl Into<Arc<str>>) {
156        let content = content.into();
157        self.apply(|contents| contents.apply_segment_override(segment, content));
158    }
159
160    /// Install or remove the line editor's repainter mechanism.
161    pub fn set_repainter(&self, new_repainter: Option<Repainter>) {
162        *self.acquire_repainter_lock() = new_repainter;
163    }
164
165    /// Fire the installed repainter, explicitly dropping the lock before executing.
166    fn request_repaint(&self) {
167        // Cloning the Option<Arc> locally ensures the MutexGuard drops immediately
168        // at the end of this statement, keeping lock contention to a minimum.
169        let local_repainter = self.acquire_repainter_lock().clone();
170
171        if let Some(repainter) = local_repainter {
172            repainter();
173        }
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use std::sync::atomic::{AtomicUsize, Ordering};
181
182    /// A `PromptState` wired to a repainter that counts how often it fires.
183    fn setup_state_with_counter() -> (Arc<PromptState>, Arc<AtomicUsize>) {
184        let state = Arc::new(PromptState::new());
185        let repainter_count = Arc::new(AtomicUsize::new(0));
186        let counter_reference = Arc::clone(&repainter_count);
187
188        state.set_repainter(Some(Arc::new(move || {
189            counter_reference.fetch_add(1, Ordering::Relaxed);
190        })));
191
192        (state, repainter_count)
193    }
194
195    #[test]
196    fn set_writes_only_the_targeted_segment() {
197        let state = PromptState::new();
198        state.set(PromptSegment::Left, "LeftSegment");
199
200        let contents = state.contents();
201        assert_eq!(contents.left.as_deref(), Some("LeftSegment"));
202        assert_eq!(contents.right, None);
203    }
204
205    #[test]
206    fn indicator_vi_insert_vi_normal_and_multiline_are_independent() {
207        let state = PromptState::new();
208        state.apply(|contents| {
209            contents.apply_segment_override(PromptSegment::Indicator, "Indicator");
210            contents.apply_segment_override(PromptSegment::ViInsert, "ViInsert");
211            contents.apply_segment_override(PromptSegment::ViNormal, "ViNormal");
212            contents.apply_segment_override(PromptSegment::Multiline, "Multiline");
213        });
214
215        let contents = state.contents();
216        assert_eq!(contents.indicator.as_deref(), Some("Indicator"));
217        assert_eq!(contents.vi_insert.as_deref(), Some("ViInsert"));
218        assert_eq!(contents.vi_normal.as_deref(), Some("ViNormal"));
219        assert_eq!(contents.multiline.as_deref(), Some("Multiline"));
220    }
221
222    #[test]
223    fn each_set_triggers_exactly_one_repaint() {
224        let (state, repainter_count) = setup_state_with_counter();
225        state.set(PromptSegment::Left, "Alpha");
226        state.set(PromptSegment::Right, "Beta");
227
228        assert_eq!(repainter_count.load(Ordering::Relaxed), 2);
229    }
230
231    #[test]
232    fn apply_batches_multiple_segments_into_one_repaint() {
233        let (state, repainter_count) = setup_state_with_counter();
234        state.apply(|contents| {
235            contents.apply_segment_override(PromptSegment::Left, "Alpha");
236            contents.apply_segment_override(PromptSegment::Right, "Beta");
237            contents.apply_segment_override(PromptSegment::Indicator, "Gamma");
238        });
239
240        let contents = state.contents();
241        assert_eq!(contents.left.as_deref(), Some("Alpha"));
242        assert_eq!(contents.right.as_deref(), Some("Beta"));
243        assert_eq!(contents.indicator.as_deref(), Some("Gamma"));
244        // Three segments, but a single repaint.
245        assert_eq!(repainter_count.load(Ordering::Relaxed), 1);
246    }
247
248    #[test]
249    fn set_contents_overwrites_a_pushed_override() {
250        let state = PromptState::new();
251        state.set(PromptSegment::Left, "Pushed");
252
253        state.set_contents(PromptContents {
254            left: Some("Baseline".into()),
255            ..Default::default()
256        });
257
258        assert_eq!(state.contents().left.as_deref(), Some("Baseline"));
259    }
260
261    #[test]
262    fn detaching_repainter_stops_repaints() {
263        let (state, repainter_count) = setup_state_with_counter();
264        state.set_repainter(None);
265        state.set(PromptSegment::Left, "Delta");
266
267        assert_eq!(repainter_count.load(Ordering::Relaxed), 0);
268    }
269}
nu_protocol/engine/prompt_state.rs

nu_protocol/engine/
prompt_state.rs
