ez_input 0.1.0

// src/scanner.rs
use std::io::{self, BufRead, BufReader, Cursor, Read, Stdin};
use std::marker::PhantomData;
use std::str::FromStr;

/// 高性能 Scanner，专门为 CP 场景优化
pub struct Scanner<R> {
    reader: BufReader<R>,
    buffer: Vec<u8>,
    position: usize,
    capacity: usize,
}

impl Scanner<Stdin> {
    /// 从标准输入创建 Scanner
    #[inline]
    pub fn stdin() -> Self {
        Scanner::new(io::stdin())
    }
}

impl Scanner<Cursor<&'static str>> {
    /// 从标准输入创建 Scanner
    #[inline]
    pub fn str(s: &'static str) -> Self {
        Scanner::new(Cursor::new(s))
    }
}

impl<R: Read> Scanner<R> {
    /// 创建新的 Scanner
    pub fn new(reader: R) -> Self {
        Scanner {
            reader: BufReader::new(reader),
            buffer: Vec::with_capacity(65536), // 64KB 缓冲区
            position: 0,
            capacity: 0,
        }
    }

    /// 读取下一个 token 并解析为指定类型
    #[inline]
    pub fn next<T: FromStr>(&mut self) -> T {
        let token = self.next_token();
        match token.parse::<T>() {
            Ok(val) => val,
            Err(_) => panic!("Failed to parse token: {}", token),
        }
    }

    /// 使用自定义解析器读取 token
    #[inline]
    pub fn next_with_parser<T, F>(&mut self, parser: F) -> T
    where
        F: FnOnce(&str) -> Option<T>,
    {
        let token = self.next_token();
        match parser(token) {
            Some(val) => val,
            None => panic!("Failed to parse token: {}", token),
        }
    }

    /// 读取下一个 token 作为字符串切片
    pub fn next_token(&mut self) -> &str {
        self.skip_whitespace();
        let start = self.position;

        while self.position < self.capacity && !self.buffer[self.position].is_ascii_whitespace() {
            self.position += 1;
        }

        let end = self.position;
        // 安全：我们只读取有效的 UTF-8 数据
        unsafe { std::str::from_utf8_unchecked(&self.buffer[start..end]) }
    }

    /// 读取一行
    pub fn next_line(&mut self) -> String {
        let mut line = String::new();

        // 先从缓冲区读取
        if self.position < self.capacity {
            let start = self.position;
            while self.position < self.capacity && self.buffer[self.position] != b'\n' {
                self.position += 1;
            }

            if let Ok(s) = std::str::from_utf8(&self.buffer[start..self.position]) {
                line.push_str(s);
            }

            // 跳过换行符
            if self.position < self.capacity && self.buffer[self.position] == b'\n' {
                self.position += 1;
            }
        }

        // 如果缓冲区没有换行符，继续从 reader 读取
        if !line.contains('\n') {
            let mut remaining = String::new();
            match self.reader.read_line(&mut remaining) {
                Ok(0) => return line, // EOF
                Ok(_) => {
                    line.push_str(&remaining);
                }
                Err(e) => panic!("IO Error: {}", e),
            }
        }

        // 处理换行符
        if line.ends_with('\n') {
            line.pop();
            if line.ends_with('\r') {
                line.pop();
            }
        }

        line
    }

    /// 跳过空白字符
    fn skip_whitespace(&mut self) {
        loop {
            while self.position < self.capacity && self.buffer[self.position].is_ascii_whitespace()
            {
                self.position += 1;
            }

            if self.position < self.capacity {
                break;
            }

            if !self.refill_buffer() {
                break;
            }
        }
    }

    /// 重新填充缓冲区，返回是否成功读取到数据
    fn refill_buffer(&mut self) -> bool {
        self.buffer.clear();
        self.buffer.resize(65536, 0);

        match self.reader.read(&mut self.buffer) {
            Ok(0) => false, // EOF
            Ok(n) => {
                self.capacity = n;
                self.position = 0;
                self.buffer.truncate(n);
                true
            }
            Err(e) => panic!("IO Error: {}", e),
        }
    }

    /// 读取所有剩余内容
    pub fn read_all(&mut self) -> String {
        let mut result = String::new();

        // 添加缓冲区剩余内容
        if self.position < self.capacity {
            if let Ok(s) = std::str::from_utf8(&self.buffer[self.position..self.capacity]) {
                result.push_str(s);
            }
        }

        // 读取剩余所有内容
        let mut remaining = String::new();
        match self.reader.read_to_string(&mut remaining) {
            Ok(_) => {
                result.push_str(&remaining);
            }
            Err(e) => panic!("IO Error: {}", e),
        }

        result
    }

    /// 检查是否还有更多数据
    pub fn has_next(&mut self) -> bool {
        self.skip_whitespace();
        self.position < self.capacity
    }
}

/// 批量读取的扩展 trait
pub trait ScannerExt {
    fn next_vec<T: FromStr>(&mut self, n: usize) -> Vec<T>;
    fn next_matrix<T: FromStr>(&mut self, rows: usize, cols: usize) -> Vec<Vec<T>>;
}

impl<R: Read> ScannerExt for Scanner<R> {
    /// 读取 Vec<T>
    fn next_vec<T: FromStr>(&mut self, n: usize) -> Vec<T> {
        let mut vec = Vec::with_capacity(n);
        for _ in 0..n {
            vec.push(self.next());
        }
        vec
    }

    /// 读取二维 Vec
    fn next_matrix<T: FromStr>(&mut self, rows: usize, cols: usize) -> Vec<Vec<T>> {
        let mut matrix = Vec::with_capacity(rows);
        for _ in 0..rows {
            let mut row = Vec::with_capacity(cols);
            for _ in 0..cols {
                row.push(self.next());
            }
            matrix.push(row);
        }
        matrix
    }
}

/// 为特定类型提供快捷读取方法
pub trait ReadNext {
    fn next_i32(&mut self) -> i32;
    fn next_i64(&mut self) -> i64;
    fn next_usize(&mut self) -> usize;
    fn next_f64(&mut self) -> f64;
    fn next_string(&mut self) -> String;
}

impl<R: Read> ReadNext for Scanner<R> {
    #[inline]
    fn next_i32(&mut self) -> i32 {
        self.next()
    }
    #[inline]
    fn next_i64(&mut self) -> i64 {
        self.next()
    }
    #[inline]
    fn next_usize(&mut self) -> usize {
        self.next()
    }
    #[inline]
    fn next_f64(&mut self) -> f64 {
        self.next()
    }
    #[inline]
    fn next_string(&mut self) -> String {
        self.next_token().to_string()
    }
}

/// 带迭代器支持的 Scanner
pub struct ScannerIter<'a, R: Read, T> {
    scanner: &'a mut Scanner<R>,
    _marker: PhantomData<T>,
}

impl<'a, R: Read, T> Iterator for ScannerIter<'a, R, T>
where
    T: FromStr,
{
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.scanner.has_next() {
            Some(self.scanner.next())
        } else {
            None
        }
    }
}

impl<R: Read> Scanner<R> {
    /// 创建迭代器
    pub fn iter<T: FromStr>(&mut self) -> ScannerIter<'_, R, T> {
        ScannerIter {
            scanner: self,
            _marker: PhantomData,
        }
    }
}