diff --git a/src/lib.rs b/src/lib.rs index ea54eac8..d85035af 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -198,6 +198,9 @@ fn _byte_pair_merge(ranks: &HashMap, Rank>, piece: &[u8]) -> Vec<(usize, pub fn byte_pair_encode(piece: &[u8], ranks: &HashMap, Rank>) -> Vec { let piece_len = piece.len(); + if piece_len == 0 { + return vec![]; + } if piece_len == 1 { return vec![ranks[piece]]; } @@ -680,12 +683,19 @@ mod tests { use fancy_regex::Regex; use rustc_hash::FxHashMap as HashMap; - use crate::{Rank, byte_pair_split}; + use crate::{Rank, byte_pair_encode, byte_pair_split}; fn setup_ranks() -> HashMap, Rank> { HashMap::from_iter([(b"ab".to_vec(), 0), (b"cd".to_vec(), 1)]) } + #[test] + fn test_empty_piece() { + let ranks = setup_ranks(); + let res = byte_pair_encode(b"", &ranks); + assert_eq!(res, Vec::::new()); + } + #[test] fn test_simple_characters() { let ranks = setup_ranks();