lb_tantivy/tokenizer/
tokenized_string.rs

1use std::cmp::Ordering;
2use std::io;
3use std::io::{Read, Write};
4
5use common::*;
6
7use crate::tokenizer::{Token, TokenStream};
8
9/// Struct representing pre-tokenized text
10#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Eq, PartialEq)]
11pub struct PreTokenizedString {
12    /// Original text
13    pub text: String,
14    /// Tokens derived from the text
15    pub tokens: Vec<Token>,
16}
17
18impl Ord for PreTokenizedString {
19    fn cmp(&self, other: &Self) -> Ordering {
20        self.text.cmp(&other.text)
21    }
22}
23
24impl PartialOrd for PreTokenizedString {
25    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
26        Some(self.cmp(other))
27    }
28}
29
30impl BinarySerializable for PreTokenizedString {
31    fn serialize<W: Write + ?Sized>(&self, writer: &mut W) -> io::Result<()> {
32        if let Ok(text) = serde_json::to_string(self) {
33            <String as BinarySerializable>::serialize(&text, writer)
34        } else {
35            Err(io::Error::other(
36                "Failed to dump PreTokenizedString to json.",
37            ))
38        }
39    }
40
41    fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
42        let json_text = <String as BinarySerializable>::deserialize(reader)?;
43
44        if let Ok(value) = serde_json::from_str(&json_text) {
45            Ok(value)
46        } else {
47            Err(io::Error::other(
48                "Failed to parse string data as PreTokenizedString.",
49            ))
50        }
51    }
52}
53
54/// [`TokenStream`] implementation which wraps [`PreTokenizedString`]
55pub struct PreTokenizedStream {
56    tokenized_string: PreTokenizedString,
57    current_token: i64,
58}
59
60impl From<PreTokenizedString> for PreTokenizedStream {
61    fn from(s: PreTokenizedString) -> PreTokenizedStream {
62        PreTokenizedStream {
63            tokenized_string: s,
64            current_token: -1,
65        }
66    }
67}
68
69impl TokenStream for PreTokenizedStream {
70    fn advance(&mut self) -> bool {
71        self.current_token += 1;
72        self.current_token < self.tokenized_string.tokens.len() as i64
73    }
74
75    fn token(&self) -> &Token {
76        assert!(
77            self.current_token >= 0,
78            "TokenStream not initialized. You should call advance() at least once."
79        );
80        &self.tokenized_string.tokens[self.current_token as usize]
81    }
82
83    fn token_mut(&mut self) -> &mut Token {
84        assert!(
85            self.current_token >= 0,
86            "TokenStream not initialized. You should call advance() at least once."
87        );
88        &mut self.tokenized_string.tokens[self.current_token as usize]
89    }
90}
91
92#[cfg(test)]
93mod tests {
94
95    use super::*;
96
97    #[test]
98    fn test_tokenized_stream() {
99        let tok_text = PreTokenizedString {
100            text: String::from("A a"),
101            tokens: vec![
102                Token {
103                    offset_from: 0,
104                    offset_to: 1,
105                    position: 0,
106                    text: String::from("A"),
107                    position_length: 1,
108                },
109                Token {
110                    offset_from: 2,
111                    offset_to: 3,
112                    position: 1,
113                    text: String::from("a"),
114                    position_length: 1,
115                },
116            ],
117        };
118
119        let mut token_stream = PreTokenizedStream::from(tok_text.clone());
120
121        for expected_token in tok_text.tokens {
122            assert!(token_stream.advance());
123            assert_eq!(token_stream.token(), &expected_token);
124        }
125        assert!(!token_stream.advance());
126    }
127}
lb_tantivy/tokenizer/tokenized_string.rs

lb_tantivy/tokenizer/
tokenized_string.rs
