use crate::sanitize::{sanitize_with_policy, str_display_width, truncate_str, SanitizePolicy};
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct OutputOptions {
pub width: usize,
pub wrap: bool,
pub code_line_numbers: bool,
pub max_rows: usize,
pub sanitize: SanitizePolicy,
}
impl OutputOptions {
pub fn for_width(width: usize) -> Self {
Self {
width: width.max(1),
wrap: true,
code_line_numbers: true,
max_rows: usize::MAX,
sanitize: SanitizePolicy::code_output(SanitizePolicy::UNLIMITED_WIDTH, usize::MAX),
}
}
pub fn with_wrap(mut self, wrap: bool) -> Self {
self.wrap = wrap;
self
}
pub fn with_code_line_numbers(mut self, show: bool) -> Self {
self.code_line_numbers = show;
self
}
pub fn with_max_rows(mut self, max_rows: usize) -> Self {
self.max_rows = max_rows.max(1);
self
}
}
impl Default for OutputOptions {
fn default() -> Self {
Self::for_width(88)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct OutputDocument {
pub rows: Vec<OutputRow>,
pub copy_text: String,
pub removed_control_chars: usize,
pub removed_escape_sequences: usize,
pub removed_zero_width_chars: usize,
pub truncated: bool,
}
impl OutputDocument {
pub fn visible_rows(&self, scroll: usize, height: usize) -> &[OutputRow] {
let height = height.max(1);
let max_scroll = self.rows.len().saturating_sub(height);
let scroll = scroll.min(max_scroll);
let end = self.rows.len().saturating_sub(scroll);
let start = end.saturating_sub(height);
&self.rows[start..end]
}
pub fn code_blocks(&self) -> Vec<CodeBlockOutput> {
let mut blocks = Vec::new();
let mut language = String::new();
let mut lines = Vec::new();
for row in &self.rows {
match row {
OutputRow::CodeFenceStart { language: lang } => {
language = lang.clone();
lines.clear();
}
OutputRow::Code { text, .. } => lines.push(text.clone()),
OutputRow::CodeFenceEnd => {
blocks.push(CodeBlockOutput {
language: if language.is_empty() {
None
} else {
Some(language.clone())
},
code: lines.join("\n"),
line_count: lines.len(),
});
language.clear();
lines.clear();
}
_ => {}
}
}
blocks
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum OutputRow {
Blank,
Paragraph {
text: String,
continuation: bool,
},
Heading {
level: usize,
text: String,
},
Bullet {
text: String,
continuation: bool,
},
Quote {
text: String,
continuation: bool,
},
CodeFenceStart {
language: String,
},
Code {
language: String,
line_number: usize,
gutter: Option<String>,
text: String,
tokens: Vec<CodeToken>,
},
CodeFenceEnd,
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct CodeBlockOutput {
pub language: Option<String>,
pub code: String,
pub line_count: usize,
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct CodeToken {
pub kind: CodeTokenKind,
pub text: String,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum CodeTokenKind {
Text,
Keyword,
String,
Number,
Comment,
Punctuation,
}
pub fn materialize_output(input: &str, options: OutputOptions) -> OutputDocument {
let sanitized = sanitize_with_policy(input, options.sanitize);
let mut rows = Vec::new();
let mut copy_lines = Vec::new();
let mut in_code = false;
let mut code_line = 1usize;
let mut language = String::new();
let mut truncated = sanitized.truncated;
for raw in sanitized.text.lines() {
if rows.len() >= options.max_rows {
truncated = true;
break;
}
let trimmed_end = raw.trim_end();
if let Some(lang) = trimmed_end.strip_prefix("```") {
if in_code {
rows.push(OutputRow::CodeFenceEnd);
copy_lines.push("```".to_string());
in_code = false;
language.clear();
} else {
language = sanitize_info_string(lang.trim());
rows.push(OutputRow::CodeFenceStart {
language: language.clone(),
});
copy_lines.push(if language.is_empty() {
"```".to_string()
} else {
format!("```{language}")
});
in_code = true;
code_line = 1;
}
continue;
}
if in_code {
let code_width = code_text_width(&options, code_line);
let text = if options.wrap {
truncate_str(trimmed_end, code_width)
} else {
trimmed_end.to_string()
};
if str_display_width(trimmed_end) > str_display_width(&text) {
truncated = true;
}
rows.push(OutputRow::Code {
language: language.clone(),
line_number: code_line,
gutter: if options.code_line_numbers {
Some(format!("{:>3} |", code_line))
} else {
None
},
tokens: tokenize_code_line(&text),
text: text.clone(),
});
copy_lines.push(trimmed_end.to_string());
code_line += 1;
continue;
}
let trimmed = trimmed_end.trim_start();
if trimmed.is_empty() {
rows.push(OutputRow::Blank);
copy_lines.push(String::new());
} else if let Some(text) = heading_text(trimmed, 1) {
push_heading(&mut rows, &mut copy_lines, 1, text, options.width);
} else if let Some(text) = heading_text(trimmed, 2) {
push_heading(&mut rows, &mut copy_lines, 2, text, options.width);
} else if let Some(text) = heading_text(trimmed, 3) {
push_heading(&mut rows, &mut copy_lines, 3, text, options.width);
} else if let Some(text) = trimmed
.strip_prefix("- ")
.or_else(|| trimmed.strip_prefix("* "))
.or_else(|| trimmed.strip_prefix("+ "))
{
push_wrapped(
&mut rows,
text,
options.width.saturating_sub(2),
options.wrap,
OutputKind::Bullet,
);
copy_lines.push(format!("- {text}"));
} else if let Some(text) = trimmed.strip_prefix("> ") {
push_wrapped(
&mut rows,
text,
options.width.saturating_sub(2),
options.wrap,
OutputKind::Quote,
);
copy_lines.push(format!("> {text}"));
} else {
push_wrapped(
&mut rows,
trimmed,
options.width,
options.wrap,
OutputKind::Paragraph,
);
copy_lines.push(trimmed.to_string());
}
}
if in_code {
rows.push(OutputRow::CodeFenceEnd);
copy_lines.push("```".to_string());
truncated = true;
}
if rows.is_empty() {
rows.push(OutputRow::Blank);
}
OutputDocument {
rows,
copy_text: copy_lines.join("\n"),
removed_control_chars: sanitized.removed_control_chars,
removed_escape_sequences: sanitized.removed_escape_sequences,
removed_zero_width_chars: sanitized.removed_zero_width_chars,
truncated,
}
}
pub fn tokenize_code_line(line: &str) -> Vec<CodeToken> {
let trimmed = line.trim_start();
if trimmed.starts_with("//") || trimmed.starts_with('#') || trimmed.starts_with("--") {
return vec![CodeToken {
kind: CodeTokenKind::Comment,
text: line.to_string(),
}];
}
let mut tokens = Vec::new();
let mut chars = line.chars().peekable();
while let Some(ch) = chars.next() {
if ch == '"' || ch == '\'' || ch == '`' {
let mut text = String::from(ch);
let mut escaped = false;
for next in chars.by_ref() {
text.push(next);
if escaped {
escaped = false;
} else if next == '\\' {
escaped = true;
} else if next == ch {
break;
}
}
tokens.push(CodeToken {
kind: CodeTokenKind::String,
text,
});
} else if ch.is_ascii_digit() {
let mut text = String::from(ch);
while let Some(next) = chars.peek().copied() {
if next.is_ascii_digit() || matches!(next, '_' | '.') {
text.push(next);
chars.next();
} else {
break;
}
}
tokens.push(CodeToken {
kind: CodeTokenKind::Number,
text,
});
} else if is_ident_start(ch) {
let mut text = String::from(ch);
while let Some(next) = chars.peek().copied() {
if is_ident_continue(next) {
text.push(next);
chars.next();
} else {
break;
}
}
tokens.push(CodeToken {
kind: if is_keyword(&text) {
CodeTokenKind::Keyword
} else {
CodeTokenKind::Text
},
text,
});
} else {
tokens.push(CodeToken {
kind: if ch.is_ascii_punctuation() {
CodeTokenKind::Punctuation
} else {
CodeTokenKind::Text
},
text: ch.to_string(),
});
}
}
tokens
}
#[derive(Clone, Copy)]
enum OutputKind {
Paragraph,
Bullet,
Quote,
}
fn push_heading(
rows: &mut Vec<OutputRow>,
copy_lines: &mut Vec<String>,
level: usize,
text: &str,
width: usize,
) {
rows.push(OutputRow::Heading {
level,
text: truncate_str(text, width.saturating_sub(level + 1)),
});
copy_lines.push(format!("{} {text}", "#".repeat(level)));
}
fn push_wrapped(rows: &mut Vec<OutputRow>, text: &str, width: usize, wrap: bool, kind: OutputKind) {
if !wrap || width == 0 || str_display_width(text) <= width {
rows.push(make_output_row(kind, truncate_str(text, width), false));
return;
}
let mut line = String::new();
let mut line_width = 0usize;
let mut continuation = false;
for word in text.split_whitespace() {
let word_width = str_display_width(word);
let spacing = if line.is_empty() { 0 } else { 1 };
if !line.is_empty() && line_width + spacing + word_width > width {
rows.push(make_output_row(kind, line, continuation));
line = String::new();
line_width = 0;
continuation = true;
}
if !line.is_empty() {
line.push(' ');
line_width += 1;
}
if word_width > width {
line.push_str(&truncate_str(word, width.saturating_sub(line_width)));
line_width = width;
} else {
line.push_str(word);
line_width += word_width;
}
}
if !line.is_empty() {
rows.push(make_output_row(kind, line, continuation));
}
}
fn make_output_row(kind: OutputKind, text: String, continuation: bool) -> OutputRow {
match kind {
OutputKind::Paragraph => OutputRow::Paragraph { text, continuation },
OutputKind::Bullet => OutputRow::Bullet { text, continuation },
OutputKind::Quote => OutputRow::Quote { text, continuation },
}
}
fn heading_text(input: &str, level: usize) -> Option<&str> {
let prefix = match level {
1 => "# ",
2 => "## ",
3 => "### ",
_ => return None,
};
input.strip_prefix(prefix)
}
fn code_text_width(options: &OutputOptions, line_number: usize) -> usize {
if !options.code_line_numbers {
return options.width;
}
let gutter = format!("{:>3} |", line_number);
options.width.saturating_sub(gutter.len() + 1).max(1)
}
fn sanitize_info_string(input: &str) -> String {
input
.chars()
.filter(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-' | '+' | '.' | '#'))
.take(32)
.collect()
}
fn is_ident_start(ch: char) -> bool {
ch == '_' || ch.is_ascii_alphabetic()
}
fn is_ident_continue(ch: char) -> bool {
ch == '_' || ch.is_ascii_alphanumeric()
}
fn is_keyword(token: &str) -> bool {
matches!(
token,
"as" | "async"
| "await"
| "break"
| "case"
| "class"
| "const"
| "continue"
| "crate"
| "default"
| "else"
| "enum"
| "export"
| "false"
| "fn"
| "for"
| "from"
| "function"
| "if"
| "impl"
| "import"
| "in"
| "let"
| "loop"
| "match"
| "mod"
| "mut"
| "pub"
| "return"
| "self"
| "static"
| "struct"
| "super"
| "trait"
| "true"
| "type"
| "use"
| "where"
| "while"
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn materialized_output_strips_ansi_and_parses_code() {
let doc = materialize_output(
"# Hi\n\x1b[31mred\x1b[0m\n```rust\nfn main() { println!(\"ok\"); }\n```",
OutputOptions::for_width(80),
);
assert_eq!(doc.removed_escape_sequences, 2);
assert!(matches!(doc.rows[0], OutputRow::Heading { .. }));
assert_eq!(doc.code_blocks().len(), 1);
assert!(!doc.copy_text.contains('\x1b'));
}
#[test]
fn code_tokens_mark_keywords_strings_and_numbers() {
let tokens = tokenize_code_line("let answer = 42; // comment");
assert!(tokens
.iter()
.any(|token| token.kind == CodeTokenKind::Keyword));
assert!(tokens
.iter()
.any(|token| token.kind == CodeTokenKind::Number));
assert!(tokens
.iter()
.any(|token| token.kind == CodeTokenKind::Punctuation));
}
#[test]
fn visible_rows_clamp_scroll() {
let doc = materialize_output("one\ntwo\nthree\nfour", OutputOptions::for_width(20));
let rows = doc.visible_rows(999, 2);
assert_eq!(rows.len(), 2);
}
}