use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::task::Context;
use std::task::Poll;
use crossterm::event::Event;
use tokio::sync::broadcast;
use tokio::sync::watch;
use tokio_stream::Stream;
use tokio_stream::wrappers::BroadcastStream;
use tokio_stream::wrappers::WatchStream;
use tokio_stream::wrappers::errors::BroadcastStreamRecvError;
use super::TuiEvent;
pub type EventResult = std::io::Result<Event>;
pub trait EventSource: Send + 'static {
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<EventResult>>;
}
pub struct EventBroker<S: EventSource = CrosstermEventSource> {
state: Mutex<EventBrokerState<S>>,
resume_events_tx: watch::Sender<()>,
}
enum EventBrokerState<S: EventSource> {
Paused, Start, Running(S), }
impl<S: EventSource + Default> EventBrokerState<S> {
fn active_event_source_mut(&mut self) -> Option<&mut S> {
match self {
EventBrokerState::Paused => None,
EventBrokerState::Start => {
*self = EventBrokerState::Running(S::default());
match self {
EventBrokerState::Running(events) => Some(events),
EventBrokerState::Paused | EventBrokerState::Start => unreachable!(),
}
}
EventBrokerState::Running(events) => Some(events),
}
}
}
impl<S: EventSource + Default> EventBroker<S> {
pub fn new() -> Self {
let (resume_events_tx, _resume_events_rx) = watch::channel(());
Self {
state: Mutex::new(EventBrokerState::Start),
resume_events_tx,
}
}
pub fn pause_events(&self) {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
*state = EventBrokerState::Paused;
}
pub fn resume_events(&self) {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
*state = EventBrokerState::Start;
let _ = self.resume_events_tx.send(());
}
pub fn resume_events_rx(&self) -> watch::Receiver<()> {
self.resume_events_tx.subscribe()
}
}
pub struct CrosstermEventSource(pub crossterm::event::EventStream);
impl Default for CrosstermEventSource {
fn default() -> Self {
Self(crossterm::event::EventStream::new())
}
}
impl EventSource for CrosstermEventSource {
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<EventResult>> {
Pin::new(&mut self.get_mut().0).poll_next(cx)
}
}
pub struct TuiEventStream<S: EventSource + Default + Unpin = CrosstermEventSource> {
broker: Arc<EventBroker<S>>,
draw_stream: BroadcastStream<()>,
resume_stream: WatchStream<()>,
terminal_focused: Arc<AtomicBool>,
poll_draw_first: bool,
#[cfg(unix)]
suspend_context: crate::tui_internal::tui::job_control::SuspendContext,
#[cfg(unix)]
alt_screen_active: Arc<AtomicBool>,
}
impl<S: EventSource + Default + Unpin> TuiEventStream<S> {
pub fn new(
broker: Arc<EventBroker<S>>,
draw_rx: broadcast::Receiver<()>,
terminal_focused: Arc<AtomicBool>,
#[cfg(unix)] suspend_context: crate::tui_internal::tui::job_control::SuspendContext,
#[cfg(unix)] alt_screen_active: Arc<AtomicBool>,
) -> Self {
let resume_stream = WatchStream::from_changes(broker.resume_events_rx());
Self {
broker,
draw_stream: BroadcastStream::new(draw_rx),
resume_stream,
terminal_focused,
poll_draw_first: false,
#[cfg(unix)]
suspend_context,
#[cfg(unix)]
alt_screen_active,
}
}
pub fn poll_crossterm_event(&mut self, cx: &mut Context<'_>) -> Poll<Option<TuiEvent>> {
loop {
let poll_result = {
let mut state = self
.broker
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let events = match state.active_event_source_mut() {
Some(events) => events,
None => {
drop(state);
match Pin::new(&mut self.resume_stream).poll_next(cx) {
Poll::Ready(Some(())) => continue,
Poll::Ready(None) => return Poll::Ready(None),
Poll::Pending => return Poll::Pending,
}
}
};
match Pin::new(events).poll_next(cx) {
Poll::Ready(Some(Ok(event))) => Some(event),
Poll::Ready(Some(Err(_))) | Poll::Ready(None) => {
*state = EventBrokerState::Start;
return Poll::Ready(None);
}
Poll::Pending => {
drop(state);
match Pin::new(&mut self.resume_stream).poll_next(cx) {
Poll::Ready(Some(())) => continue,
Poll::Ready(None) => return Poll::Ready(None),
Poll::Pending => return Poll::Pending,
}
}
}
};
if let Some(mapped) = poll_result.and_then(|event| self.map_crossterm_event(event)) {
return Poll::Ready(Some(mapped));
}
}
}
pub fn poll_draw_event(&mut self, cx: &mut Context<'_>) -> Poll<Option<TuiEvent>> {
match Pin::new(&mut self.draw_stream).poll_next(cx) {
Poll::Ready(Some(Ok(()))) => Poll::Ready(Some(TuiEvent::Draw)),
Poll::Ready(Some(Err(BroadcastStreamRecvError::Lagged(_)))) => {
Poll::Ready(Some(TuiEvent::Draw))
}
Poll::Ready(None) => Poll::Ready(None),
Poll::Pending => Poll::Pending,
}
}
fn map_crossterm_event(&mut self, event: Event) -> Option<TuiEvent> {
match event {
Event::Key(key_event) => {
#[cfg(unix)]
if crate::tui_internal::tui::job_control::SUSPEND_KEY.is_press(key_event) {
self.broker.pause_events();
let suspend_result = self.suspend_context.suspend(&self.alt_screen_active);
self.broker.resume_events();
if let Err(err) = suspend_result {
tracing::warn!(
event = "tui_suspend_failed",
error = %err,
"failed to suspend TUI process"
);
}
return Some(TuiEvent::Draw);
}
Some(TuiEvent::Key(key_event))
}
Event::Resize(_, _) => Some(TuiEvent::Resize),
Event::Paste(pasted) => Some(TuiEvent::Paste(pasted)),
Event::FocusGained => {
self.terminal_focused.store(true, Ordering::Relaxed);
crate::tui_internal::terminal_palette::requery_default_colors();
Some(TuiEvent::Draw)
}
Event::FocusLost => {
self.terminal_focused.store(false, Ordering::Relaxed);
None
}
_ => None,
}
}
}
impl<S: EventSource + Default + Unpin> Unpin for TuiEventStream<S> {}
impl<S: EventSource + Default + Unpin> Stream for TuiEventStream<S> {
type Item = TuiEvent;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let draw_first = self.poll_draw_first;
self.poll_draw_first = !self.poll_draw_first;
if draw_first {
if let Poll::Ready(event) = self.poll_draw_event(cx) {
return Poll::Ready(event);
}
if let Poll::Ready(event) = self.poll_crossterm_event(cx) {
return Poll::Ready(event);
}
} else {
if let Poll::Ready(event) = self.poll_crossterm_event(cx) {
return Poll::Ready(event);
}
if let Poll::Ready(event) = self.poll_draw_event(cx) {
return Poll::Ready(event);
}
}
Poll::Pending
}
}
#[cfg(test)]
mod tests {
use super::*;
use crossterm::event::Event;
use crossterm::event::KeyCode;
use crossterm::event::KeyEvent;
use crossterm::event::KeyModifiers;
use pretty_assertions::assert_eq;
use std::task::Context;
use std::task::Poll;
use std::time::Duration;
use tokio::sync::broadcast;
use tokio::sync::mpsc;
use tokio::time::timeout;
use tokio_stream::StreamExt;
struct FakeEventSource {
rx: mpsc::UnboundedReceiver<EventResult>,
tx: mpsc::UnboundedSender<EventResult>,
}
struct FakeEventSourceHandle {
broker: Arc<EventBroker<FakeEventSource>>,
}
impl FakeEventSource {
fn new() -> Self {
let (tx, rx) = mpsc::unbounded_channel();
Self { rx, tx }
}
}
impl Default for FakeEventSource {
fn default() -> Self {
Self::new()
}
}
impl FakeEventSourceHandle {
fn new(broker: Arc<EventBroker<FakeEventSource>>) -> Self {
Self { broker }
}
fn send(&self, event: EventResult) {
let mut state = self
.broker
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let Some(source) = state.active_event_source_mut() else {
return;
};
let _ = source.tx.send(event);
}
}
impl EventSource for FakeEventSource {
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<EventResult>> {
Pin::new(&mut self.get_mut().rx).poll_recv(cx)
}
}
fn make_stream(
broker: Arc<EventBroker<FakeEventSource>>,
draw_rx: broadcast::Receiver<()>,
terminal_focused: Arc<AtomicBool>,
) -> TuiEventStream<FakeEventSource> {
TuiEventStream::new(
broker,
draw_rx,
terminal_focused,
#[cfg(unix)]
crate::tui_internal::tui::job_control::SuspendContext::new(),
#[cfg(unix)]
Arc::new(AtomicBool::new(false)),
)
}
type SetupState = (
Arc<EventBroker<FakeEventSource>>,
FakeEventSourceHandle,
broadcast::Sender<()>,
broadcast::Receiver<()>,
Arc<AtomicBool>,
);
fn setup() -> SetupState {
let source = FakeEventSource::new();
let broker = Arc::new(EventBroker::new());
*broker.state.lock().unwrap() = EventBrokerState::Running(source);
let handle = FakeEventSourceHandle::new(broker.clone());
let (draw_tx, draw_rx) = broadcast::channel(1);
let terminal_focused = Arc::new(AtomicBool::new(true));
(broker, handle, draw_tx, draw_rx, terminal_focused)
}
#[tokio::test(flavor = "current_thread")]
async fn key_event_skips_unmapped() {
let (broker, handle, _draw_tx, draw_rx, terminal_focused) = setup();
let mut stream = make_stream(broker, draw_rx, terminal_focused);
handle.send(Ok(Event::FocusLost));
handle.send(Ok(Event::Key(KeyEvent::new(
KeyCode::Char('a'),
KeyModifiers::NONE,
))));
let next = stream.next().await.unwrap();
match next {
TuiEvent::Key(key) => {
assert_eq!(key, KeyEvent::new(KeyCode::Char('a'), KeyModifiers::NONE));
}
other => panic!("expected key event, got {other:?}"),
}
}
#[tokio::test(flavor = "current_thread")]
async fn draw_and_key_events_yield_both() {
let (broker, handle, draw_tx, draw_rx, terminal_focused) = setup();
let mut stream = make_stream(broker, draw_rx, terminal_focused);
let expected_key = KeyEvent::new(KeyCode::Char('a'), KeyModifiers::NONE);
let _ = draw_tx.send(());
handle.send(Ok(Event::Key(expected_key)));
let first = stream.next().await.unwrap();
let second = stream.next().await.unwrap();
let mut saw_draw = false;
let mut saw_key = false;
for event in [first, second] {
match event {
TuiEvent::Draw => {
saw_draw = true;
}
TuiEvent::Key(key) => {
assert_eq!(key, expected_key);
saw_key = true;
}
other => panic!("expected draw or key event, got {other:?}"),
}
}
assert!(saw_draw && saw_key, "expected both draw and key events");
}
#[tokio::test(flavor = "current_thread")]
async fn lagged_draw_maps_to_draw() {
let (broker, _handle, draw_tx, draw_rx, terminal_focused) = setup();
let mut stream = make_stream(broker, draw_rx.resubscribe(), terminal_focused);
let _ = draw_tx.send(());
let _ = draw_tx.send(());
let first = stream.next().await;
assert!(matches!(first, Some(TuiEvent::Draw)));
}
#[tokio::test(flavor = "current_thread")]
async fn resize_event_maps_to_resize() {
let (broker, handle, _draw_tx, draw_rx, terminal_focused) = setup();
let mut stream = make_stream(broker, draw_rx, terminal_focused);
handle.send(Ok(Event::Resize(80, 24)));
let next = stream.next().await;
assert!(matches!(next, Some(TuiEvent::Resize)));
}
#[tokio::test(flavor = "current_thread")]
async fn error_or_eof_ends_stream() {
let (broker, handle, _draw_tx, draw_rx, terminal_focused) = setup();
let mut stream = make_stream(broker, draw_rx, terminal_focused);
handle.send(Err(std::io::Error::other("boom")));
let next = stream.next().await;
assert!(next.is_none());
}
#[tokio::test(flavor = "current_thread")]
async fn resume_wakes_paused_stream() {
let (broker, handle, _draw_tx, draw_rx, terminal_focused) = setup();
let mut stream = make_stream(broker.clone(), draw_rx, terminal_focused);
broker.pause_events();
let task = tokio::spawn(async move { stream.next().await });
tokio::task::yield_now().await;
broker.resume_events();
let expected_key = KeyEvent::new(KeyCode::Char('r'), KeyModifiers::NONE);
handle.send(Ok(Event::Key(expected_key)));
let event = timeout(Duration::from_millis(100), task)
.await
.expect("timed out waiting for resumed event")
.expect("join failed");
match event {
Some(TuiEvent::Key(key)) => assert_eq!(key, expected_key),
other => panic!("expected key event, got {other:?}"),
}
}
#[tokio::test(flavor = "current_thread")]
async fn resume_wakes_pending_stream() {
let (broker, handle, _draw_tx, draw_rx, terminal_focused) = setup();
let mut stream = make_stream(broker.clone(), draw_rx, terminal_focused);
let task = tokio::spawn(async move { stream.next().await });
tokio::task::yield_now().await;
broker.pause_events();
broker.resume_events();
let expected_key = KeyEvent::new(KeyCode::Char('p'), KeyModifiers::NONE);
handle.send(Ok(Event::Key(expected_key)));
let event = timeout(Duration::from_millis(100), task)
.await
.expect("timed out waiting for resumed event")
.expect("join failed");
match event {
Some(TuiEvent::Key(key)) => assert_eq!(key, expected_key),
other => panic!("expected key event, got {other:?}"),
}
}
}