icu_segmenter/complex/
mod.rs

1// This file is part of ICU4X. For terms of use, please see the file
2// called LICENSE at the top level of the ICU4X source tree
3// (online at: https://github.com/unicode-org/icu4x/blob/main/LICENSE ).
4
5use crate::provider::*;
6use crate::{GraphemeClusterSegmenter, GraphemeClusterSegmenterBorrowed};
7use alloc::vec::Vec;
8use icu_provider::prelude::*;
9
10mod dictionary;
11use dictionary::*;
12mod language;
13use language::*;
14#[cfg(feature = "lstm")]
15mod lstm;
16#[cfg(feature = "lstm")]
17use lstm::*;
18
19#[derive(Debug, Clone)]
20#[expect(clippy::large_enum_variant)]
21enum DictOrLstm {
22    Dict(DataPayload<UCharDictionaryBreakDataV1>),
23    #[cfg(feature = "lstm")]
24    Lstm(DataPayload<SegmenterLstmAutoV1>),
25}
26
27#[derive(Debug, Clone, Copy)]
28enum DictOrLstmBorrowed<'data> {
29    Dict(&'data UCharDictionaryBreakData<'data>),
30    #[cfg(feature = "lstm")]
31    Lstm(&'data LstmData<'data>),
32}
33
34fn borrow_dictor(dict_or: &DictOrLstm) -> DictOrLstmBorrowed<'_> {
35    match dict_or {
36        DictOrLstm::Dict(dict) => DictOrLstmBorrowed::Dict(dict.get()),
37        #[cfg(feature = "lstm")]
38        DictOrLstm::Lstm(lstm) => DictOrLstmBorrowed::Lstm(lstm.get()),
39    }
40}
41
42fn fromstatic_dictor(dict_or: DictOrLstmBorrowed<'static>) -> DictOrLstm {
43    match dict_or {
44        DictOrLstmBorrowed::Dict(dict) => DictOrLstm::Dict(DataPayload::from_static_ref(dict)),
45        #[cfg(feature = "lstm")]
46        DictOrLstmBorrowed::Lstm(lstm) => DictOrLstm::Lstm(DataPayload::from_static_ref(lstm)),
47    }
48}
49
50#[derive(Debug)]
51pub(crate) struct ComplexPayloads {
52    grapheme: GraphemeClusterSegmenter,
53    my: Option<DictOrLstm>,
54    km: Option<DictOrLstm>,
55    lo: Option<DictOrLstm>,
56    th: Option<DictOrLstm>,
57    ja: Option<DataPayload<UCharDictionaryBreakDataV1>>,
58}
59
60#[derive(Debug, Clone, Copy)]
61pub(crate) struct ComplexPayloadsBorrowed<'data> {
62    grapheme: GraphemeClusterSegmenterBorrowed<'data>,
63    my: Option<DictOrLstmBorrowed<'data>>,
64    km: Option<DictOrLstmBorrowed<'data>>,
65    lo: Option<DictOrLstmBorrowed<'data>>,
66    th: Option<DictOrLstmBorrowed<'data>>,
67    ja: Option<&'data UCharDictionaryBreakData<'data>>,
68}
69
70#[cfg(feature = "lstm")]
71const MY_LSTM: &DataMarkerAttributes = DataMarkerAttributes::from_str_or_panic("Burmese_");
72#[cfg(feature = "lstm")]
73const KM_LSTM: &DataMarkerAttributes = DataMarkerAttributes::from_str_or_panic("Khmer_");
74#[cfg(feature = "lstm")]
75const LO_LSTM: &DataMarkerAttributes = DataMarkerAttributes::from_str_or_panic("Lao_");
76#[cfg(feature = "lstm")]
77const TH_LSTM: &DataMarkerAttributes = DataMarkerAttributes::from_str_or_panic("Thai_");
78
79const MY_DICT: &DataMarkerAttributes = DataMarkerAttributes::from_str_or_panic("burmesedict");
80const KM_DICT: &DataMarkerAttributes = DataMarkerAttributes::from_str_or_panic("khmerdict");
81const LO_DICT: &DataMarkerAttributes = DataMarkerAttributes::from_str_or_panic("laodict");
82const TH_DICT: &DataMarkerAttributes = DataMarkerAttributes::from_str_or_panic("thaidict");
83const CJ_DICT: &DataMarkerAttributes = DataMarkerAttributes::from_str_or_panic("cjdict");
84
85impl<'data> ComplexPayloadsBorrowed<'data> {
86    fn select(&self, language: Language) -> Option<DictOrLstmBorrowed<'data>> {
87        const ERR: DataError = DataError::custom("No segmentation model for language");
88        match language {
89            Language::Burmese => self.my.or_else(|| {
90                ERR.with_display_context("my");
91                None
92            }),
93            Language::Khmer => self.km.or_else(|| {
94                ERR.with_display_context("km");
95                None
96            }),
97            Language::Lao => self.lo.or_else(|| {
98                ERR.with_display_context("lo");
99                None
100            }),
101            Language::Thai => self.th.or_else(|| {
102                ERR.with_display_context("th");
103                None
104            }),
105            Language::ChineseOrJapanese => self.ja.map(DictOrLstmBorrowed::Dict).or_else(|| {
106                ERR.with_display_context("ja");
107                None
108            }),
109            Language::Unknown => None,
110        }
111    }
112    pub(crate) fn complex_language_segment_str(&self, input: &str) -> Vec<usize> {
113        let mut result = Vec::new();
114        let mut offset = 0;
115        for (slice, lang) in LanguageIterator::new(input) {
116            match self.select(lang) {
117                Some(DictOrLstmBorrowed::Dict(dict)) => {
118                    let seg = DictionarySegmenter::new(dict, self.grapheme);
119                    result.extend(seg.segment_str(slice).map(|n| offset + n));
120                }
121                #[cfg(feature = "lstm")]
122                Some(DictOrLstmBorrowed::Lstm(lstm)) => {
123                    let seg = LstmSegmenter::new(lstm, self.grapheme);
124                    result.extend(seg.segment_str(slice).map(|n| offset + n));
125                }
126                None => {
127                    result.push(offset + slice.len());
128                }
129            }
130            offset += slice.len();
131        }
132        result
133    }
134    /// Return UTF-16 segment offset array using dictionary or lstm segmenter.
135    pub(crate) fn complex_language_segment_utf16(&self, input: &[u16]) -> Vec<usize> {
136        let mut result = Vec::new();
137        let mut offset = 0;
138        for (slice, lang) in LanguageIteratorUtf16::new(input) {
139            match self.select(lang) {
140                Some(DictOrLstmBorrowed::Dict(dict)) => {
141                    let seg = DictionarySegmenter::new(dict, self.grapheme);
142                    result.extend(seg.segment_utf16(slice).map(|n| offset + n));
143                }
144                #[cfg(feature = "lstm")]
145                Some(DictOrLstmBorrowed::Lstm(lstm)) => {
146                    let seg = LstmSegmenter::new(lstm, self.grapheme);
147                    result.extend(seg.segment_utf16(slice).map(|n| offset + n));
148                }
149                None => {
150                    result.push(offset + slice.len());
151                }
152            }
153            offset += slice.len();
154        }
155        result
156    }
157}
158impl ComplexPayloadsBorrowed<'static> {
159    #[cfg(feature = "lstm")]
160    #[cfg(feature = "compiled_data")]
161    pub(crate) fn new_lstm() -> Self {
162        #[expect(clippy::unwrap_used)]
163        // try_load is infallible if the provider only returns `MissingLocale`.
164        Self {
165            grapheme: GraphemeClusterSegmenter::new(),
166            my: try_load_static::<SegmenterLstmAutoV1, _>(&crate::provider::Baked, MY_LSTM)
167                .unwrap()
168                .map(DictOrLstmBorrowed::Lstm),
169            km: try_load_static::<SegmenterLstmAutoV1, _>(&crate::provider::Baked, KM_LSTM)
170                .unwrap()
171                .map(DictOrLstmBorrowed::Lstm),
172            lo: try_load_static::<SegmenterLstmAutoV1, _>(&crate::provider::Baked, LO_LSTM)
173                .unwrap()
174                .map(DictOrLstmBorrowed::Lstm),
175            th: try_load_static::<SegmenterLstmAutoV1, _>(&crate::provider::Baked, TH_LSTM)
176                .unwrap()
177                .map(DictOrLstmBorrowed::Lstm),
178            ja: None,
179        }
180    }
181    #[cfg(feature = "auto")]
182    #[cfg(feature = "compiled_data")]
183    #[expect(clippy::unwrap_used)]
184    pub(crate) fn new_auto() -> Self {
185        let mut this = Self::new_lstm();
186        this.ja = try_load_static::<SegmenterDictionaryAutoV1, _>(&crate::provider::Baked, CJ_DICT)
187            .unwrap();
188        this
189    }
190    #[cfg(feature = "compiled_data")]
191    pub(crate) fn new_dict() -> Self {
192        #[expect(clippy::unwrap_used)]
193        // try_load is infallible if the provider only returns `MissingLocale`.
194        Self {
195            grapheme: GraphemeClusterSegmenter::new(),
196            my: try_load_static::<SegmenterDictionaryExtendedV1, _>(
197                &crate::provider::Baked,
198                MY_DICT,
199            )
200            .unwrap()
201            .map(DictOrLstmBorrowed::Dict),
202            km: try_load_static::<SegmenterDictionaryExtendedV1, _>(
203                &crate::provider::Baked,
204                KM_DICT,
205            )
206            .unwrap()
207            .map(DictOrLstmBorrowed::Dict),
208            lo: try_load_static::<SegmenterDictionaryExtendedV1, _>(
209                &crate::provider::Baked,
210                LO_DICT,
211            )
212            .unwrap()
213            .map(DictOrLstmBorrowed::Dict),
214            th: try_load_static::<SegmenterDictionaryExtendedV1, _>(
215                &crate::provider::Baked,
216                TH_DICT,
217            )
218            .unwrap()
219            .map(DictOrLstmBorrowed::Dict),
220            ja: try_load_static::<SegmenterDictionaryAutoV1, _>(&crate::provider::Baked, CJ_DICT)
221                .unwrap(),
222        }
223    }
224
225    #[cfg(feature = "compiled_data")]
226    pub(crate) fn new_southeast_asian() -> Self {
227        #[expect(clippy::unwrap_used)]
228        // try_load is infallible if the provider only returns `MissingLocale`.
229        Self {
230            grapheme: GraphemeClusterSegmenter::new(),
231            my: try_load_static::<SegmenterDictionaryExtendedV1, _>(
232                &crate::provider::Baked,
233                MY_DICT,
234            )
235            .unwrap()
236            .map(DictOrLstmBorrowed::Dict),
237            km: try_load_static::<SegmenterDictionaryExtendedV1, _>(
238                &crate::provider::Baked,
239                KM_DICT,
240            )
241            .unwrap()
242            .map(DictOrLstmBorrowed::Dict),
243            lo: try_load_static::<SegmenterDictionaryExtendedV1, _>(
244                &crate::provider::Baked,
245                LO_DICT,
246            )
247            .unwrap()
248            .map(DictOrLstmBorrowed::Dict),
249            th: try_load_static::<SegmenterDictionaryExtendedV1, _>(
250                &crate::provider::Baked,
251                TH_DICT,
252            )
253            .unwrap()
254            .map(DictOrLstmBorrowed::Dict),
255            ja: None,
256        }
257    }
258
259    #[cfg(feature = "compiled_data")]
260    pub(crate) const fn empty() -> Self {
261        Self {
262            grapheme: GraphemeClusterSegmenter::new(),
263            my: None,
264            km: None,
265            lo: None,
266            th: None,
267            ja: None,
268        }
269    }
270
271    pub(crate) fn static_to_owned(self) -> ComplexPayloads {
272        ComplexPayloads {
273            grapheme: self.grapheme.static_to_owned(),
274            my: self.my.map(fromstatic_dictor),
275            km: self.km.map(fromstatic_dictor),
276            lo: self.lo.map(fromstatic_dictor),
277            th: self.th.map(fromstatic_dictor),
278            ja: self.ja.map(DataPayload::from_static_ref),
279        }
280    }
281}
282
283impl ComplexPayloads {
284    pub(crate) fn as_borrowed(&self) -> ComplexPayloadsBorrowed<'_> {
285        ComplexPayloadsBorrowed {
286            grapheme: self.grapheme.as_borrowed(),
287            my: self.my.as_ref().map(borrow_dictor),
288            km: self.km.as_ref().map(borrow_dictor),
289            lo: self.lo.as_ref().map(borrow_dictor),
290            th: self.th.as_ref().map(borrow_dictor),
291            ja: self.ja.as_ref().map(|p| p.get()),
292        }
293    }
294
295    #[cfg(feature = "lstm")]
296    pub(crate) fn try_new_lstm<D>(provider: &D) -> Result<Self, DataError>
297    where
298        D: DataProvider<SegmenterBreakGraphemeClusterV1>
299            + DataProvider<SegmenterLstmAutoV1>
300            + ?Sized,
301    {
302        Ok(Self {
303            grapheme: GraphemeClusterSegmenter::try_new_unstable(provider)?,
304            my: try_load::<SegmenterLstmAutoV1, D>(provider, MY_LSTM)?
305                .map(DataPayload::cast)
306                .map(DictOrLstm::Lstm),
307            km: try_load::<SegmenterLstmAutoV1, D>(provider, KM_LSTM)?
308                .map(DataPayload::cast)
309                .map(DictOrLstm::Lstm),
310            lo: try_load::<SegmenterLstmAutoV1, D>(provider, LO_LSTM)?
311                .map(DataPayload::cast)
312                .map(DictOrLstm::Lstm),
313            th: try_load::<SegmenterLstmAutoV1, D>(provider, TH_LSTM)?
314                .map(DataPayload::cast)
315                .map(DictOrLstm::Lstm),
316            ja: None,
317        })
318    }
319
320    pub(crate) fn try_new_dict<D>(provider: &D) -> Result<Self, DataError>
321    where
322        D: DataProvider<SegmenterBreakGraphemeClusterV1>
323            + DataProvider<SegmenterDictionaryExtendedV1>
324            + DataProvider<SegmenterDictionaryAutoV1>
325            + ?Sized,
326    {
327        Ok(Self {
328            grapheme: GraphemeClusterSegmenter::try_new_unstable(provider)?,
329            my: try_load::<SegmenterDictionaryExtendedV1, D>(provider, MY_DICT)?
330                .map(DataPayload::cast)
331                .map(DictOrLstm::Dict),
332            km: try_load::<SegmenterDictionaryExtendedV1, D>(provider, KM_DICT)?
333                .map(DataPayload::cast)
334                .map(DictOrLstm::Dict),
335            lo: try_load::<SegmenterDictionaryExtendedV1, D>(provider, LO_DICT)?
336                .map(DataPayload::cast)
337                .map(DictOrLstm::Dict),
338            th: try_load::<SegmenterDictionaryExtendedV1, D>(provider, TH_DICT)?
339                .map(DataPayload::cast)
340                .map(DictOrLstm::Dict),
341            ja: try_load::<SegmenterDictionaryAutoV1, D>(provider, CJ_DICT)?.map(DataPayload::cast),
342        })
343    }
344
345    #[cfg(feature = "auto")] // Use by WordSegmenter with "auto" enabled.
346    pub(crate) fn try_new_auto<D>(provider: &D) -> Result<Self, DataError>
347    where
348        D: DataProvider<SegmenterBreakGraphemeClusterV1>
349            + DataProvider<SegmenterLstmAutoV1>
350            + DataProvider<SegmenterDictionaryAutoV1>
351            + ?Sized,
352    {
353        Ok(Self {
354            grapheme: GraphemeClusterSegmenter::try_new_unstable(provider)?,
355            my: try_load::<SegmenterLstmAutoV1, D>(provider, MY_LSTM)?
356                .map(DataPayload::cast)
357                .map(DictOrLstm::Lstm),
358            km: try_load::<SegmenterLstmAutoV1, D>(provider, KM_LSTM)?
359                .map(DataPayload::cast)
360                .map(DictOrLstm::Lstm),
361            lo: try_load::<SegmenterLstmAutoV1, D>(provider, LO_LSTM)?
362                .map(DataPayload::cast)
363                .map(DictOrLstm::Lstm),
364            th: try_load::<SegmenterLstmAutoV1, D>(provider, TH_LSTM)?
365                .map(DataPayload::cast)
366                .map(DictOrLstm::Lstm),
367            ja: try_load::<SegmenterDictionaryAutoV1, D>(provider, CJ_DICT)?.map(DataPayload::cast),
368        })
369    }
370
371    pub(crate) fn try_new_southeast_asian<D>(provider: &D) -> Result<Self, DataError>
372    where
373        D: DataProvider<SegmenterDictionaryExtendedV1>
374            + DataProvider<SegmenterBreakGraphemeClusterV1>
375            + ?Sized,
376    {
377        Ok(Self {
378            grapheme: GraphemeClusterSegmenter::try_new_unstable(provider)?,
379            my: try_load::<SegmenterDictionaryExtendedV1, _>(provider, MY_DICT)?
380                .map(DataPayload::cast)
381                .map(DictOrLstm::Dict),
382            km: try_load::<SegmenterDictionaryExtendedV1, _>(provider, KM_DICT)?
383                .map(DataPayload::cast)
384                .map(DictOrLstm::Dict),
385            lo: try_load::<SegmenterDictionaryExtendedV1, _>(provider, LO_DICT)?
386                .map(DataPayload::cast)
387                .map(DictOrLstm::Dict),
388            th: try_load::<SegmenterDictionaryExtendedV1, _>(provider, TH_DICT)?
389                .map(DataPayload::cast)
390                .map(DictOrLstm::Dict),
391            ja: None,
392        })
393    }
394
395    pub(crate) fn try_new_empty<D>(provider: &D) -> Result<Self, DataError>
396    where
397        D: DataProvider<SegmenterBreakGraphemeClusterV1> + ?Sized,
398    {
399        Ok(Self {
400            grapheme: GraphemeClusterSegmenter::try_new_unstable(provider)?,
401            my: None,
402            km: None,
403            lo: None,
404            th: None,
405            ja: None,
406        })
407    }
408}
409fn try_load<M: DataMarker, P: DataProvider<M> + ?Sized>(
410    provider: &P,
411    model: &'static DataMarkerAttributes,
412) -> Result<Option<DataPayload<M>>, DataError> {
413    provider
414        .load(DataRequest {
415            id: DataIdentifierBorrowed::for_marker_attributes(model),
416            metadata: {
417                let mut m = DataRequestMetadata::default();
418                m.silent = true;
419                m.attributes_prefix_match = true;
420                m
421            },
422        })
423        .allow_identifier_not_found()
424        .map(|r| r.map(|r| r.payload))
425}
426
427#[cfg(feature = "compiled_data")]
428fn try_load_static<M: DataMarker, P: DataProvider<M> + ?Sized>(
429    provider: &P,
430    model: &'static DataMarkerAttributes,
431) -> Result<Option<&'static <M::DataStruct as yoke::Yokeable<'static>>::Output>, DataError> {
432    provider
433        .load(DataRequest {
434            id: DataIdentifierBorrowed::for_marker_attributes(model),
435            metadata: {
436                let mut m = DataRequestMetadata::default();
437                m.silent = true;
438                m.attributes_prefix_match = true;
439                m
440            },
441        })
442        .allow_identifier_not_found()
443        .map(|r| r.and_then(|r| r.payload.get_static()))
444}
445
446#[cfg(test)]
447#[cfg(feature = "serde")]
448mod tests {
449    use super::*;
450
451    #[test]
452    fn thai_word_break() {
453        const TEST_STR: &str = "ภาษาไทยภาษาไทย";
454        let utf16: Vec<u16> = TEST_STR.encode_utf16().collect();
455
456        let lstm = ComplexPayloadsBorrowed::new_lstm();
457        let dict = ComplexPayloadsBorrowed::new_dict();
458
459        assert_eq!(
460            lstm.complex_language_segment_str(TEST_STR),
461            [12, 21, 33, 42]
462        );
463        assert_eq!(lstm.complex_language_segment_utf16(&utf16), [4, 7, 11, 14]);
464
465        assert_eq!(
466            dict.complex_language_segment_str(TEST_STR),
467            [12, 21, 33, 42]
468        );
469        assert_eq!(dict.complex_language_segment_utf16(&utf16), [4, 7, 11, 14]);
470    }
471}
icu_segmenter/complex/mod.rs

icu_segmenter/complex/
mod.rs
