IdeA/crates/backend/src/stream.rs

use std::collections::HashMap;
use std::hash::Hash;
use std::sync::{Arc, Mutex};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum OutputSinkError {
    Closed,
    Full,
    Other(String),
}

pub trait OutputSink<T>: Send + Sync + 'static {
    fn send(&self, item: T) -> Result<(), OutputSinkError>;
}

type SharedOutputSink<T> = Arc<dyn OutputSink<T>>;
type SinkRegistry<K, T> = HashMap<K, (u64, SharedOutputSink<T>)>;

pub struct OutputBridge<K, T>
where
    K: Copy + Eq + Hash,
{
    sinks: Mutex<SinkRegistry<K, T>>,
}

impl<K, T> Default for OutputBridge<K, T>
where
    K: Copy + Eq + Hash,
    T: 'static,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<K, T> OutputBridge<K, T>
where
    K: Copy + Eq + Hash,
    T: 'static,
{
    #[must_use]
    pub fn new() -> Self {
        Self {
            sinks: Mutex::new(HashMap::new()),
        }
    }

    pub fn register(&self, key: K, sink: SharedOutputSink<T>) -> u64 {
        if let Ok(mut map) = self.sinks.lock() {
            let gen = map.get(&key).map_or(0, |(g, _)| g.wrapping_add(1));
            map.insert(key, (gen, sink));
            gen
        } else {
            0
        }
    }

    pub fn unregister(&self, key: &K) {
        if let Ok(mut map) = self.sinks.lock() {
            map.remove(key);
        }
    }

    pub fn unregister_if(&self, key: &K, gen: u64) {
        if let Ok(mut map) = self.sinks.lock() {
            if matches!(map.get(key), Some((g, _)) if *g == gen) {
                map.remove(key);
            }
        }
    }

    pub fn try_send_output(&self, key: &K, item: T) -> Result<(), OutputSinkError> {
        let sink = {
            let Ok(map) = self.sinks.lock() else {
                return Err(OutputSinkError::Closed);
            };
            map.get(key)
                .map(|(_, sink)| Arc::clone(sink))
                .ok_or(OutputSinkError::Closed)?
        };
        sink.send(item)
    }

    pub fn send_output(&self, key: &K, item: T) -> bool {
        self.try_send_output(key, item).is_ok()
    }

    #[must_use]
    pub fn active_sessions(&self) -> usize {
        self.sinks.lock().map(|m| m.len()).unwrap_or(0)
    }
}

struct ReplayEntry<T> {
    generation: u64,
    sink: Option<SharedOutputSink<T>>,
    scrollback: Vec<T>,
}

pub struct ReplayOutputBridge<K, T>
where
    K: Copy + Eq + Hash,
{
    entries: Mutex<HashMap<K, ReplayEntry<T>>>,
    max_chunks: usize,
    max_bytes: usize,
    measure: fn(&T) -> usize,
}

impl<K, T> ReplayOutputBridge<K, T>
where
    K: Copy + Eq + Hash,
    T: Clone + 'static,
{
    #[must_use]
    pub fn new(max_chunks: usize, max_bytes: usize, measure: fn(&T) -> usize) -> Self {
        Self {
            entries: Mutex::new(HashMap::new()),
            max_chunks,
            max_bytes,
            measure,
        }
    }

    pub fn register(&self, key: K, sink: SharedOutputSink<T>) -> u64 {
        if let Ok(mut map) = self.entries.lock() {
            match map.get_mut(&key) {
                Some(entry) => {
                    entry.generation = entry.generation.wrapping_add(1);
                    entry.sink = Some(sink);
                    entry.generation
                }
                None => {
                    map.insert(
                        key,
                        ReplayEntry {
                            generation: 0,
                            sink: Some(sink),
                            scrollback: Vec::new(),
                        },
                    );
                    0
                }
            }
        } else {
            0
        }
    }

    #[must_use]
    pub fn scrollback(&self, key: &K) -> Vec<T> {
        self.entries
            .lock()
            .ok()
            .and_then(|m| m.get(key).map(|e| e.scrollback.clone()))
            .unwrap_or_default()
    }

    pub fn clear_scrollback(&self, key: &K) -> usize {
        self.entries
            .lock()
            .ok()
            .and_then(|mut m| {
                m.get_mut(key).map(|entry| {
                    let cleared = entry.scrollback.len();
                    entry.scrollback.clear();
                    cleared
                })
            })
            .unwrap_or_default()
    }

    /// Moves the retained stream state from `from` to `to`.
    ///
    /// When `to` already exists, its live sink wins and the source scrollback is
    /// prepended to the destination scrollback. This is used when a logical chat
    /// survives a provider restart under a fresh transport session id.
    pub fn rekey(&self, from: &K, to: K) {
        if *from == to {
            return;
        }
        if let Ok(mut map) = self.entries.lock() {
            let Some(source) = map.remove(from) else {
                return;
            };
            match map.get_mut(&to) {
                Some(destination) => {
                    // Invalidate both pre-rekey pump generations. In particular,
                    // the dying source pump must never detach a sink that was
                    // already attached to the replacement key.
                    let mut generation = destination.generation.wrapping_add(1);
                    if generation == source.generation {
                        generation = generation.wrapping_add(1);
                    }
                    destination.generation = generation;
                    let mut scrollback = source.scrollback;
                    scrollback.append(&mut destination.scrollback);
                    destination.scrollback = scrollback;
                    trim_scrollback(destination, self.max_chunks, self.max_bytes, self.measure);
                }
                None => {
                    map.insert(to, source);
                }
            }
        }
    }

    pub fn unregister(&self, key: &K) {
        if let Ok(mut map) = self.entries.lock() {
            map.remove(key);
        }
    }

    pub fn detach(&self, key: &K) {
        if let Ok(mut map) = self.entries.lock() {
            if let Some(entry) = map.get_mut(key) {
                entry.sink = None;
                entry.generation = entry.generation.wrapping_add(1);
            }
        }
    }

    pub fn detach_if(&self, key: &K, gen: u64) {
        if let Ok(mut map) = self.entries.lock() {
            if let Some(entry) = map.get_mut(key) {
                if entry.generation == gen {
                    entry.sink = None;
                }
            }
        }
    }

    pub fn try_send_output(&self, key: &K, item: T) -> Result<(), OutputSinkError> {
        let sink = {
            let Ok(mut map) = self.entries.lock() else {
                return Err(OutputSinkError::Closed);
            };
            let Some(entry) = map.get_mut(key) else {
                return Err(OutputSinkError::Closed);
            };
            entry.scrollback.push(item.clone());
            trim_scrollback(entry, self.max_chunks, self.max_bytes, self.measure);
            entry.sink.as_ref().map(Arc::clone)
        };

        match sink {
            Some(sink) => sink.send(item),
            None => Err(OutputSinkError::Closed),
        }
    }

    pub fn send_output(&self, key: &K, item: T) -> bool {
        self.try_send_output(key, item).is_ok()
    }

    #[must_use]
    pub fn active_sessions(&self) -> usize {
        self.entries.lock().map(|m| m.len()).unwrap_or(0)
    }
}

fn trim_scrollback<T>(
    entry: &mut ReplayEntry<T>,
    max_chunks: usize,
    max_bytes: usize,
    measure: fn(&T) -> usize,
) {
    while !entry.scrollback.is_empty()
        && (entry.scrollback.len() > max_chunks
            || entry.scrollback.iter().map(measure).sum::<usize>() > max_bytes)
    {
        entry.scrollback.remove(0);
    }
}