1 files changed, 35 insertions, 19 deletions

diff --git a/src/ansi/mod.rs b/src/ansi/mod.rs
index 96fb87a..4265b45 100644
--- a/src/ansi/mod.rs
+++ b/src/ansi/mod.rs
@@ -67,39 +67,55 @@ impl TerminalInput {
         self.buf.pop_front()
     }
 
-    fn read_all_input(&mut self) {
-        while {
-            let buf_len = self.buf.len();
-            self.read_internal();
-            self.buf.len() != buf_len
-        } {}
-    }
-
     pub fn query_cursor_position(&mut self) -> Option<CursorPos> {
-        self.read_all_input();
-
         // query to request cursor position
         let mut stdout = std::io::stdout().lock();
         stdout.write_all(b"\x1b[6n").ok()?;
         stdout.flush().ok()?;
         drop(stdout);
 
-        loop {
+        self.read_internal();
+
+        while self
+            .buf
+            .back()
+            .map(|b| !b.as_bytes().ends_with(b"R"))
+            .unwrap_or(true)
+        {
             let buf_len = self.buf.len();
             self.read_internal();
 
             // no new input, give up
             if self.buf.len() <= buf_len {
-                return None;
+                break;
             }
+        }
 
-            // buf.len() > buf_len >= 0, unwrap is safe.
-            let input = self.buf.back().unwrap();
-            if let Some(pos) = try_parse_cursor_position_message(input.as_bytes()) {
-                self.buf.pop_back();
-                return Some(pos);
-            }
+        if self.buf.is_empty() {
+            return None;
+        }
+
+        if !self.buf.back().unwrap().as_bytes().ends_with(b"R") {
+            return None;
         }
+
+        let i = self
+            .buf
+            .iter()
+            .enumerate()
+            .rev()
+            .find(|(_, x)| x.as_bytes().starts_with(b"\x1b"))
+            .map(|x| x.0)?;
+
+        let input = self.buf.iter().skip(i).fold(BString::new(), |mut acc, x| {
+            acc.push_all(x.as_bytes());
+            acc
+        });
+
+        let pos = try_parse_cursor_position_message(&input)?;
+        self.buf.truncate(i);
+
+        Some(pos)
     }
 }
 
@@ -193,7 +209,7 @@ pub struct Escape<'a> {
 
 use terminfo_lean::parse::Terminfo;
 
-use crate::cursor::CursorPos;
+use crate::{BString, PushAll, cursor::CursorPos};
 
 static TERMINFO: RwLock<Option<&'static Terminfo<'static>>> = RwLock::new(None);
 static ESCAPE_TRIE: RwLock<Option<&'static EscapeTrie>> = RwLock::new(None);