1use std::collections::HashMap;
2
3use crate::fingerprint::WeightedFeature;
4
5#[derive(Debug, Clone)]
6pub struct WebFeatureOptions {
7 pub include_bigrams: bool,
8 pub include_length_bucket: bool,
9 pub drop_stop_words: bool,
10 pub drop_numeric_tokens: bool,
11 pub stem_tokens: bool,
12}
13
14impl Default for WebFeatureOptions {
15 fn default() -> Self {
16 Self {
17 include_bigrams: true,
18 include_length_bucket: true,
19 drop_stop_words: true,
20 drop_numeric_tokens: true,
21 stem_tokens: true,
22 }
23 }
24}
25
26#[derive(Debug, Clone, Default)]
27pub struct WebFeatureExtractor {
28 options: WebFeatureOptions,
29}
30
31impl WebFeatureExtractor {
32 pub fn new(options: WebFeatureOptions) -> Self {
33 Self { options }
34 }
35
36 pub fn options(&self) -> &WebFeatureOptions {
37 &self.options
38 }
39
40 pub fn text_from_html(&self, html: &str) -> String {
41 decode_html_entities(&strip_html(html))
42 }
43
44 pub fn extract_from_html(&self, html: &str) -> Vec<WeightedFeature> {
45 let text = self.text_from_html(html);
46 self.extract_from_text(&text)
47 }
48
49 pub fn extract_from_text(&self, text: &str) -> Vec<WeightedFeature> {
50 let tokens = self.tokenize(text);
51 let mut weights: HashMap<String, u32> = HashMap::new();
52
53 for token in &tokens {
54 *weights.entry(format!("term:{token}")).or_insert(0) += 3;
55 }
56
57 if self.options.include_bigrams {
58 for pair in tokens.windows(2) {
59 *weights
60 .entry(format!("phrase:{} {}", pair[0], pair[1]))
61 .or_insert(0) += 2;
62 }
63 }
64
65 if self.options.include_length_bucket {
66 *weights
67 .entry(format!("len:{}", length_bucket(tokens.len())))
68 .or_insert(0) += 4;
69 }
70
71 let mut features: Vec<_> = weights
72 .into_iter()
73 .map(|(key, weight)| WeightedFeature::new(key, weight))
74 .collect();
75 features.sort_by(|a, b| a.key.cmp(&b.key));
76 features
77 }
78
79 pub fn tokenize(&self, text: &str) -> Vec<String> {
80 let mut tokens = Vec::new();
81 let mut current = String::new();
82
83 for ch in text.chars() {
84 if ch.is_alphanumeric() {
85 current.extend(ch.to_lowercase());
86 } else {
87 self.push_token(&mut current, &mut tokens);
88 }
89 }
90 self.push_token(&mut current, &mut tokens);
91
92 tokens
93 }
94
95 fn push_token(&self, current: &mut String, tokens: &mut Vec<String>) {
96 if current.is_empty() {
97 return;
98 }
99
100 let raw = std::mem::take(current);
101 if self.options.drop_numeric_tokens && raw.chars().all(|ch| ch.is_ascii_digit()) {
102 return;
103 }
104 if self.options.drop_stop_words && is_stop_word(&raw) {
105 return;
106 }
107
108 let token = if self.options.stem_tokens {
109 stem(&raw)
110 } else {
111 raw
112 };
113
114 if token.len() >= 2 {
115 tokens.push(token);
116 }
117 }
118}
119
120fn strip_html(html: &str) -> String {
121 let mut out = String::with_capacity(html.len());
122 let mut pos = 0usize;
123 let lower: Vec<u8> = html.as_bytes().iter().map(u8::to_ascii_lowercase).collect();
124 let bytes = html.as_bytes();
125
126 while pos < html.len() {
127 if lower[pos..].starts_with(b"<!--") {
128 if let Some(end) = find_bytes(&lower[pos + 4..], b"-->") {
129 pos += 4 + end + 3;
130 } else {
131 break;
132 }
133 out.push(' ');
134 continue;
135 }
136
137 if bytes[pos] == b'<' {
138 let Some(end_rel) = html[pos..].find('>') else {
139 break;
140 };
141 let end = pos + end_rel;
142 let inside = String::from_utf8_lossy(&lower[pos + 1..end]);
143 let tag = tag_name(&inside);
144 let closing = inside.trim_start().starts_with('/');
145
146 if !closing && should_skip_tag(&tag) {
147 let close_pat = format!("</{tag}");
148 if let Some(close_start_rel) = find_bytes(&lower[end + 1..], close_pat.as_bytes()) {
149 let close_start = end + 1 + close_start_rel;
150 if let Some(close_end_rel) = find_bytes(&lower[close_start..], b">") {
151 pos = close_start + close_end_rel + 1;
152 } else {
153 break;
154 }
155 } else {
156 pos = end + 1;
157 }
158 out.push(' ');
159 continue;
160 }
161
162 if is_boundary_tag(&tag) || inside.trim_end().ends_with('/') {
163 out.push(' ');
164 }
165 pos = end + 1;
166 continue;
167 }
168
169 let ch = html[pos..].chars().next().expect("valid char boundary");
170 out.push(ch);
171 pos += ch.len_utf8();
172 }
173
174 out
175}
176
177fn tag_name(inside: impl AsRef<str>) -> String {
178 let inside = inside.as_ref();
179 let trimmed = inside.trim_start().trim_start_matches('/');
180 trimmed
181 .chars()
182 .take_while(|ch| ch.is_ascii_alphanumeric())
183 .collect()
184}
185
186fn should_skip_tag(tag: &str) -> bool {
187 matches!(
188 tag,
189 "script"
190 | "style"
191 | "noscript"
192 | "template"
193 | "svg"
194 | "head"
195 | "nav"
196 | "footer"
197 | "aside"
198 | "form"
199 | "iframe"
200 )
201}
202
203fn is_boundary_tag(tag: &str) -> bool {
204 matches!(
205 tag,
206 "address"
207 | "article"
208 | "br"
209 | "div"
210 | "h1"
211 | "h2"
212 | "h3"
213 | "h4"
214 | "h5"
215 | "h6"
216 | "header"
217 | "li"
218 | "main"
219 | "p"
220 | "section"
221 | "td"
222 | "th"
223 | "tr"
224 | "ul"
225 | "ol"
226 )
227}
228
229fn decode_html_entities(text: &str) -> String {
230 let mut out = String::with_capacity(text.len());
231 let mut chars = text.chars().peekable();
232
233 while let Some(ch) = chars.next() {
234 if ch != '&' {
235 out.push(ch);
236 continue;
237 }
238
239 let mut entity = String::new();
240 while let Some(next) = chars.peek().copied() {
241 chars.next();
242 if next == ';' {
243 break;
244 }
245 if entity.len() > 16 {
246 break;
247 }
248 entity.push(next);
249 }
250
251 match decode_entity(&entity) {
252 Some(decoded) => out.push(decoded),
253 None => {
254 out.push('&');
255 out.push_str(&entity);
256 out.push(';');
257 }
258 }
259 }
260
261 out
262}
263
264fn decode_entity(entity: &str) -> Option<char> {
265 match entity {
266 "amp" => Some('&'),
267 "lt" => Some('<'),
268 "gt" => Some('>'),
269 "quot" => Some('"'),
270 "apos" => Some('\''),
271 "nbsp" => Some(' '),
272 _ if entity.starts_with("#x") || entity.starts_with("#X") => {
273 u32::from_str_radix(&entity[2..], 16)
274 .ok()
275 .and_then(char::from_u32)
276 }
277 _ if entity.starts_with('#') => entity[1..].parse::<u32>().ok().and_then(char::from_u32),
278 _ => None,
279 }
280}
281
282fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option<usize> {
283 if needle.is_empty() {
284 return Some(0);
285 }
286 haystack
287 .windows(needle.len())
288 .position(|window| window == needle)
289}
290
291fn is_stop_word(token: &str) -> bool {
292 matches!(
293 token,
294 "a" | "an"
295 | "and"
296 | "are"
297 | "as"
298 | "at"
299 | "be"
300 | "by"
301 | "for"
302 | "from"
303 | "has"
304 | "have"
305 | "in"
306 | "is"
307 | "it"
308 | "its"
309 | "of"
310 | "on"
311 | "or"
312 | "that"
313 | "the"
314 | "this"
315 | "to"
316 | "was"
317 | "were"
318 | "will"
319 | "with"
320 )
321}
322
323fn stem(token: &str) -> String {
324 let len = token.len();
325 if len > 6 && token.ends_with("ingly") {
326 return token[..len - 5].to_string();
327 }
328 if len > 5 && token.ends_with("ing") {
329 return token[..len - 3].to_string();
330 }
331 if len > 5 && token.ends_with("edly") {
332 return token[..len - 4].to_string();
333 }
334 if len > 4 && token.ends_with("ed") {
335 return token[..len - 2].to_string();
336 }
337 if len > 4 && token.ends_with("ies") {
338 return format!("{}y", &token[..len - 3]);
339 }
340 if len > 5
341 && (token.ends_with("sses")
342 || token.ends_with("ches")
343 || token.ends_with("shes")
344 || token.ends_with("xes")
345 || token.ends_with("zes"))
346 {
347 return token[..len - 2].to_string();
348 }
349 if len > 4 && token.ends_with('s') {
350 return token[..len - 1].to_string();
351 }
352 token.to_string()
353}
354
355fn length_bucket(len: usize) -> &'static str {
356 match len {
357 0 => "empty",
358 1..=15 => "tiny",
359 16..=63 => "short",
360 64..=255 => "medium",
361 256..=1023 => "long",
362 _ => "huge",
363 }
364}
365
366#[cfg(test)]
367mod tests {
368 use super::*;
369
370 #[test]
371 fn strips_html_boilerplate_and_decodes_entities() {
372 let extractor = WebFeatureExtractor::default();
373 let text = extractor.text_from_html(
374 r#"
375 <html>
376 <head><title>hidden title</title></head>
377 <body>
378 <!-- hidden comment -->
379 <nav>hidden navigation</nav>
380 <main><h1>Visible & useful</h1><p>Text here.</p></main>
381 <script>hiddenCode()</script>
382 <style>.hidden { color: red }</style>
383 </body>
384 </html>
385 "#,
386 );
387
388 assert!(text.contains("Visible & useful"));
389 assert!(text.contains("Text here"));
390 assert!(!text.contains("hidden title"));
391 assert!(!text.contains("hidden navigation"));
392 assert!(!text.contains("hiddenCode"));
393 }
394
395 #[test]
396 fn strips_html_with_unicode_text_before_tags() {
397 let extractor = WebFeatureExtractor::default();
398 let text = extractor.text_from_html("İstanbul <script>hidden</script><main>café</main>");
399
400 assert!(text.contains("İstanbul"));
401 assert!(text.contains("café"));
402 assert!(!text.contains("hidden"));
403 }
404
405 #[test]
406 fn tokenizes_normalizes_and_filters_web_text() {
407 let extractor = WebFeatureExtractor::default();
408 let tokens =
409 extractor.tokenize("The crawlers were RUNNING on pages 123 and indexed pages.");
410
411 assert!(!tokens.contains(&"the".to_string()));
412 assert!(!tokens.contains(&"123".to_string()));
413 assert!(tokens.contains(&"crawler".to_string()));
414 assert!(tokens.contains(&"runn".to_string()));
415 assert!(tokens.contains(&"page".to_string()));
416 }
417
418 #[test]
419 fn extraction_adds_terms_phrases_and_length_bucket() {
420 let extractor = WebFeatureExtractor::default();
421 let features = extractor.extract_from_text("near duplicate web documents");
422 let keys: Vec<_> = features
423 .iter()
424 .map(|feature| feature.key.as_str())
425 .collect();
426
427 assert!(keys.contains(&"term:near"));
428 assert!(keys.contains(&"term:duplicate"));
429 assert!(keys.contains(&"phrase:near duplicate"));
430 assert!(keys.iter().any(|key| key.starts_with("len:")));
431 }
432
433 #[test]
434 fn repeated_terms_increase_weight() {
435 let extractor = WebFeatureExtractor::default();
436 let features = extractor.extract_from_text("crawler crawler crawler");
437 let crawler = features
438 .iter()
439 .find(|feature| feature.key == "term:crawler")
440 .expect("crawler feature");
441
442 assert_eq!(crawler.weight, 9);
443 }
444}