1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use ik_rs::core::ik_segmenter::{IKSegmenter, TokenMode};
use once_cell::sync::Lazy;

cfg_if::cfg_if! {
    if #[cfg(feature="use-parking-lot")] {
        use parking_lot::RwLock;
    }
    else /*if #[cfg(feature="use-std-sync")]*/ {
        use std::sync::RwLock;
    }
}

use tantivy_tokenizer_api::{Token, TokenStream, Tokenizer};

pub static GLOBAL_IK: Lazy<RwLock<IKSegmenter>> = Lazy::new(|| {
    let ik = IKSegmenter::new();
    RwLock::new(ik)
});

#[derive(Clone)]
pub struct IkTokenizer {
    mode: TokenMode,
}

pub struct IkTokenStream {
    tokens: Vec<Token>,
    index: usize,
}

impl TokenStream for IkTokenStream {
    fn advance(&mut self) -> bool {
        if self.index < self.tokens.len() {
            self.index = self.index + 1;
            true
        } else {
            false
        }
    }
    fn token(&self) -> &Token {
        &self.tokens[self.index - 1]
    }

    fn token_mut(&mut self) -> &mut Token {
        &mut self.tokens[self.index - 1]
    }
}

impl IkTokenizer {
    pub fn new(mode: TokenMode) -> Self {
        Self { mode }
    }
}

impl Tokenizer for IkTokenizer {
    type TokenStream<'a> = IkTokenStream;
    fn token_stream<'a>(&mut self, text: &'a str) -> Self::TokenStream<'a> {
        let mut indices = text.char_indices().collect::<Vec<_>>();
        indices.push((text.len(), '\0'));

        let lock_guard = {cfg_if::cfg_if! {
            if #[cfg(feature="use-parking-lot")] {Some(GLOBAL_IK.read())}
            else /*if #[cfg(feature="use-std-sync")]*/ {
                match GLOBAL_IK.read() {
                    Err(_err) => None,
                    Ok(lck) => Some(lck)
                }
            }
        }};
        let orig_tokens = lock_guard.map_or(vec![],|seg|seg.tokenize(text, self.mode.clone()));

        let mut tokens = Vec::new();
        for token in orig_tokens.iter() {
            tokens.push(Token {
                offset_from: indices[token.begin_pos()].0,
                offset_to: indices[token.end_pos()].0,
                position: token.begin_pos(),
                text: String::from(
                    &text[(indices[token.begin_pos()].0)..(indices[token.end_pos()].0)],
                ),
                position_length: token.len(),
            });
        }
        IkTokenStream { tokens, index: 0 }
    }
}

#[cfg(test)]
mod tests {
    use crate::TokenMode;
    use tantivy_tokenizer_api::{Token, TokenStream, Tokenizer};
    #[test]
    fn tantivy_ik_works() {
        let mut tokenizer = crate::IkTokenizer::new(TokenMode::SEARCH);
        let mut token_stream = tokenizer.token_stream(
            "张华考上了北京大学；李萍进了中等技术学校；我在百货公司当售货员：我们都有光明的前途",
        );
        let mut tokens = Vec::new();
        let mut token_text = Vec::new();
        while let Some(token) = token_stream.next() {
            tokens.push(token.clone());
            token_text.push(token.text.clone());
        }
        // offset should be byte-indexed
        assert_eq!(tokens[0].offset_from, 0);
        assert_eq!(tokens[0].offset_to, "张华".bytes().len());
        assert_eq!(tokens[1].offset_from, "张华".bytes().len());
        // check tokenized text
        assert_eq!(
            token_text,
            vec![
                "张华",
                "考",
                "上了",
                "北京大学",
                "李萍",
                "进了",
                "中等",
                "技术学校",
                "我",
                "在",
                "百货公司",
                "当",
                "售货员",
                "我们",
                "都有",
                "光明",
                "的",
                "前途"
            ]
        );
    }
}