use languages::{
Amharic, Arabic, Armenian, Bengali, Bulgarian, Burmese, Catalan, Danish, Dutch, English,
Finnish, French, German, Greek, Gujarati, Hindi, Italian, Japanese, Kannada, Kazakh, Language,
Malayalam, Marathi, Polish, Portuguese, Punjabi, Russian, Slovak, Spanish, Tamil, Telugu, Ukrainian,
};
use regex::Regex;
use std::sync::LazyLock;
mod constants;
pub mod languages;
use serde::Serialize;
static PARA_SPLIT_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\n[\r]*\n").unwrap());
#[derive(Debug, Clone, Serialize)]
pub struct SentenceBoundary<'a> {
pub start_index: usize,
pub end_index: usize,
pub start_byte: usize,
pub end_byte: usize,
pub text: &'a str,
pub boundary_symbol: Option<String>,
pub is_paragraph_break: bool,
}
pub fn language_factory(language_code: &str) -> Box<dyn Language> {
let mut current_code = language_code;
let mut visited = std::collections::HashSet::new();
loop {
if visited.contains(current_code) {
current_code = "en"; } else {
visited.insert(current_code);
}
match current_code {
"am" => return Box::new(Amharic {}),
"ar" => return Box::new(Arabic {}),
"bg" => return Box::new(Bulgarian {}),
"bn" => return Box::new(Bengali {}),
"ca" => return Box::new(Catalan {}),
"da" => return Box::new(Danish {}),
"de" => return Box::new(German {}),
"el" => return Box::new(Greek {}),
"en" => return Box::new(English {}),
"es" => return Box::new(Spanish {}),
"fi" => return Box::new(Finnish {}),
"fr" => return Box::new(French {}),
"gu" => return Box::new(Gujarati {}),
"hi" => return Box::new(Hindi {}),
"hy" => return Box::new(Armenian {}),
"it" => return Box::new(Italian {}),
"ja" => return Box::new(Japanese {}),
"kk" => return Box::new(Kazakh {}),
"kn" => return Box::new(Kannada {}),
"ml" => return Box::new(Malayalam {}),
"mr" => return Box::new(Marathi {}),
"my" => return Box::new(Burmese {}),
"nl" => return Box::new(Dutch {}),
"pa" => return Box::new(Punjabi {}),
"pl" => return Box::new(Polish {}),
"pt" => return Box::new(Portuguese {}),
"ru" => return Box::new(Russian {}),
"sk" => return Box::new(Slovak {}),
"ta" => return Box::new(Tamil {}),
"te" => return Box::new(Telugu {}),
"uk" => return Box::new(Ukrainian {}),
_ => {
if let Some(fallbacks) = languages::get_fallbacks(current_code) {
for next_code in fallbacks {
if !visited.contains(next_code) {
current_code = next_code;
break;
}
}
} else {
current_code = "en"; }
}
}
}
}
fn chunk_text(text: &str, chunk_size: usize) -> Vec<&str> {
if chunk_size == 0 || text.len() <= chunk_size {
return vec![text];
}
let mut chunks = Vec::new();
let mut paragraphs = Vec::new();
let mut last_end = 0;
for mat in PARA_SPLIT_REGEX.find_iter(text) {
paragraphs.push((last_end, mat.start()));
last_end = mat.end();
}
if last_end < text.len() {
paragraphs.push((last_end, text.len()));
}
if paragraphs.is_empty() {
return vec![text];
}
let mut current_start = 0;
let mut current_end = 0;
let mut i = 0;
while i < paragraphs.len() {
let (para_start, para_end) = paragraphs[i];
let para_size = para_end - para_start;
if para_size > chunk_size {
if current_end > current_start {
let safe_end = text.ceil_char_boundary(current_end);
chunks.push(&text[current_start..safe_end]);
current_start = 0;
current_end = 0;
}
chunks.push(&text[para_start..para_end]);
i += 1;
continue;
}
if current_end == current_start {
current_start = para_start;
current_end = para_end;
i += 1;
continue;
}
let potential_size = para_end - current_start;
if potential_size > chunk_size {
let safe_end = text.ceil_char_boundary(current_end);
chunks.push(&text[current_start..safe_end]);
current_start = para_start;
current_end = para_end;
} else {
current_end = para_end;
}
i += 1;
}
if current_end > current_start {
let safe_end = text.ceil_char_boundary(current_end);
chunks.push(&text[current_start..safe_end]);
}
chunks
}
pub fn segment<'a>(language_code: &str, text: &'a str) -> Vec<&'a str> {
const CHUNK_SIZE: usize = 10 * 1024;
let language = language_factory(language_code);
if text.len() > CHUNK_SIZE {
let chunks = chunk_text(text, CHUNK_SIZE);
let mut all_sentences = Vec::new();
for chunk in chunks {
let chunk_sentences = language.segment(chunk);
all_sentences.extend(chunk_sentences);
}
all_sentences
} else {
language.segment(text)
}
}
pub fn get_sentence_boundaries<'a>(
language_code: &str,
text: &'a str,
) -> Vec<SentenceBoundary<'a>> {
const CHUNK_SIZE: usize = 10 * 1024;
let language = language_factory(language_code);
if text.len() > CHUNK_SIZE {
let chunks = chunk_text(text, CHUNK_SIZE);
let mut all_boundaries = Vec::new();
let mut chunk_offset = 0;
for chunk in chunks {
let chunk_boundaries = language.get_sentence_boundaries(chunk);
let mut prev_end_index = 0;
for boundary in chunk_boundaries {
let start_byte = boundary.start_byte + chunk_offset;
let end_byte = boundary.end_byte + chunk_offset;
let start_index = if prev_end_index > 0 {
prev_end_index
} else {
text[..start_byte].chars().count()
};
let end_index = start_index + boundary.text.chars().count();
prev_end_index = end_index;
all_boundaries.push(SentenceBoundary {
start_index,
end_index,
start_byte,
end_byte,
text: boundary.text,
boundary_symbol: boundary.boundary_symbol,
is_paragraph_break: boundary.is_paragraph_break,
});
}
chunk_offset += chunk.len();
}
all_boundaries
} else {
language.get_sentence_boundaries(text)
}
}
#[cfg(test)]
mod tests {
use std::fs;
use super::*;
pub fn run_language_tests_for_language(language: &str, test_file: &str) {
let content = fs::read_to_string(test_file).expect("Failed to read test file");
let test_cases: Vec<&str> = content.split("===").collect();
for case in test_cases {
let case = case.trim();
if case.is_empty() || case.starts_with('#') {
continue; }
let parts: Vec<&str> = case.split("---")
.map(|part| part.trim())
.filter(|part| !part.is_empty())
.collect();
if parts.is_empty() {
continue; }
assert_eq!(parts.len(), 2, "Malformed test case: \n{}", case);
let input = parts[0];
let expected: Vec<&str> = parts[1].lines()
.map(|line| line.trim())
.filter(|line| !line.is_empty())
.collect();
let result = segment(language, input);
let trimmed_result: Vec<String> =
result.iter().map(|item| item.trim().to_string()).collect();
assert_eq!(trimmed_result, expected, "Failed for input: \n{}", input);
}
}
#[test]
fn test_urdu_segment() {
run_language_tests_for_language("ur", "tests/ur.txt");
}
#[test]
fn test_chinese_segment() {
run_language_tests_for_language("zh", "tests/zh.txt");
}
#[test]
fn test_chunk_text_basic() {
let text = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph.";
let chunks = chunk_text(text, 20);
assert_eq!(chunks.len(), 3);
assert_eq!(chunks[0], "First paragraph.");
assert_eq!(chunks[1], "Second paragraph.");
assert_eq!(chunks[2], "Third paragraph.");
}
#[test]
fn test_chunk_text_no_paragraph_breaks() {
let text =
"This is a long text without paragraph breaks that should be returned as one chunk.";
let chunks = chunk_text(text, 20);
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0], text);
}
#[test]
fn test_segment_no_word_split_at_chunk_boundary() {
const CHUNK_SIZE: usize = 10 * 1024;
let padding = "a".repeat(CHUNK_SIZE - 4);
let text = format!("{}Christopher.", padding);
assert!(text.len() > CHUNK_SIZE, "test input must exceed chunk size");
assert_eq!(&text[CHUNK_SIZE - 4..CHUNK_SIZE], "Chri");
let sentences = segment("en", &text);
for s in &sentences {
assert!(
!s.trim_end().ends_with("Chri"),
"Word 'Christopher' was split: sentence ends with 'Chri': {:?}",
s
);
assert!(
!s.trim_start().starts_with("stopher"),
"Word 'Christopher' was split: sentence starts with 'stopher': {:?}",
s
);
}
assert!(
sentences.iter().any(|s| s.contains("Christopher")),
"Expected 'Christopher' to appear whole in a sentence, got: {:?}",
sentences.last()
);
}
#[test]
fn test_get_sentence_boundaries_no_word_split_at_chunk_boundary() {
const CHUNK_SIZE: usize = 10 * 1024;
let padding = "a".repeat(CHUNK_SIZE - 4);
let text = format!("{}Christopher.", padding);
assert!(text.len() > CHUNK_SIZE, "test input must exceed chunk size");
let boundaries = get_sentence_boundaries("en", &text);
for b in &boundaries {
assert!(
!b.text.trim_end().ends_with("Chri"),
"Boundary ends with split fragment 'Chri': {:?}",
b.text
);
assert!(
!b.text.trim_start().starts_with("stopher"),
"Boundary starts with split fragment 'stopher': {:?}",
b.text
);
}
let reconstructed: String = boundaries.iter().map(|b| b.text).collect();
assert_eq!(
reconstructed, text,
"Text reconstruction failed after chunking"
);
}
#[test]
fn test_segment_automatic_chunking() {
let small_text = "First sentence. Second sentence.\n\nThird sentence. Fourth sentence.";
let large_text = small_text.repeat(10000);
let result = segment("en", &large_text);
let expected_per_repetition = segment("en", small_text);
assert!(result.len() >= expected_per_repetition.len() * 9000);
let small_result = segment("en", small_text);
assert_eq!(small_result, expected_per_repetition);
}
#[test]
fn test_get_sentence_boundaries_with_paragraph_breaks() {
let text = "Title\n\nSentence 1.\n\nSentence 2.";
let boundaries = get_sentence_boundaries("en", text);
assert!(boundaries.len() >= 2);
for i in 1..boundaries.len() {
assert!(
boundaries[i].start_index >= boundaries[i - 1].end_index,
"Boundary {} starts at {} but previous ends at {}",
i,
boundaries[i].start_index,
boundaries[i - 1].end_index
);
}
let reconstructed: String = boundaries.iter().map(|b| b.text).collect();
assert_eq!(
reconstructed, text,
"Reconstructed text doesn't match original"
);
let paragraph_breaks: Vec<_> = boundaries.iter().filter(|b| b.is_paragraph_break).collect();
assert!(
paragraph_breaks.len() >= 2,
"Expected at least 2 paragraph breaks, found {}",
paragraph_breaks.len()
);
}
#[test]
fn test_get_sentence_boundaries_with_multibyte_cjk() {
let text = "日本語です。\n\n中文文章。";
let boundaries = get_sentence_boundaries("en", text);
assert!(
boundaries.len() >= 2,
"Expected at least 2 boundaries, got {}",
boundaries.len()
);
for i in 1..boundaries.len() {
assert!(
boundaries[i].start_index >= boundaries[i - 1].end_index,
"Boundary {} starts at {} but previous ends at {}",
i,
boundaries[i].start_index,
boundaries[i - 1].end_index
);
}
let reconstructed: String = boundaries.iter().map(|b| b.text).collect();
assert_eq!(
reconstructed, text,
"Reconstructed text doesn't match original.\nOriginal: {:?}\nReconstructed: {:?}",
text, reconstructed
);
}
#[test]
fn test_get_sentence_boundaries_with_emoji() {
let text = "Hello world 👋.\n\nGoodbye 👋.";
let boundaries = get_sentence_boundaries("en", text);
assert!(
boundaries.len() >= 2,
"Expected at least 2 boundaries, got {}",
boundaries.len()
);
let reconstructed: String = boundaries.iter().map(|b| b.text).collect();
assert_eq!(
reconstructed, text,
"Reconstructed text doesn't match original with emojis"
);
for boundary in &boundaries {
assert!(
boundary.start_index <= boundary.end_index,
"Invalid boundary: start_index {} > end_index {}",
boundary.start_index,
boundary.end_index
);
}
}
#[test]
fn test_get_sentence_boundaries_with_mixed_scripts() {
let text = "English text. Café résumé.\n\n日本語のテキスト。";
let boundaries = get_sentence_boundaries("en", text);
let reconstructed: String = boundaries.iter().map(|b| b.text).collect();
assert_eq!(
reconstructed, text,
"Mixed script text failed reconstruction"
);
for i in 1..boundaries.len() {
assert!(
boundaries[i].start_index >= boundaries[i - 1].end_index,
"Boundaries not ordered correctly at index {}",
i
);
}
for (i, boundary) in boundaries.iter().enumerate() {
assert!(
boundary.start_index <= boundary.end_index,
"Boundary {} has invalid indices: {} > {}",
i,
boundary.start_index,
boundary.end_index
);
}
}
#[test]
fn test_get_sentence_boundaries_character_vs_byte_offsets() {
let text = "Short. 日本語.";
let boundaries = get_sentence_boundaries("en", text);
let total_chars = text.chars().count();
let total_bytes = text.len();
let mut last_char_end = 0;
for boundary in &boundaries {
assert_eq!(
boundary.start_index, last_char_end,
"Gap in character indices: expected start {}, got {}",
last_char_end, boundary.start_index
);
last_char_end = boundary.end_index;
}
assert_eq!(
last_char_end, total_chars,
"Character coverage mismatch: ended at {}, total chars = {}",
last_char_end, total_chars
);
let mut last_byte_end = 0;
for boundary in &boundaries {
assert_eq!(
boundary.start_byte, last_byte_end,
"Gap in byte indices: expected start {}, got {}",
last_byte_end, boundary.start_byte
);
last_byte_end = boundary.end_byte;
}
assert_eq!(
last_byte_end, total_bytes,
"Byte coverage mismatch: ended at {}, total bytes = {}",
last_byte_end, total_bytes
);
}
#[test]
fn test_segment_with_multibyte_characters() {
let text = "日本語です。中文文章。";
let sentences = segment("en", text);
assert!(sentences.len() > 0, "Should find at least one sentence");
let reconstructed: String = sentences.join("");
assert_eq!(
reconstructed, text,
"Segment reconstruction failed for multi-byte text"
);
}
#[test]
fn test_segment_quote_followed_by_spaced_dots_does_not_panic() {
let text = "\"x.\" . .";
let sentences = segment("en", text);
assert_eq!(sentences, vec![text]);
}
#[test]
fn test_get_sentence_boundaries_quote_followed_by_spaced_dots_does_not_panic() {
let text = "\"x.\" . .";
let boundaries = get_sentence_boundaries("en", text);
assert_eq!(boundaries.len(), 1);
assert_eq!(boundaries[0].text, text);
assert_eq!(boundaries[0].start_index, 0);
assert_eq!(boundaries[0].end_index, text.chars().count());
assert_eq!(boundaries[0].start_byte, 0);
assert_eq!(boundaries[0].end_byte, text.len());
assert_eq!(boundaries[0].boundary_symbol.as_deref(), Some("."));
assert!(!boundaries[0].is_paragraph_break);
}
#[test]
fn test_boundary_symbol_detection_with_trailing_space() {
let text = "Hello world. This is a test.Another test. And another test.";
let boundaries = get_sentence_boundaries("en", text);
assert!(
boundaries.len() >= 4,
"Expected at least 4 boundaries, got {}",
boundaries.len()
);
let non_paragraph: Vec<_> = boundaries
.iter()
.filter(|b| !b.is_paragraph_break)
.collect();
for (i, boundary) in non_paragraph.iter().enumerate() {
assert!(
boundary.boundary_symbol.is_some(),
"Boundary {} should have a boundary_symbol, got None",
i
);
assert_eq!(
boundary.boundary_symbol.as_deref(),
Some("."),
"Boundary {} should have period as symbol",
i
);
}
}
#[test]
fn test_boundary_symbol_with_multiple_trailing_spaces() {
let text = "Hello. This is another. Yet another.";
let boundaries = get_sentence_boundaries("en", text);
let non_paragraph = boundaries
.iter()
.filter(|b| !b.is_paragraph_break)
.collect::<Vec<_>>();
for (i, boundary) in non_paragraph.iter().enumerate() {
assert_eq!(
boundary.boundary_symbol.as_deref(),
Some("."),
"Boundary {} should have period symbol despite trailing spaces",
i
);
}
}
#[test]
fn test_boundary_symbol_with_mixed_terminators() {
let text = "Hello! How are you? I'm fine. Yes.";
let boundaries = get_sentence_boundaries("en", text);
let non_paragraph = boundaries
.iter()
.filter(|b| !b.is_paragraph_break)
.collect::<Vec<_>>();
let symbols: Vec<_> = non_paragraph
.iter()
.filter_map(|b| b.boundary_symbol.as_deref())
.collect();
assert!(
symbols.iter().any(|&s| s == "!"),
"Should detect exclamation mark"
);
assert!(
symbols.iter().any(|&s| s == "?"),
"Should detect question mark"
);
assert!(symbols.iter().any(|&s| s == "."), "Should detect period");
}
#[test]
fn test_boundary_symbol_with_cjk_terminator() {
let text = "日本語です。中文です。";
let boundaries = get_sentence_boundaries("en", text);
let non_paragraph = boundaries
.iter()
.filter(|b| !b.is_paragraph_break)
.collect::<Vec<_>>();
let with_cjk_stop = non_paragraph
.iter()
.filter(|b| b.boundary_symbol.as_deref() == Some("。"))
.count();
assert!(
with_cjk_stop >= 1,
"Should detect at least one CJK full stop, got {}",
with_cjk_stop
);
}
#[test]
fn test_boundary_symbol_with_tabs_and_spaces() {
let text = "First sentence.\t\n Second sentence. Third one!";
let boundaries = get_sentence_boundaries("en", text);
let non_paragraph = boundaries
.iter()
.filter(|b| !b.is_paragraph_break)
.collect::<Vec<_>>();
for (i, boundary) in non_paragraph.iter().enumerate() {
assert!(
boundary.boundary_symbol.is_some(),
"Boundary {} should have boundary_symbol despite mixed whitespace",
i
);
}
let reconstructed: String = boundaries.iter().map(|b| b.text).collect();
assert_eq!(
reconstructed, text,
"Text reconstruction failed with mixed whitespace"
);
}
}