mecab_ko_dict/trie/
owned.rs

1//! Owned-byte Trie + `TrieBuilder`
2
3use std::borrow::Cow;
4#[cfg(feature = "zstd")]
5use std::io::{Read, Write as IoWrite};
6use std::path::Path;
7
8use yada::{builder::DoubleArrayBuilder, DoubleArray};
9
10use crate::error::{DictError, Result};
11
12use super::backend::PrefixSearchResult;
13
14/// Double-Array Trie
15///
16/// 문자열 키를 효율적으로 검색하는 자료구조입니다.
17/// 형태소 분석에서 사전 검색에 사용됩니다.
18pub struct Trie<'a> {
19    /// 내부 Double-Array
20    da: DoubleArray<Cow<'a, [u8]>>,
21}
22
23impl<'a> Trie<'a> {
24    /// 바이트 슬라이스에서 Trie 생성
25    #[must_use]
26    pub fn new(bytes: &'a [u8]) -> Self {
27        Self {
28            da: DoubleArray::new(Cow::Borrowed(bytes)),
29        }
30    }
31
32    /// 소유 바이트 벡터에서 Trie 생성
33    #[must_use]
34    pub fn from_vec(bytes: Vec<u8>) -> Trie<'static> {
35        Trie {
36            da: DoubleArray::new(Cow::Owned(bytes)),
37        }
38    }
39
40    /// 파일에서 Trie 로드
41    ///
42    /// # Errors
43    ///
44    /// 파일을 읽을 수 없는 경우 에러를 반환합니다.
45    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Trie<'static>> {
46        let bytes = std::fs::read(path.as_ref()).map_err(DictError::Io)?;
47        Ok(Self::from_vec(bytes))
48    }
49
50    /// 압축된 파일에서 Trie 로드 (zstd)
51    ///
52    /// # Errors
53    ///
54    /// 파일을 읽거나 압축 해제할 수 없는 경우 에러를 반환합니다.
55    #[cfg(feature = "zstd")]
56    pub fn from_compressed_file<P: AsRef<Path>>(path: P) -> Result<Trie<'static>> {
57        let file = std::fs::File::open(path.as_ref()).map_err(DictError::Io)?;
58        let mut decoder = zstd::Decoder::new(file).map_err(DictError::Io)?;
59        let mut bytes = Vec::new();
60        decoder.read_to_end(&mut bytes).map_err(DictError::Io)?;
61        Ok(Self::from_vec(bytes))
62    }
63
64    /// 압축된 파일에서 Trie 로드 (zstd feature 비활성화 시)
65    ///
66    /// # Errors
67    ///
68    /// zstd feature가 비활성화된 경우 항상 에러를 반환합니다.
69    #[cfg(not(feature = "zstd"))]
70    pub fn from_compressed_file<P: AsRef<Path>>(_path: P) -> Result<Trie<'static>> {
71        Err(DictError::Format(
72            "zstd feature is not enabled. Use uncompressed files or enable the 'zstd' feature."
73                .to_string(),
74        ))
75    }
76
77    /// 정확히 일치하는 키 검색
78    ///
79    /// ```rust
80    /// use mecab_ko_dict::trie::{Trie, TrieBuilder};
81    ///
82    /// let bytes = TrieBuilder::build(&[("가다", 1u32)]).unwrap();
83    /// let trie = Trie::new(&bytes);
84    /// assert_eq!(trie.exact_match("가다"), Some(1));
85    /// ```
86    #[must_use]
87    pub fn exact_match(&self, key: &str) -> Option<u32> {
88        self.da.exact_match_search(key.as_bytes())
89    }
90
91    /// 바이트 키로 정확히 일치하는 키 검색
92    #[must_use]
93    pub fn exact_match_bytes(&self, key: &[u8]) -> Option<u32> {
94        self.da.exact_match_search(key)
95    }
96
97    /// 공통 접두사 검색
98    ///
99    /// 주어진 텍스트의 접두사와 일치하는 모든 키를 찾습니다.
100    ///
101    /// ```rust
102    /// use mecab_ko_dict::trie::{Trie, TrieBuilder};
103    ///
104    /// let entries = vec![("가", 0u32), ("가방", 2)];
105    /// let bytes = TrieBuilder::build(&entries).unwrap();
106    /// let trie = Trie::new(&bytes);
107    /// let results: Vec<_> = trie.common_prefix_search("가방에").collect();
108    /// assert_eq!(results.len(), 2);
109    /// ```
110    pub fn common_prefix_search<'b>(
111        &'b self,
112        text: &'b str,
113    ) -> impl Iterator<Item = (u32, usize)> + 'b {
114        self.da.common_prefix_search(text.as_bytes())
115    }
116
117    /// 바이트 키로 공통 접두사 검색
118    pub fn common_prefix_search_bytes<'b>(
119        &'b self,
120        key: &'b [u8],
121    ) -> impl Iterator<Item = (u32, usize)> + 'b {
122        self.da.common_prefix_search(key)
123    }
124
125    /// 텍스트의 특정 위치에서 공통 접두사 검색
126    #[must_use]
127    pub fn common_prefix_search_at(
128        &self,
129        text: &str,
130        start_byte: usize,
131    ) -> PrefixSearchResult {
132        if start_byte >= text.len() {
133            return PrefixSearchResult::new();
134        }
135
136        let suffix = &text[start_byte..];
137        self.da
138            .common_prefix_search(suffix.as_bytes())
139            .map(|(value, len)| (value, start_byte + len))
140            .collect()
141    }
142}
143
144/// Trie 빌더
145///
146/// 키-값 쌍에서 Double-Array Trie를 빌드합니다.
147pub struct TrieBuilder;
148
149impl TrieBuilder {
150    /// 정렬된 키-값 쌍에서 Trie 빌드
151    ///
152    /// # Errors
153    ///
154    /// 엔트리가 비어있거나 Trie 빌드에 실패한 경우 에러를 반환합니다.
155    ///
156    /// ```rust
157    /// use mecab_ko_dict::trie::TrieBuilder;
158    ///
159    /// let entries = vec![("가", 0u32), ("가다", 1), ("가방", 2)];
160    /// let bytes = TrieBuilder::build(&entries).unwrap();
161    /// assert!(!bytes.is_empty());
162    /// ```
163    pub fn build(entries: &[(&str, u32)]) -> Result<Vec<u8>> {
164        if entries.is_empty() {
165            return Err(DictError::Format(
166                "Cannot build Trie from empty entries".to_string(),
167            ));
168        }
169
170        let keyset: Vec<_> = entries.iter().map(|(k, v)| (k.as_bytes(), *v)).collect();
171
172        DoubleArrayBuilder::build(&keyset)
173            .ok_or_else(|| DictError::Format("Failed to build Double-Array Trie".to_string()))
174    }
175
176    /// 바이트 키-값 쌍에서 Trie 빌드
177    ///
178    /// # Errors
179    ///
180    /// 엔트리가 비어있거나 Trie 빌드에 실패한 경우 에러를 반환합니다.
181    pub fn build_bytes(entries: &[(&[u8], u32)]) -> Result<Vec<u8>> {
182        if entries.is_empty() {
183            return Err(DictError::Format(
184                "Cannot build Trie from empty entries".to_string(),
185            ));
186        }
187
188        DoubleArrayBuilder::build(entries)
189            .ok_or_else(|| DictError::Format("Failed to build Double-Array Trie".to_string()))
190    }
191
192    /// 정렬되지 않은 엔트리에서 Trie 빌드
193    ///
194    /// # Errors
195    ///
196    /// 엔트리가 비어있거나 Trie 빌드에 실패한 경우 에러를 반환합니다.
197    pub fn build_unsorted(entries: &mut [(&str, u32)]) -> Result<Vec<u8>> {
198        entries.sort_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes()));
199        Self::build(entries)
200    }
201
202    /// Trie를 파일로 저장
203    ///
204    /// # Errors
205    ///
206    /// 파일을 쓸 수 없는 경우 에러를 반환합니다.
207    pub fn save_to_file<P: AsRef<Path>>(bytes: &[u8], path: P) -> Result<()> {
208        std::fs::write(path.as_ref(), bytes).map_err(DictError::Io)
209    }
210
211    /// Trie를 압축하여 파일로 저장 (zstd)
212    ///
213    /// # Errors
214    ///
215    /// 파일을 쓰거나 압축할 수 없는 경우 에러를 반환합니다.
216    #[cfg(feature = "zstd")]
217    pub fn save_to_compressed_file<P: AsRef<Path>>(
218        bytes: &[u8],
219        path: P,
220        level: i32,
221    ) -> Result<()> {
222        let file = std::fs::File::create(path.as_ref()).map_err(DictError::Io)?;
223        let mut encoder = zstd::Encoder::new(file, level).map_err(DictError::Io)?;
224        encoder.write_all(bytes).map_err(DictError::Io)?;
225        encoder.finish().map_err(DictError::Io)?;
226        Ok(())
227    }
228
229    /// Trie를 압축하여 파일로 저장 (zstd feature 비활성화 시)
230    ///
231    /// # Errors
232    ///
233    /// zstd feature가 비활성화된 경우 항상 에러를 반환합니다.
234    #[cfg(not(feature = "zstd"))]
235    pub fn save_to_compressed_file<P: AsRef<Path>>(
236        _bytes: &[u8],
237        _path: P,
238        _level: i32,
239    ) -> Result<()> {
240        Err(DictError::Format(
241            "zstd feature is not enabled. Use uncompressed files or enable the 'zstd' feature."
242                .to_string(),
243        ))
244    }
245}
mecab_ko_dict/trie/owned.rs

mecab_ko_dict/trie/
owned.rs
