1use crate::{Diagnostic, Document, Limits, Node, Position, Section, Variable};
2use std::collections::{BTreeMap, BTreeSet};
3use xmlparser::{ElementEnd, Token, Tokenizer};
4
5pub fn parse_limit(value: &str) -> Result<u64, String> {
8 let error = || format!("Invalid token limit {value:?}; use 500, 10k, or 1.5k");
9 let parse_digits = |s: &str| -> Result<u64, String> {
10 if s.is_empty() || !s.bytes().all(|b| b.is_ascii_digit()) {
11 return Err(error());
12 }
13 s.parse::<u64>().map_err(|_| error())
14 };
15 if let Some(number) = value.strip_suffix('k').or_else(|| value.strip_suffix('K')) {
16 let (whole, fraction) = number.split_once('.').unwrap_or((number, ""));
17 let mut n = parse_digits(whole)?.checked_mul(1000).ok_or_else(error)?;
18 if number.contains('.') {
19 if fraction.is_empty() || fraction.len() > 3 {
20 return Err(error());
21 }
22 n = n
23 .checked_add(parse_digits(fraction)? * 10u64.pow(3 - fraction.len() as u32))
24 .ok_or_else(error)?;
25 }
26 Ok(n)
27 } else {
28 parse_digits(value)
29 }
30}
31struct Locations<'a> {
32 source: &'a str,
33 starts: Vec<usize>,
34 wide: Vec<(usize, usize)>,
35}
36impl<'a> Locations<'a> {
37 fn new(source: &'a str) -> Self {
38 let mut extra = 0;
39 let wide = source
40 .char_indices()
41 .filter_map(|(i, c)| {
42 if c.is_ascii() {
43 None
44 } else {
45 extra += c.len_utf8() - 1;
46 Some((i + c.len_utf8(), extra))
47 }
48 })
49 .collect();
50 Self {
51 source,
52 starts: std::iter::once(0)
53 .chain(source.match_indices('\n').map(|(i, _)| i + 1))
54 .collect(),
55 wide,
56 }
57 }
58 fn extra_bytes(&self, offset: usize) -> usize {
59 let n = self.wide.partition_point(|&(end, _)| end <= offset);
60 if n == 0 { 0 } else { self.wide[n - 1].1 }
61 }
62 fn at(&self, offset: usize) -> Position {
63 let row = self.starts.partition_point(|&i| i <= offset) - 1;
64 Position {
65 line: row + 1,
66 column: offset - self.starts[row] + 1
67 - (self.extra_bytes(offset) - self.extra_bytes(self.starts[row])),
68 offset,
69 }
70 }
71 fn error(&self, row: usize, column: usize) -> Position {
72 let start = self
73 .starts
74 .get(row.saturating_sub(1))
75 .copied()
76 .unwrap_or(self.source.len());
77 let offset = self.source[start..]
78 .char_indices()
79 .nth(column.saturating_sub(1))
80 .map_or(self.source.len(), |(i, _)| start + i);
81 Position {
82 line: row,
83 column,
84 offset,
85 }
86 }
87}
88fn valid_char(c: char) -> bool {
89 matches!(c, '\t' | '\n' | '\r') || (c >= '\u{20}' && c != '\u{fffe}' && c != '\u{ffff}')
90}
91fn decode(value: &str) -> Result<String, String> {
92 let mut out = String::new();
93 let mut rest = value;
94 while let Some(index) = rest.find('&') {
95 out.push_str(&rest[..index]);
96 rest = &rest[index + 1..];
97 let end = rest
98 .find(';')
99 .ok_or("Unclosed entity; escape a literal & as &")?;
100 let entity = &rest[..end];
101 let ch = match entity {
102 "amp" => '&',
103 "lt" => '<',
104 "gt" => '>',
105 "quot" => '"',
106 "apos" => '\'',
107 _ => {
108 let n = if let Some(hex) = entity.strip_prefix("#x") {
109 (!hex.is_empty() && hex.bytes().all(|b| b.is_ascii_hexdigit()))
110 .then(|| u32::from_str_radix(hex, 16).ok())
111 .flatten()
112 } else if let Some(dec) = entity.strip_prefix('#') {
113 (!dec.is_empty() && dec.bytes().all(|b| b.is_ascii_digit()))
114 .then(|| dec.parse::<u32>().ok())
115 .flatten()
116 } else {
117 None
118 };
119 n.and_then(char::from_u32)
120 .filter(|c| valid_char(*c) && !('\u{80}'..='\u{9f}').contains(c))
121 .ok_or_else(|| format!("Unsupported entity &{entity};"))?
122 }
123 };
124 out.push(ch);
125 rest = &rest[end + 1..];
126 }
127 out.push_str(rest);
128 Ok(out.replace("\r\n", "\n").replace('\r', "\n"))
129}
130pub(crate) fn valid_id(id: &str) -> bool {
131 !id.is_empty()
132 && id
133 .bytes()
134 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b'.' | b':'))
135}
136struct Frame {
137 name: String,
138 attrs: BTreeMap<String, String>,
139 children: Vec<Node>,
140 position: Position,
141}
142fn limits(frame: &Frame) -> Result<Limits, Diagnostic> {
143 let get = |key: &str| {
144 frame
145 .attrs
146 .get(key)
147 .map(|s| parse_limit(s).map_err(|e| Diagnostic::new("limit", e, frame.position)))
148 .transpose()
149 };
150 let limits = Limits {
151 max_tokens: get("max-tokens")?,
152 per_item: get("per-item")?,
153 reason: frame.attrs.get("reason").cloned(),
154 sig: frame.attrs.get("sig").cloned(),
155 };
156 if (limits.max_tokens.is_some() || limits.per_item.is_some())
157 && limits.reason.as_ref().is_none_or(|r| r.trim().is_empty())
158 {
159 return Err(Diagnostic::new(
160 "reason",
161 "Every declared token limit requires a nonblank reason",
162 frame.position,
163 ));
164 }
165 Ok(limits)
166}
167fn append(nodes: &mut Vec<Node>, text: String) {
168 if let Some(Node::Text { value }) = nodes.last_mut() {
169 value.push_str(&text);
170 } else if !text.is_empty() {
171 nodes.push(Node::Text { value: text });
172 }
173}
174
175pub fn parse(source: &str) -> Result<Document, Diagnostic> {
179 if source.len() > 4 * 1024 * 1024 {
180 return Err(Diagnostic::new(
181 "source-size",
182 "Source exceeds 4 MiB",
183 Position::default(),
184 ));
185 }
186 let locations = Locations::new(source);
187 if let Some((i, _)) = source.char_indices().find(|(_, c)| !valid_char(*c)) {
188 return Err(Diagnostic::new(
189 "character",
190 "Invalid control character",
191 locations.at(i),
192 ));
193 }
194 let mut stack: Vec<Frame> = Vec::new();
195 let mut pending: Option<Frame> = None;
196 let mut result = None;
197 let mut ids = BTreeSet::new();
198 let mut variables = BTreeSet::new();
199 for token in Tokenizer::from(source) {
200 let token = token.map_err(|e| {
201 Diagnostic::new(
202 "syntax",
203 e.to_string(),
204 locations.error(e.pos().row as usize, e.pos().col as usize),
205 )
206 })?;
207 match token {
208 Token::ElementStart {
209 prefix,
210 local,
211 span,
212 } => {
213 let p = locations.at(span.start());
214 if !prefix.is_empty() || !matches!(local.as_str(), "htmlp" | "section") {
215 return Err(Diagnostic::new(
216 "element",
217 "Only lowercase htmlp and section are supported",
218 p,
219 ));
220 }
221 if stack.len() >= 64 {
222 return Err(Diagnostic::new("depth", "Maximum element depth is 64", p));
223 }
224 if (stack.is_empty() && (local.as_str() != "htmlp" || result.is_some()))
225 || (!stack.is_empty() && local.as_str() == "htmlp")
226 {
227 return Err(Diagnostic::new(
228 "root",
229 "Expected exactly one htmlp root",
230 p,
231 ));
232 }
233 pending = Some(Frame {
234 name: local.to_string(),
235 attrs: BTreeMap::new(),
236 children: Vec::new(),
237 position: p,
238 });
239 }
240 Token::Attribute {
241 prefix,
242 local,
243 value,
244 span,
245 } => {
246 let frame = pending.as_mut().ok_or_else(|| {
247 Diagnostic::new("syntax", "Unexpected attribute", locations.at(span.start()))
248 })?;
249 let allowed: &[&str] = match frame.name.as_str() {
250 "htmlp" => &[
251 "version",
252 "tokenizer",
253 "max-tokens",
254 "per-item",
255 "reason",
256 "sig",
257 ],
258 "section" => &["id", "max-tokens", "per-item", "reason", "sig"],
259 _ => &["id"],
260 };
261 if !prefix.is_empty() || !allowed.contains(&local.as_str()) {
262 return Err(Diagnostic::new(
263 "attribute",
264 format!("Unknown attribute {local}"),
265 locations.at(span.start()),
266 ));
267 }
268 let decoded = decode(value.as_str())
269 .map_err(|e| Diagnostic::new("entity", e, locations.at(value.start())))?;
270 if frame.attrs.insert(local.to_string(), decoded).is_some() {
271 return Err(Diagnostic::new(
272 "attribute",
273 "Duplicate attribute",
274 locations.at(span.start()),
275 ));
276 }
277 }
278 Token::ElementEnd { end, span } => {
279 if matches!(end, ElementEnd::Open | ElementEnd::Empty) {
280 let frame = pending.take().ok_or_else(|| {
281 Diagnostic::new(
282 "syntax",
283 "Unexpected opening tag",
284 locations.at(span.start()),
285 )
286 })?;
287 if let Some(id) = frame.attrs.get("id") {
288 if !valid_id(id) || variables.contains(id) || !ids.insert(id.clone()) {
289 return Err(Diagnostic::new(
290 "id",
291 "IDs must be unique and contain only ASCII letters, digits, _, -, ., or :",
292 frame.position,
293 ));
294 }
295 }
296 stack.push(frame);
297 }
298 if !matches!(end, ElementEnd::Open) {
299 let frame = stack.pop().ok_or_else(|| {
300 Diagnostic::new(
301 "syntax",
302 "Unexpected closing tag",
303 locations.at(span.start()),
304 )
305 })?;
306 if let ElementEnd::Close(prefix, local) = end {
307 if !prefix.is_empty() || frame.name != local.as_str() {
308 return Err(Diagnostic::new(
309 "syntax",
310 format!("Expected </{}>, found </{local}>", frame.name),
311 locations.at(span.start()),
312 ));
313 }
314 }
315 let budget = limits(&frame)?;
316 if frame.name == "htmlp" {
317 if frame.attrs.get("version").is_some_and(|v| v != "0.2") {
318 return Err(Diagnostic::new(
319 "version",
320 "Only version 0.2 is supported",
321 frame.position,
322 ));
323 }
324 if budget.max_tokens.is_none() {
325 return Err(Diagnostic::new(
326 "limit",
327 "htmlp requires max-tokens",
328 frame.position,
329 ));
330 }
331 let tokenizer = frame
332 .attrs
333 .get("tokenizer")
334 .cloned()
335 .unwrap_or_else(|| "cl100k_base".into());
336 if !valid_id(&tokenizer) {
337 return Err(Diagnostic::new(
338 "tokenizer",
339 "Invalid tokenizer identifier",
340 frame.position,
341 ));
342 }
343 result = Some(Document {
344 version: "0.2".into(),
345 tokenizer,
346 limits: budget,
347 children: frame.children,
348 position: frame.position,
349 });
350 } else {
351 let node = Node::Section(Section {
352 id: frame.attrs.get("id").cloned(),
353 limits: budget,
354 children: frame.children,
355 position: frame.position,
356 });
357 stack
358 .last_mut()
359 .ok_or_else(|| Diagnostic::new("root", "Missing root", frame.position))?
360 .children
361 .push(node);
362 }
363 }
364 }
365 Token::Text { text } => {
366 if let Some(frame) = stack.last_mut() {
367 let raw = text.as_str();
368 let mut cursor = 0;
369 while let Some(relative) = raw[cursor..].find("{{") {
370 let start = cursor + relative;
371 let p = locations.at(text.start() + start);
372 let end = raw[start + 2..]
373 .find("}}")
374 .map(|i| start + 2 + i)
375 .ok_or_else(|| {
376 Diagnostic::new("variable", "Unclosed {{name}} placeholder", p)
377 })?;
378 let id = &raw[start + 2..end];
379 if !valid_id(id) || ids.contains(id) {
380 return Err(Diagnostic::new(
381 "variable",
382 "Variable names must be valid IDs and cannot match a section ID",
383 p,
384 ));
385 }
386 let decoded = decode(&raw[cursor..start]).map_err(|e| {
387 Diagnostic::new("entity", e, locations.at(text.start() + cursor))
388 })?;
389 append(&mut frame.children, decoded);
390 variables.insert(id.to_string());
391 frame.children.push(Node::Variable(Variable {
392 id: id.to_string(),
393 position: p,
394 }));
395 cursor = end + 2;
396 }
397 let decoded = decode(&raw[cursor..]).map_err(|e| {
398 Diagnostic::new("entity", e, locations.at(text.start() + cursor))
399 })?;
400 append(&mut frame.children, decoded);
401 } else if !text.as_str().trim().is_empty() {
402 return Err(Diagnostic::new(
403 "root",
404 "Text must be inside htmlp",
405 locations.at(text.start()),
406 ));
407 }
408 }
409 Token::Comment { .. } => {}
410 _ => {
411 return Err(Diagnostic::new(
412 "syntax",
413 "DOCTYPE, declarations, CDATA, and processing instructions are not supported",
414 Position::default(),
415 ));
416 }
417 }
418 }
419 if !stack.is_empty() {
420 return Err(Diagnostic::new(
421 "syntax",
422 "Unclosed element",
423 stack.last().unwrap().position,
424 ));
425 }
426 result.ok_or_else(|| Diagnostic::new("root", "Missing htmlp root", Position::default()))
427}