use std::io::{self, BufRead, BufReader, Read, Write};
use std::time::{Duration, Instant};
use bytes::BytesMut;
use miette::{IntoDiagnostic, Result, WrapErr};
use crate::defaults::io::*;
use crate::defaults::memory::*;
use crate::defaults::processing::BLOCK_SIZE;
use crate::{config, process_line, strip_line_ending};
pub fn handle_stdin() -> Result<()> {
let offset = config::offset();
let offset_unit = config::offset_unit();
let mut processor = StdinProcessor::new();
if offset == 0 {
return processor.tail();
}
match (offset.is_positive(), offset_unit) {
(true, config::OffsetUnit::Lines) => processor.skip_lines(offset as u64),
(true, config::OffsetUnit::Bytes) => processor.skip_bytes(offset as u64),
(true, config::OffsetUnit::Blocks) => {
let bytes_to_skip = (offset as u64) * BLOCK_SIZE;
processor.skip_bytes(bytes_to_skip)
}
(false, config::OffsetUnit::Lines) => processor.backtrack_lines((-offset) as u64),
(false, config::OffsetUnit::Bytes) => processor.backtrack_bytes((-offset) as u64),
(false, config::OffsetUnit::Blocks) => processor.backtrack_bytes(((-offset) as u64) * BLOCK_SIZE),
}
}
pub struct StdinProcessor<'a> {
inlock: io::StdinLock<'a>,
outlock: io::StdoutLock<'a>,
buffer: BytesMut,
line: String,
count: u16,
}
impl<'a> Default for StdinProcessor<'a> {
fn default() -> Self {
Self::new()
}
}
impl<'a> StdinProcessor<'a> {
pub fn new() -> Self {
Self {
inlock: io::stdin().lock(),
outlock: io::stdout().lock(),
buffer: BytesMut::with_capacity(OUTPUT_BUFFER_CAPACITY),
line: String::with_capacity(LINE_CAPACITY),
count: 0,
}
}
pub fn process_line(&mut self, line: &str) -> Result<()> {
process_line(line, &mut self.buffer, &mut self.outlock).with_context(|| "Failed to process line")?;
self.count += 1;
self.flush_if_needed()
}
pub fn flush_if_needed(&mut self) -> Result<()> {
if self.count >= FLUSH_LINE_COUNT {
self.outlock.flush().into_diagnostic()?;
self.count = 0;
}
Ok(())
}
pub fn flush(&mut self) -> Result<()> {
self.outlock.flush().into_diagnostic()?;
self.count = 0;
Ok(())
}
pub fn read_line(&mut self) -> Result<usize> {
self.line.clear();
let bytes_read = self.inlock.read_line(&mut self.line).into_diagnostic()?;
if bytes_read > 0 {
strip_line_ending(&mut self.line);
}
Ok(bytes_read)
}
pub fn line(&self) -> &str {
&self.line
}
pub fn process_to_end(&mut self) -> Result<()> {
while self.read_line()? != 0 {
let line = self.line().to_string();
self.process_line(&line)?;
}
self.flush()
}
pub fn handle_overshoot(&mut self, overshoot: &[u8]) -> Result<()> {
let mut start = 0;
for (i, &byte) in overshoot.iter().enumerate() {
if byte == b'\n' {
let line_bytes = &overshoot[start..i];
let line = String::from_utf8_lossy(line_bytes);
self.process_line(&line)?;
start = i + 1;
}
}
if start < overshoot.len() {
let remaining_bytes = &overshoot[start..];
let remaining_str = String::from_utf8_lossy(remaining_bytes);
self.line.push_str(&remaining_str);
}
if !self.line.is_empty() && self.inlock.read_line(&mut self.line).into_diagnostic()? > 0 {
strip_line_ending(&mut self.line);
let line = self.line().to_string();
self.process_line(&line)?;
}
self.tail()
}
pub fn tail(&mut self) -> Result<()> {
self.process_to_end()?;
if !config::tailing() {
return Ok(());
}
let mut last_flush = Instant::now();
loop {
std::thread::sleep(Duration::from_millis(100));
match self.read_line()? {
0 => continue, _ => {
let line = self.line().to_string();
self.process_line(&line)?;
if last_flush.elapsed() >= TAIL_FLUSH_INTERVAL {
self.flush()?;
last_flush = Instant::now();
}
}
}
}
}
pub fn skip_lines(&mut self, count: u64) -> Result<()> {
let mut lines_skipped = 0u64;
while lines_skipped < count {
match self.read_line()? {
0 => {
return Ok(());
}
_ => {
lines_skipped += 1;
}
}
}
self.tail()
}
pub fn skip_bytes(&mut self, to_skip: u64) -> Result<()> {
let mut buffer = [0u8; READ_BUFFER_SIZE];
let mut bytes_skipped = 0u64;
while bytes_skipped < to_skip {
let bytes_read = self.inlock.read(&mut buffer).into_diagnostic()?;
if bytes_read == 0 {
return Ok(());
}
let bytes_to_consume = std::cmp::min(bytes_read as u64, to_skip - bytes_skipped);
bytes_skipped += bytes_to_consume;
if bytes_skipped == to_skip && bytes_to_consume < bytes_read as u64 {
let overshoot_start = bytes_to_consume as usize;
let overshoot = &buffer[overshoot_start..bytes_read];
return self.handle_overshoot(overshoot);
}
}
self.tail()
}
pub fn backtrack_bytes(&mut self, bytes_to_show: u64) -> Result<()> {
let mut circular_buffer = CircularByteBuffer::new(bytes_to_show as usize);
loop {
let bytes_read = self.inlock.read(&mut self.buffer).into_diagnostic()?;
if bytes_read == 0 {
break; }
circular_buffer.write(&self.buffer[..bytes_read]);
}
if circular_buffer.is_empty() {
return Ok(()); }
let mut overshoot: Vec<u8> = Vec::new();
let output_bytes = circular_buffer.extract_last_bytes();
let process_this = match find_last_char(output_bytes.as_slice(), b'\n') {
Some(last_line_ending) => {
overshoot = output_bytes[last_line_ending + 1..].to_vec();
output_bytes[..last_line_ending].to_vec()
}
None => output_bytes[..].to_vec(),
};
let output_str = String::from_utf8_lossy(&process_this);
for line in output_str.lines() {
self.process_line(line)?;
}
if !overshoot.is_empty() {
self.handle_overshoot(overshoot.as_slice())
} else {
self.tail()
}
}
pub fn backtrack_lines(&mut self, lines_to_show: u64) -> Result<()> {
use std::collections::VecDeque;
use tempfile::NamedTempFile;
let mut line_buffer: VecDeque<String> = VecDeque::with_capacity(lines_to_show as usize);
let mut memory_used = 0usize;
let mut temp_file: Option<NamedTempFile> = None;
loop {
let bytes_read = self.read_line()?;
if bytes_read == 0 {
break; }
if memory_used > MEMORY_LIMIT_BYTES && temp_file.is_none() {
let mut temp = NamedTempFile::new()
.into_diagnostic()
.wrap_err("Failed to create temporary file for large stdin backtrack")?;
for line in &line_buffer {
writeln!(temp, "{}", line)
.into_diagnostic()
.wrap_err("Failed to write to temporary file")?;
}
temp_file = Some(temp);
line_buffer.clear();
memory_used = 0;
}
match &mut temp_file {
Some(temp) => {
writeln!(temp, "{}", self.line)
.into_diagnostic()
.wrap_err("Failed to write to temporary file")?;
}
None => {
if line_buffer.len() >= lines_to_show as usize {
if let Some(old_line) = line_buffer.pop_front() {
memory_used -= old_line.len();
}
}
memory_used += self.line.len();
line_buffer.push_back(self.line.clone());
}
}
}
match temp_file {
Some(mut temp) => {
temp.flush()
.into_diagnostic()
.wrap_err("Failed to flush temporary file")?;
self.read_last_n_lines_from_temp_file(temp, lines_to_show)?;
}
None => {
for buffered_line in line_buffer {
self.process_line(&buffered_line)?;
}
}
}
self.flush()?;
Ok(())
}
fn read_last_n_lines_from_temp_file(
&mut self,
temp_file: tempfile::NamedTempFile,
lines_to_show: u64,
) -> Result<()> {
use std::collections::VecDeque;
use std::fs::File;
let file = File::open(temp_file.path())
.into_diagnostic()
.wrap_err("Failed to open temporary file for reading")?;
let reader = BufReader::new(file);
let mut line_buffer: VecDeque<String> = VecDeque::with_capacity(lines_to_show as usize);
for line_result in reader.lines() {
let line = line_result
.into_diagnostic()
.wrap_err("Failed to read line from temporary file")?;
if line_buffer.len() >= lines_to_show as usize {
line_buffer.pop_front();
}
line_buffer.push_back(line);
}
for line in line_buffer {
self.process_line(&line)?;
}
Ok(())
}
}
pub struct CircularByteBuffer {
buffer: Vec<u8>,
pos: usize,
total_read: u64,
capacity: usize,
}
impl CircularByteBuffer {
pub fn new(capacity: usize) -> Self {
Self {
buffer: vec![0u8; capacity],
pos: 0,
total_read: 0,
capacity,
}
}
pub fn write(&mut self, data: &[u8]) {
for &byte in data {
self.buffer[self.pos % self.capacity] = byte;
self.pos += 1;
self.total_read += 1;
}
}
pub fn extract_last_bytes(&self) -> Vec<u8> {
if self.total_read == 0 {
return Vec::new();
}
let bytes_to_output = std::cmp::min(self.total_read, self.capacity as u64) as usize;
if self.total_read >= self.capacity as u64 {
let start_pos = self.pos % self.capacity;
let mut result = Vec::with_capacity(bytes_to_output);
for i in 0..bytes_to_output {
result.push(self.buffer[(start_pos + i) % self.capacity]);
}
result
} else {
self.buffer[..bytes_to_output].to_vec()
}
}
pub fn is_empty(&self) -> bool {
self.total_read == 0
}
pub fn total_written(&self) -> u64 {
self.total_read
}
}
fn find_last_char(buffer: &[u8], c: u8) -> Option<usize> {
buffer.iter().rposition(|&b| b == c)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn can_find_last_char() {
assert_eq!(find_last_char(b"hello\nworld\n", b'\n'), Some(11));
assert_eq!(find_last_char(b"hello\nworld", b'\n'), Some(5));
assert_eq!(find_last_char(b"hello world", b'\n'), None);
assert_eq!(find_last_char(b"", b'\n'), None);
assert_eq!(find_last_char(b"\n", b'\n'), Some(0));
assert_eq!(find_last_char(b"a", b'\n'), None);
assert_eq!(find_last_char(b"\n\n\n", b'\n'), Some(2));
}
#[test]
fn circular_buffer_edge_cases() {
let mut buffer = CircularByteBuffer::new(5);
buffer.write(b"12345");
assert_eq!(buffer.extract_last_bytes(), b"12345");
buffer.write(b"67890");
assert_eq!(buffer.extract_last_bytes(), b"67890");
let mut buffer2 = CircularByteBuffer::new(10);
buffer2.write(b"abc");
assert_eq!(buffer2.extract_last_bytes(), b"abc");
buffer2.write(b"def");
buffer2.write(b"ghi");
assert_eq!(buffer2.extract_last_bytes(), b"abcdefghi");
}
#[test]
fn backtracking_with_partial_lines() {
let buffer_with_newline = b"line1\nline2\nline3\n";
assert_eq!(find_last_char(buffer_with_newline, b'\n'), Some(17));
let buffer_without_newline = b"line1\nline2\nline3";
assert_eq!(find_last_char(buffer_without_newline, b'\n'), Some(11));
let buffer_no_newlines = b"single long line without newlines";
assert_eq!(find_last_char(buffer_no_newlines, b'\n'), None);
}
}