1use std::cmp::Ordering;
8
9use rupg_common::SqlState;
10
11use crate::error::TypeError;
12use crate::json::{self, Kind, Open, Sink, Stop};
13use crate::jsonb::escape;
14use crate::numeric::numeric_in;
15use crate::tsquery::too_deep;
16
17const MAX_DEPTH: usize = 10_000;
19
20#[derive(Debug, Clone, PartialEq, Eq)]
22pub enum JsonItem {
23 Null,
24 Bool(bool),
25 Number,
27 String(String),
29 Array(Vec<usize>),
31 Object(Vec<(String, usize)>),
33}
34
35#[derive(Debug, Clone)]
37pub struct JsonNode {
38 pub item: JsonItem,
39 pub start: usize,
40 pub end: usize,
41}
42
43#[derive(Debug)]
45pub struct JsonTree<'a> {
46 text: &'a str,
47 nodes: Vec<JsonNode>,
48}
49
50#[derive(Default)]
52struct Builder {
53 nodes: Vec<JsonNode>,
54 open: Vec<usize>,
56 key: Option<String>,
58}
59
60impl Builder {
61 fn push(&mut self, item: JsonItem, start: usize, end: usize) -> usize {
63 let at = self.nodes.len();
64 self.nodes.push(JsonNode { item, start, end });
65 if let Some(&parent) = self.open.last() {
66 match &mut self.nodes[parent].item {
67 JsonItem::Array(items) => items.push(at),
68 JsonItem::Object(fields) => fields.push((self.key.take().unwrap_or_default(), at)),
69 _ => {}
70 }
71 }
72 at
73 }
74}
75
76impl Sink for Builder {
77 fn open(&mut self, open: Open) -> Result<(), TypeError> {
78 self.open_at(open, 0)
79 }
80
81 fn close(&mut self, open: Open) -> Result<(), TypeError> {
82 self.close_at(open, 0)
83 }
84
85 fn key(&mut self, key: &str) -> Result<(), TypeError> {
86 self.key = Some(key.to_owned());
87 Ok(())
88 }
89
90 fn scalar(&mut self, kind: Kind, token: &str, text: &str) -> Result<(), TypeError> {
91 self.scalar_at(kind, 0, token, text)
92 }
93
94 fn open_at(&mut self, open: Open, start: usize) -> Result<(), TypeError> {
95 let item = match open {
96 Open::Array => JsonItem::Array(Vec::new()),
97 Open::Object => JsonItem::Object(Vec::new()),
98 };
99 let at = self.push(item, start, start);
100 self.open.push(at);
101 Ok(())
102 }
103
104 fn close_at(&mut self, _: Open, end: usize) -> Result<(), TypeError> {
105 if let Some(at) = self.open.pop() {
106 self.nodes[at].end = end;
107 }
108 Ok(())
109 }
110
111 fn scalar_at(
112 &mut self,
113 kind: Kind,
114 start: usize,
115 token: &str,
116 text: &str,
117 ) -> Result<(), TypeError> {
118 let item = match kind {
119 Kind::String => JsonItem::String(text.to_owned()),
120 Kind::Number => JsonItem::Number,
121 Kind::True => JsonItem::Bool(true),
122 Kind::False => JsonItem::Bool(false),
123 _ => JsonItem::Null,
124 };
125 self.push(item, start, start + token.len());
126 Ok(())
127 }
128}
129
130impl<'a> JsonTree<'a> {
131 pub fn parse(text: &'a str) -> Result<JsonTree<'a>, TypeError> {
137 JsonTree::parse_with(text, true)
138 }
139
140 pub fn parse_with(text: &'a str, escapes: bool) -> Result<JsonTree<'a>, TypeError> {
146 let mut builder = Builder::default();
147 match json::parse(text, escapes, &mut builder) {
148 Ok(()) => Ok(JsonTree { text, nodes: builder.nodes }),
149 Err(Stop::Syntax(fail)) => Err(json::syntax(text, fail)),
150 Err(Stop::Action(error)) => Err(error),
151 }
152 }
153
154 pub fn source(&self) -> &'a str {
156 self.text
157 }
158
159 pub fn into_nodes(self) -> Vec<JsonNode> {
161 self.nodes
162 }
163
164 pub fn node(&self, at: usize) -> &JsonNode {
165 &self.nodes[at]
166 }
167
168 pub fn item(&self, at: usize) -> &JsonItem {
169 &self.nodes[at].item
170 }
171
172 pub fn text(&self, at: usize) -> &'a str {
174 let node = &self.nodes[at];
175 &self.text[node.start..node.end]
176 }
177
178 pub fn field(&self, at: usize, key: &str) -> Option<usize> {
180 match self.item(at) {
181 JsonItem::Object(fields) => fields.iter().rev().find(|(k, _)| k == key).map(|f| f.1),
182 _ => None,
183 }
184 }
185
186 pub fn element(&self, at: usize, index: i64) -> Option<usize> {
188 let JsonItem::Array(items) = self.item(at) else { return None };
189 let index = if index < 0 { index + items.len() as i64 } else { index };
190 usize::try_from(index).ok().and_then(|i| items.get(i).copied())
191 }
192
193 pub fn as_text(&self, at: usize) -> Option<String> {
195 match self.item(at) {
196 JsonItem::Null => None,
197 JsonItem::String(s) => Some(s.clone()),
198 _ => Some(self.text(at).to_owned()),
199 }
200 }
201
202 pub fn type_name(&self, at: usize) -> &'static str {
204 match self.item(at) {
205 JsonItem::Null => "null",
206 JsonItem::Bool(_) => "boolean",
207 JsonItem::Number => "number",
208 JsonItem::String(_) => "string",
209 JsonItem::Array(_) => "array",
210 JsonItem::Object(_) => "object",
211 }
212 }
213
214 pub fn tokens(&self) -> Tokens<'_, 'a> {
216 Tokens { tree: self, stack: Vec::new(), start: true, value: None }
217 }
218
219 pub fn strip_nulls(&self, arrays: bool, spaces: bool) -> String {
221 let (comma, colon) = if spaces { (", ", ": ") } else { (",", ":") };
222 let mut out = String::with_capacity(self.text.len());
223 let mut first = true;
225 let mut skip = false;
227 for token in self.tokens() {
228 match token {
229 Token::Key(name, value) => {
230 if self.item(value) == &JsonItem::Null {
231 skip = true;
232 continue;
233 }
234 if !first {
235 out.push_str(comma);
236 }
237 escape(name, &mut out);
238 out.push_str(colon);
239 first = true;
240 }
241 Token::Scalar { node, elem } => {
242 if std::mem::take(&mut skip) {
243 continue;
244 }
245 if elem {
246 if arrays && self.item(node) == &JsonItem::Null {
247 continue;
248 }
249 if !first {
250 out.push_str(comma);
251 }
252 }
253 self.put_scalar(node, &mut out);
254 first = false;
255 }
256 Token::Begin { node, raw, elem } => {
257 if raw {
258 continue;
259 }
260 if elem && !first {
261 out.push_str(comma);
262 }
263 out.push(if self.is_array(node) { '[' } else { '{' });
264 first = true;
265 }
266 Token::End { node, raw } => {
267 if raw {
268 continue;
269 }
270 out.push(if self.is_array(node) { ']' } else { '}' });
271 first = false;
272 }
273 }
274 }
275 out
276 }
277
278 pub fn pretty(&self) -> String {
280 fn indent(out: &mut String, on: bool, level: usize) {
281 if on {
282 out.push('\n');
283 out.extend(std::iter::repeat_n(' ', level * 4));
284 }
285 }
286 let mut out = String::with_capacity(self.text.len() * 2);
287 let (mut first, mut use_indent, mut last_was_key) = (true, false, false);
288 let mut level = 0;
289 for token in self.tokens() {
290 let mut key = false;
291 match token {
292 Token::Begin { node, raw, .. } => {
293 if !first {
294 out.push(',');
295 }
296 if !raw {
297 indent(&mut out, use_indent && !last_was_key, level);
298 out.push(if self.is_array(node) { '[' } else { '{' });
299 }
300 first = true;
301 level += 1;
302 }
303 Token::Key(name, value) => {
304 if !first {
305 out.push(',');
306 }
307 indent(&mut out, use_indent, level);
308 escape(name, &mut out);
309 out.push_str(": ");
310 first = true;
311 key = !self.is_scalar(value);
312 }
313 Token::Scalar { node, elem } => {
314 if elem {
315 if !first {
316 out.push(',');
317 }
318 indent(&mut out, use_indent, level);
319 }
320 self.put_scalar(node, &mut out);
321 first = false;
322 }
323 Token::End { node, raw } => {
324 level -= 1;
325 if !raw {
326 indent(&mut out, use_indent, level);
327 out.push(if self.is_array(node) { ']' } else { '}' });
328 }
329 first = false;
330 }
331 }
332 use_indent = true;
333 last_was_key = key;
334 }
335 out
336 }
337
338 pub fn compare(&self, other: &JsonTree<'_>) -> Ordering {
340 let (mut a, mut b) = (self.tokens(), other.tokens());
341 loop {
342 let (ta, tb) = match (a.next(), b.next()) {
343 (Some(ta), Some(tb)) => (ta, tb),
344 (None, None) => return Ordering::Equal,
345 (None, Some(_)) => return Ordering::Less,
346 (Some(_), None) => return Ordering::Greater,
347 };
348 let (ra, rb) = (self.token_rank(ta), other.token_rank(tb));
349 let order = match (ta, tb) {
350 _ if ra != rb => ra.cmp(&rb),
351 (Token::Key(x, _), Token::Key(y, _)) => x.as_bytes().cmp(y.as_bytes()),
352 (Token::Begin { node: x, raw: rx, .. }, Token::Begin { node: y, raw: ry, .. }) => {
353 let mut order = Ordering::Equal;
354 if rx != ry {
355 order = if rx { Ordering::Less } else { Ordering::Greater };
356 }
357 let (nx, ny) = (self.count(x, rx), other.count(y, ry));
359 if nx != ny {
360 order = nx.cmp(&ny);
361 }
362 order
363 }
364 (Token::Scalar { node: x, .. }, Token::Scalar { node: y, .. }) => {
365 self.compare_scalar(x, other, y)
366 }
367 _ => Ordering::Equal,
368 };
369 if order != Ordering::Equal {
370 return order;
371 }
372 }
373 }
374
375 pub fn contains(&self, other: &JsonTree<'_>) -> Result<bool, TypeError> {
381 if self.is_scalar(0) && !other.is_scalar(0) {
382 return Ok(false);
383 }
384 self.contains_at(0, other, 0, 0)
385 }
386
387 pub fn exists(&self, key: &str) -> bool {
389 match self.item(0) {
390 JsonItem::Object(fields) => fields.iter().any(|(k, _)| k == key),
391 _ => self
392 .members(0)
393 .iter()
394 .any(|&x| matches!(self.item(x), JsonItem::String(s) if s == key)),
395 }
396 }
397
398 pub fn set_path(
404 &self,
405 path: &[Option<&str>],
406 new: &str,
407 op: PathOp,
408 ) -> Result<String, TypeError> {
409 let mut out = String::with_capacity(self.text.len() + new.len());
410 self.set_at(0, path, 0, new, op, &mut out)?;
411 Ok(out)
412 }
413
414 fn set_at(
416 &self,
417 at: usize,
418 path: &[Option<&str>],
419 level: usize,
420 new: &str,
421 op: PathOp,
422 out: &mut String,
423 ) -> Result<(), TypeError> {
424 if level > MAX_DEPTH {
425 return Err(too_deep());
426 }
427 let Some(step) = path[level] else {
428 return Err(TypeError::new(
429 SqlState::NULL_VALUE_NOT_ALLOWED,
430 format!("path element at position {} is null", level + 1),
431 ));
432 };
433 match self.item(at) {
434 JsonItem::Object(fields) => self.set_object(fields, step, path, level, new, op, out),
435 JsonItem::Array(items) => self.set_array(items, step, path, level, new, op, out),
436 _ => {
437 out.push_str(self.text(at));
438 Ok(())
439 }
440 }
441 }
442
443 #[allow(clippy::too_many_arguments)]
445 fn set_object(
446 &self,
447 fields: &[(String, usize)],
448 step: &str,
449 path: &[Option<&str>],
450 level: usize,
451 new: &str,
452 op: PathOp,
453 out: &mut String,
454 ) -> Result<(), TypeError> {
455 let last = level + 1 == path.len();
456 let mut first = true;
457 out.push('{');
458 if fields.is_empty() && op.creates() && last {
459 put_field(out, &mut first, step, new);
460 }
461 let mut done = false;
462 for (i, (name, value)) in fields.iter().enumerate() {
463 if !done && name == step {
464 done = true;
465 if !last {
466 put_key(out, &mut first, name);
467 self.set_at(*value, path, level + 1, new, op, out)?;
468 } else if op.inserts() {
469 let mut error = TypeError::new(
470 SqlState::INVALID_PARAMETER_VALUE,
471 "cannot replace existing key".to_string(),
472 );
473 error.hint =
474 Some("Try using the function jsonb_set to replace key value.".to_string());
475 return Err(error);
476 } else if op != PathOp::Delete {
477 put_field(out, &mut first, name, new);
478 }
479 continue;
480 }
481 if op.creates() && !done && last && i + 1 == fields.len() {
482 put_field(out, &mut first, step, new);
483 }
484 put_field(out, &mut first, name, self.text(*value));
485 }
486 out.push('}');
487 Ok(())
488 }
489
490 #[allow(clippy::too_many_arguments)]
492 fn set_array(
493 &self,
494 items: &[usize],
495 step: &str,
496 path: &[Option<&str>],
497 level: usize,
498 new: &str,
499 op: PathOp,
500 out: &mut String,
501 ) -> Result<(), TypeError> {
502 let Some(index) = json_index(step) else {
503 return Err(TypeError::new(
504 SqlState::INVALID_TEXT_REPRESENTATION,
505 format!("path element at position {} is not an integer: \"{step}\"", level + 1),
506 ));
507 };
508 let count = items.len() as i64;
509 let mut index = i64::from(index);
510 let mut front = false;
512 if index < 0 {
513 if -index > count {
514 front = true;
515 } else {
516 index += count;
517 }
518 }
519 index = index.min(count);
520 let last = level + 1 == path.len();
521 let mut first = true;
522 let mut done = false;
523 out.push('[');
524 if (front || items.is_empty()) && last && op.creates() {
525 put_element(out, &mut first, new);
526 done = true;
527 }
528 for (i, &item) in items.iter().enumerate() {
529 if front || i as i64 != index {
530 put_element(out, &mut first, self.text(item));
531 continue;
532 }
533 done = true;
534 if !last {
535 put_element(out, &mut first, "");
536 self.set_at(item, path, level + 1, new, op, out)?;
537 continue;
538 }
539 if matches!(op, PathOp::InsertBefore | PathOp::Create) {
540 put_element(out, &mut first, new);
541 }
542 if op.inserts() {
543 put_element(out, &mut first, self.text(item));
544 }
545 if matches!(op, PathOp::InsertAfter | PathOp::Replace) {
546 put_element(out, &mut first, new);
547 }
548 }
549 if op.creates() && !done && last {
550 put_element(out, &mut first, new);
551 }
552 out.push(']');
553 Ok(())
554 }
555
556 fn members(&self, at: usize) -> &[usize] {
558 match self.item(at) {
559 JsonItem::Array(items) => items,
560 _ if at == 0 && self.is_scalar(0) => ROOT,
561 _ => &[],
562 }
563 }
564
565 fn contains_at(
567 &self,
568 at: usize,
569 other: &JsonTree<'_>,
570 other_at: usize,
571 depth: usize,
572 ) -> Result<bool, TypeError> {
573 if depth > MAX_DEPTH {
574 return Err(too_deep());
575 }
576 let object = matches!(self.item(at), JsonItem::Object(_));
577 if object != matches!(other.item(other_at), JsonItem::Object(_)) {
578 return Ok(false);
579 }
580 if let (JsonItem::Object(mine), JsonItem::Object(theirs)) =
581 (self.item(at), other.item(other_at))
582 {
583 if mine.len() < theirs.len() {
584 return Ok(false);
585 }
586 for (name, want) in theirs {
587 let Some(have) = self.field(at, name) else { return Ok(false) };
588 if self.is_scalar(have) || other.is_scalar(*want) {
589 if !self.scalar_equal(have, other, *want) {
590 return Ok(false);
591 }
592 } else if !self.contains_at(have, other, *want, depth + 1)? {
593 return Ok(false);
594 }
595 }
596 return Ok(true);
597 }
598 let mine = self.members(at);
599 for &want in other.members(other_at) {
600 if other.is_scalar(want) {
601 if !mine.iter().any(|&have| self.scalar_equal(have, other, want)) {
602 return Ok(false);
603 }
604 continue;
605 }
606 let mut found = false;
607 for &have in mine.iter().filter(|&&have| !self.is_scalar(have)) {
608 if self.contains_at(have, other, want, depth + 1)? {
609 found = true;
610 break;
611 }
612 }
613 if !found {
614 return Ok(false);
615 }
616 }
617 Ok(true)
618 }
619
620 fn is_scalar(&self, at: usize) -> bool {
621 !matches!(self.item(at), JsonItem::Array(_) | JsonItem::Object(_))
622 }
623
624 fn is_array(&self, at: usize) -> bool {
625 matches!(self.item(at), JsonItem::Array(_))
626 }
627
628 fn count(&self, at: usize, raw: bool) -> usize {
630 match self.item(at) {
631 _ if raw => 1,
632 JsonItem::Array(items) => items.len(),
633 JsonItem::Object(fields) => fields.len(),
634 _ => 0,
635 }
636 }
637
638 fn rank(&self, at: usize) -> u8 {
640 match self.item(at) {
641 JsonItem::Null => 0,
642 JsonItem::String(_) => 1,
643 JsonItem::Number => 2,
644 JsonItem::Bool(_) => 3,
645 JsonItem::Array(_) => 16,
646 JsonItem::Object(_) => 17,
647 }
648 }
649
650 fn token_rank(&self, token: Token<'_>) -> u8 {
652 match token {
653 Token::Key(..) => 1,
654 Token::Begin { raw: true, .. } => 16,
655 Token::Begin { node, .. } | Token::Scalar { node, .. } | Token::End { node, .. } => {
656 self.rank(node)
657 }
658 }
659 }
660
661 fn compare_scalar(&self, at: usize, other: &JsonTree<'_>, other_at: usize) -> Ordering {
663 match (self.item(at), other.item(other_at)) {
664 (JsonItem::String(x), JsonItem::String(y)) => x.as_bytes().cmp(y.as_bytes()),
665 (JsonItem::Bool(x), JsonItem::Bool(y)) => x.cmp(y),
666 (JsonItem::Number, JsonItem::Number) => {
667 match (numeric_in(self.text(at), -1), numeric_in(other.text(other_at), -1)) {
668 (Ok(x), Ok(y)) => x.compare(&y),
669 _ => self.text(at).cmp(other.text(other_at)),
670 }
671 }
672 _ => Ordering::Equal,
673 }
674 }
675
676 fn scalar_equal(&self, at: usize, other: &JsonTree<'_>, other_at: usize) -> bool {
678 self.is_scalar(at)
679 && other.is_scalar(other_at)
680 && self.rank(at) == other.rank(other_at)
681 && self.compare_scalar(at, other, other_at) == Ordering::Equal
682 }
683
684 fn put_scalar(&self, at: usize, out: &mut String) {
686 match self.item(at) {
687 JsonItem::String(s) => escape(s, out),
688 _ => out.push_str(self.text(at)),
689 }
690 }
691}
692
693const ROOT: &[usize] = &[0];
695
696#[derive(Debug, Clone, Copy, PartialEq, Eq)]
698pub enum PathOp {
699 Create,
701 Replace,
703 Delete,
705 InsertBefore,
707 InsertAfter,
709}
710
711impl PathOp {
712 fn creates(self) -> bool {
714 matches!(self, PathOp::Create | PathOp::InsertBefore | PathOp::InsertAfter)
715 }
716
717 fn inserts(self) -> bool {
718 matches!(self, PathOp::InsertBefore | PathOp::InsertAfter)
719 }
720}
721
722pub fn json_index(s: &str) -> Option<i32> {
724 let s = s.trim_start_matches([' ', '\t', '\n', '\u{b}', '\u{c}', '\r']);
725 let digits = s.strip_prefix(['+', '-']).unwrap_or(s);
726 if digits.is_empty() || !digits.bytes().all(|b| b.is_ascii_digit()) {
727 return None;
728 }
729 s.parse().ok()
730}
731
732fn put_comma(out: &mut String, first: &mut bool) {
734 if !std::mem::take(first) {
735 out.push_str(", ");
736 }
737}
738
739fn put_element(out: &mut String, first: &mut bool, text: &str) {
741 put_comma(out, first);
742 out.push_str(text);
743}
744
745fn put_key(out: &mut String, first: &mut bool, name: &str) {
747 put_comma(out, first);
748 escape(name, out);
749 out.push_str(": ");
750}
751
752fn put_field(out: &mut String, first: &mut bool, name: &str, value: &str) {
754 put_key(out, first, name);
755 out.push_str(value);
756}
757
758#[derive(Debug, Clone, Copy)]
760pub enum Token<'t> {
761 Begin {
763 node: usize,
764 raw: bool,
765 elem: bool,
766 },
767 Key(&'t str, usize),
769 Scalar {
771 node: usize,
772 elem: bool,
773 },
774 End {
775 node: usize,
776 raw: bool,
777 },
778}
779
780#[derive(Debug)]
781struct Frame {
782 node: usize,
783 at: usize,
784 raw: bool,
785}
786
787#[derive(Debug)]
789pub struct Tokens<'t, 'a> {
790 tree: &'t JsonTree<'a>,
791 stack: Vec<Frame>,
792 start: bool,
793 value: Option<usize>,
795}
796
797impl<'t> Tokens<'t, '_> {
798 fn enter(&mut self, node: usize, elem: bool) -> Token<'t> {
799 if self.tree.is_scalar(node) {
800 Token::Scalar { node, elem }
801 } else {
802 self.stack.push(Frame { node, at: 0, raw: false });
803 Token::Begin { node, raw: false, elem }
804 }
805 }
806}
807
808impl<'t> Iterator for Tokens<'t, '_> {
809 type Item = Token<'t>;
810
811 fn next(&mut self) -> Option<Token<'t>> {
812 if let Some(value) = self.value.take() {
813 return Some(self.enter(value, false));
814 }
815 if self.start {
816 self.start = false;
817 if self.tree.is_scalar(0) {
818 self.stack.push(Frame { node: 0, at: 0, raw: true });
819 return Some(Token::Begin { node: 0, raw: true, elem: false });
820 }
821 return Some(self.enter(0, false));
822 }
823 let tree = self.tree;
824 let frame = self.stack.last_mut()?;
825 let (node, at, raw) = (frame.node, frame.at, frame.raw);
826 frame.at += 1;
827 if raw {
828 if at == 0 {
829 return Some(Token::Scalar { node, elem: false });
830 }
831 } else {
832 match tree.item(node) {
833 JsonItem::Array(items) => {
834 if let Some(&child) = items.get(at) {
835 return Some(self.enter(child, true));
836 }
837 }
838 JsonItem::Object(fields) => {
839 if let Some((name, value)) = fields.get(at) {
840 self.value = Some(*value);
841 return Some(Token::Key(name, *value));
842 }
843 }
844 _ => {}
845 }
846 }
847 self.stack.pop();
848 Some(Token::End { node, raw })
849 }
850}