tenant.openagents/omega

1use gpui::{App, Context, Entity};
2use language::{self, Buffer, BufferEditSource, TransactionId};
3use std::{
4    collections::HashMap,
5    ops::Range,
6    time::{Duration, Instant},
7};
8use sum_tree::Bias;
9use text::BufferId;
10
11use crate::{Anchor, BufferState, MultiBufferOffset};
12
13use super::{Event, MultiBuffer};
14
15#[derive(Clone)]
16pub(super) struct History {
17    next_transaction_id: TransactionId,
18    undo_stack: Vec<Transaction>,
19    redo_stack: Vec<Transaction>,
20    transaction_depth: usize,
21    group_interval: Duration,
22}
23
24impl Default for History {
25    fn default() -> Self {
26        History {
27            next_transaction_id: clock::Lamport::MIN,
28            undo_stack: Vec::new(),
29            redo_stack: Vec::new(),
30            transaction_depth: 0,
31            group_interval: Duration::from_millis(300),
32        }
33    }
34}
35
36#[derive(Clone)]
37struct Transaction {
38    id: TransactionId,
39    buffer_transactions: HashMap<BufferId, text::TransactionId>,
40    first_edit_at: Instant,
41    last_edit_at: Instant,
42    suppress_grouping: bool,
43}
44
45impl History {
46    fn start_transaction(&mut self, now: Instant) -> Option<TransactionId> {
47        self.transaction_depth += 1;
48        if self.transaction_depth == 1 {
49            let id = self.next_transaction_id.tick();
50            self.undo_stack.push(Transaction {
51                id,
52                buffer_transactions: Default::default(),
53                first_edit_at: now,
54                last_edit_at: now,
55                suppress_grouping: false,
56            });
57            Some(id)
58        } else {
59            None
60        }
61    }
62
63    fn end_transaction(
64        &mut self,
65        now: Instant,
66        buffer_transactions: HashMap<BufferId, text::TransactionId>,
67    ) -> bool {
68        assert_ne!(self.transaction_depth, 0);
69        self.transaction_depth -= 1;
70        if self.transaction_depth == 0 {
71            if buffer_transactions.is_empty() {
72                self.undo_stack.pop();
73                false
74            } else {
75                self.redo_stack.clear();
76                let transaction = self.undo_stack.last_mut().unwrap();
77                transaction.last_edit_at = now;
78                for (buffer_id, transaction_id) in buffer_transactions {
79                    transaction
80                        .buffer_transactions
81                        .entry(buffer_id)
82                        .or_insert(transaction_id);
83                }
84                true
85            }
86        } else {
87            false
88        }
89    }
90
91    fn push_transaction<'a, T>(
92        &mut self,
93        buffer_transactions: T,
94        now: Instant,
95        cx: &Context<MultiBuffer>,
96    ) where
97        T: IntoIterator<Item = (&'a Entity<Buffer>, &'a language::Transaction)>,
98    {
99        assert_eq!(self.transaction_depth, 0);
100        let transaction = Transaction {
101            id: self.next_transaction_id.tick(),
102            buffer_transactions: buffer_transactions
103                .into_iter()
104                .map(|(buffer, transaction)| (buffer.read(cx).remote_id(), transaction.id))
105                .collect(),
106            first_edit_at: now,
107            last_edit_at: now,
108            suppress_grouping: false,
109        };
110        if !transaction.buffer_transactions.is_empty() {
111            self.undo_stack.push(transaction);
112            self.redo_stack.clear();
113        }
114    }
115
116    fn finalize_last_transaction(&mut self) {
117        if let Some(transaction) = self.undo_stack.last_mut() {
118            transaction.suppress_grouping = true;
119        }
120    }
121
122    fn forget(&mut self, transaction_id: TransactionId) -> Option<Transaction> {
123        if let Some(ix) = self
124            .undo_stack
125            .iter()
126            .rposition(|transaction| transaction.id == transaction_id)
127        {
128            Some(self.undo_stack.remove(ix))
129        } else if let Some(ix) = self
130            .redo_stack
131            .iter()
132            .rposition(|transaction| transaction.id == transaction_id)
133        {
134            Some(self.redo_stack.remove(ix))
135        } else {
136            None
137        }
138    }
139
140    fn transaction(&self, transaction_id: TransactionId) -> Option<&Transaction> {
141        self.undo_stack
142            .iter()
143            .find(|transaction| transaction.id == transaction_id)
144            .or_else(|| {
145                self.redo_stack
146                    .iter()
147                    .find(|transaction| transaction.id == transaction_id)
148            })
149    }
150
151    fn transaction_mut(&mut self, transaction_id: TransactionId) -> Option<&mut Transaction> {
152        self.undo_stack
153            .iter_mut()
154            .find(|transaction| transaction.id == transaction_id)
155            .or_else(|| {
156                self.redo_stack
157                    .iter_mut()
158                    .find(|transaction| transaction.id == transaction_id)
159            })
160    }
161
162    fn pop_undo(&mut self) -> Option<&mut Transaction> {
163        assert_eq!(self.transaction_depth, 0);
164        if let Some(transaction) = self.undo_stack.pop() {
165            self.redo_stack.push(transaction);
166            self.redo_stack.last_mut()
167        } else {
168            None
169        }
170    }
171
172    fn pop_redo(&mut self) -> Option<&mut Transaction> {
173        assert_eq!(self.transaction_depth, 0);
174        if let Some(transaction) = self.redo_stack.pop() {
175            self.undo_stack.push(transaction);
176            self.undo_stack.last_mut()
177        } else {
178            None
179        }
180    }
181
182    fn remove_from_undo(&mut self, transaction_id: TransactionId) -> Option<&Transaction> {
183        let ix = self
184            .undo_stack
185            .iter()
186            .rposition(|transaction| transaction.id == transaction_id)?;
187        let transaction = self.undo_stack.remove(ix);
188        self.redo_stack.push(transaction);
189        self.redo_stack.last()
190    }
191
192    fn group(&mut self) -> Option<TransactionId> {
193        let mut count = 0;
194        let mut transactions = self.undo_stack.iter();
195        if let Some(mut transaction) = transactions.next_back() {
196            while let Some(prev_transaction) = transactions.next_back() {
197                if !prev_transaction.suppress_grouping
198                    && transaction.first_edit_at - prev_transaction.last_edit_at
199                        <= self.group_interval
200                {
201                    transaction = prev_transaction;
202                    count += 1;
203                } else {
204                    break;
205                }
206            }
207        }
208        self.group_trailing(count)
209    }
210
211    fn group_until(&mut self, transaction_id: TransactionId) {
212        let mut count = 0;
213        for transaction in self.undo_stack.iter().rev() {
214            if transaction.id == transaction_id {
215                self.group_trailing(count);
216                break;
217            } else if transaction.suppress_grouping {
218                break;
219            } else {
220                count += 1;
221            }
222        }
223    }
224
225    fn group_trailing(&mut self, n: usize) -> Option<TransactionId> {
226        let new_len = self.undo_stack.len() - n;
227        let (transactions_to_keep, transactions_to_merge) = self.undo_stack.split_at_mut(new_len);
228        if let Some(last_transaction) = transactions_to_keep.last_mut() {
229            if let Some(transaction) = transactions_to_merge.last() {
230                last_transaction.last_edit_at = transaction.last_edit_at;
231            }
232            for to_merge in transactions_to_merge {
233                for (buffer_id, transaction_id) in &to_merge.buffer_transactions {
234                    last_transaction
235                        .buffer_transactions
236                        .entry(*buffer_id)
237                        .or_insert(*transaction_id);
238                }
239            }
240        }
241
242        self.undo_stack.truncate(new_len);
243        self.undo_stack.last().map(|t| t.id)
244    }
245
246    pub(super) fn transaction_depth(&self) -> usize {
247        self.transaction_depth
248    }
249
250    pub fn set_group_interval(&mut self, group_interval: Duration) {
251        self.group_interval = group_interval;
252    }
253}
254
255impl MultiBuffer {
256    pub fn start_transaction(&mut self, cx: &mut Context<Self>) -> Option<TransactionId> {
257        self.start_transaction_at(Instant::now(), cx)
258    }
259
260    pub fn start_transaction_at(
261        &mut self,
262        now: Instant,
263        cx: &mut Context<Self>,
264    ) -> Option<TransactionId> {
265        if let Some(buffer) = self.as_singleton() {
266            return buffer.update(cx, |buffer, _| buffer.start_transaction_at(now));
267        }
268
269        for BufferState { buffer, .. } in self.buffers.values() {
270            buffer.update(cx, |buffer, _| buffer.start_transaction_at(now));
271        }
272        self.history.start_transaction(now)
273    }
274
275    pub fn last_transaction_id(&self, cx: &App) -> Option<TransactionId> {
276        if let Some(buffer) = self.as_singleton() {
277            buffer
278                .read(cx)
279                .peek_undo_stack()
280                .map(|history_entry| history_entry.transaction_id())
281        } else {
282            let last_transaction = self.history.undo_stack.last()?;
283            Some(last_transaction.id)
284        }
285    }
286
287    pub fn end_transaction(&mut self, cx: &mut Context<Self>) -> Option<TransactionId> {
288        self.end_transaction_at(Instant::now(), cx)
289    }
290
291    pub fn end_transaction_with_source(
292        &mut self,
293        source: BufferEditSource,
294        cx: &mut Context<Self>,
295    ) -> Option<TransactionId> {
296        let now = Instant::now();
297        if let Some(buffer) = self.as_singleton() {
298            return buffer.update(cx, |buffer, cx| {
299                buffer.end_transaction_with_source(source, cx)
300            });
301        }
302
303        let mut buffer_transactions = HashMap::default();
304        for BufferState { buffer, .. } in self.buffers.values() {
305            if let Some(transaction_id) = buffer.update(cx, |buffer, cx| {
306                buffer.end_transaction_with_source(source, cx)
307            }) {
308                buffer_transactions.insert(buffer.read(cx).remote_id(), transaction_id);
309            }
310        }
311
312        if self.history.end_transaction(now, buffer_transactions) {
313            let transaction_id = self.history.group().unwrap();
314            Some(transaction_id)
315        } else {
316            None
317        }
318    }
319
320    pub fn end_transaction_at(
321        &mut self,
322        now: Instant,
323        cx: &mut Context<Self>,
324    ) -> Option<TransactionId> {
325        if let Some(buffer) = self.as_singleton() {
326            return buffer.update(cx, |buffer, cx| buffer.end_transaction_at(now, cx));
327        }
328
329        let mut buffer_transactions = HashMap::default();
330        for BufferState { buffer, .. } in self.buffers.values() {
331            if let Some(transaction_id) =
332                buffer.update(cx, |buffer, cx| buffer.end_transaction_at(now, cx))
333            {
334                buffer_transactions.insert(buffer.read(cx).remote_id(), transaction_id);
335            }
336        }
337
338        if self.history.end_transaction(now, buffer_transactions) {
339            let transaction_id = self.history.group().unwrap();
340            Some(transaction_id)
341        } else {
342            None
343        }
344    }
345
346    pub fn edited_ranges_for_transaction(
347        &self,
348        transaction_id: TransactionId,
349        cx: &App,
350    ) -> Vec<Range<MultiBufferOffset>> {
351        let Some(transaction) = self.history.transaction(transaction_id) else {
352            return Vec::new();
353        };
354
355        let snapshot = self.read(cx);
356        let mut buffer_anchors = Vec::new();
357
358        for (buffer_id, buffer_transaction) in &transaction.buffer_transactions {
359            let Some(buffer) = self.buffer(*buffer_id) else {
360                continue;
361            };
362            let Some(excerpt) = snapshot.first_excerpt_for_buffer(*buffer_id) else {
363                continue;
364            };
365            let buffer_snapshot = buffer.read(cx).snapshot();
366
367            for range in buffer
368                .read(cx)
369                .edited_ranges_for_transaction_id::<usize>(*buffer_transaction)
370            {
371                buffer_anchors.push(Anchor::in_buffer(
372                    excerpt.path_key_index,
373                    buffer_snapshot.anchor_at(range.start, Bias::Left),
374                ));
375                buffer_anchors.push(Anchor::in_buffer(
376                    excerpt.path_key_index,
377                    buffer_snapshot.anchor_at(range.end, Bias::Right),
378                ));
379            }
380        }
381        buffer_anchors.sort_unstable_by(|a, b| a.cmp(b, &snapshot));
382
383        snapshot
384            .summaries_for_anchors(buffer_anchors.iter())
385            .as_chunks::<2>()
386            .0
387            .iter()
388            .map(|&[s, e]| s..e)
389            .collect::<Vec<_>>()
390    }
391
392    pub fn merge_transactions(
393        &mut self,
394        transaction: TransactionId,
395        destination: TransactionId,
396        cx: &mut Context<Self>,
397    ) {
398        if let Some(buffer) = self.as_singleton() {
399            buffer.update(cx, |buffer, _| {
400                buffer.merge_transactions(transaction, destination)
401            });
402        } else if let Some(transaction) = self.history.forget(transaction)
403            && let Some(destination) = self.history.transaction_mut(destination)
404        {
405            for (buffer_id, buffer_transaction_id) in transaction.buffer_transactions {
406                if let Some(destination_buffer_transaction_id) =
407                    destination.buffer_transactions.get(&buffer_id)
408                {
409                    if let Some(state) = self.buffers.get(&buffer_id) {
410                        state.buffer.update(cx, |buffer, _| {
411                            buffer.merge_transactions(
412                                buffer_transaction_id,
413                                *destination_buffer_transaction_id,
414                            )
415                        });
416                    }
417                } else {
418                    destination
419                        .buffer_transactions
420                        .insert(buffer_id, buffer_transaction_id);
421                }
422            }
423        }
424    }
425
426    pub fn finalize_last_transaction(&mut self, cx: &mut Context<Self>) {
427        self.history.finalize_last_transaction();
428        for BufferState { buffer, .. } in self.buffers.values() {
429            buffer.update(cx, |buffer, _| {
430                buffer.finalize_last_transaction();
431            });
432        }
433    }
434
435    pub fn push_transaction<'a, T>(&mut self, buffer_transactions: T, cx: &Context<Self>)
436    where
437        T: IntoIterator<Item = (&'a Entity<Buffer>, &'a language::Transaction)>,
438    {
439        self.history
440            .push_transaction(buffer_transactions, Instant::now(), cx);
441        self.history.finalize_last_transaction();
442    }
443
444    pub fn group_until_transaction(
445        &mut self,
446        transaction_id: TransactionId,
447        cx: &mut Context<Self>,
448    ) {
449        if let Some(buffer) = self.as_singleton() {
450            buffer.update(cx, |buffer, _| {
451                buffer.group_until_transaction(transaction_id)
452            });
453        } else {
454            self.history.group_until(transaction_id);
455        }
456    }
457    pub fn undo(&mut self, cx: &mut Context<Self>) -> Option<TransactionId> {
458        let mut transaction_id = None;
459        if let Some(buffer) = self.as_singleton() {
460            transaction_id = buffer.update(cx, |buffer, cx| buffer.undo(cx));
461        } else {
462            while let Some(transaction) = self.history.pop_undo() {
463                let mut undone = false;
464                for (buffer_id, buffer_transaction_id) in &mut transaction.buffer_transactions {
465                    if let Some(BufferState { buffer, .. }) = self.buffers.get(buffer_id) {
466                        undone |= buffer.update(cx, |buffer, cx| {
467                            let undo_to = *buffer_transaction_id;
468                            if let Some(entry) = buffer.peek_undo_stack() {
469                                *buffer_transaction_id = entry.transaction_id();
470                            }
471                            buffer.undo_to_transaction(undo_to, cx)
472                        });
473                    }
474                }
475
476                if undone {
477                    transaction_id = Some(transaction.id);
478                    break;
479                }
480            }
481        }
482
483        if let Some(transaction_id) = transaction_id {
484            cx.emit(Event::TransactionUndone { transaction_id });
485        }
486
487        transaction_id
488    }
489
490    pub fn redo(&mut self, cx: &mut Context<Self>) -> Option<TransactionId> {
491        if let Some(buffer) = self.as_singleton() {
492            return buffer.update(cx, |buffer, cx| buffer.redo(cx));
493        }
494
495        while let Some(transaction) = self.history.pop_redo() {
496            let mut redone = false;
497            for (buffer_id, buffer_transaction_id) in transaction.buffer_transactions.iter_mut() {
498                if let Some(BufferState { buffer, .. }) = self.buffers.get(buffer_id) {
499                    redone |= buffer.update(cx, |buffer, cx| {
500                        let redo_to = *buffer_transaction_id;
501                        if let Some(entry) = buffer.peek_redo_stack() {
502                            *buffer_transaction_id = entry.transaction_id();
503                        }
504                        buffer.redo_to_transaction(redo_to, cx)
505                    });
506                }
507            }
508
509            if redone {
510                return Some(transaction.id);
511            }
512        }
513
514        None
515    }
516
517    pub fn undo_transaction(&mut self, transaction_id: TransactionId, cx: &mut Context<Self>) {
518        if let Some(buffer) = self.as_singleton() {
519            buffer.update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
520        } else if let Some(transaction) = self.history.remove_from_undo(transaction_id) {
521            for (buffer_id, transaction_id) in &transaction.buffer_transactions {
522                if let Some(BufferState { buffer, .. }) = self.buffers.get(buffer_id) {
523                    buffer.update(cx, |buffer, cx| {
524                        buffer.undo_transaction(*transaction_id, cx)
525                    });
526                }
527            }
528        }
529    }
530
531    pub fn forget_transaction(&mut self, transaction_id: TransactionId, cx: &mut Context<Self>) {
532        if let Some(buffer) = self.as_singleton() {
533            buffer.update(cx, |buffer, _| {
534                buffer.forget_transaction(transaction_id);
535            });
536        } else if let Some(transaction) = self.history.forget(transaction_id) {
537            for (buffer_id, buffer_transaction_id) in transaction.buffer_transactions {
538                if let Some(state) = self.buffers.get_mut(&buffer_id) {
539                    state.buffer.update(cx, |buffer, _| {
540                        buffer.forget_transaction(buffer_transaction_id);
541                    });
542                }
543            }
544        }
545    }
546}
547