From 22f324ce648c3a0c4f4e151eb59bc7378e31b3af Mon Sep 17 00:00:00 2001
From: sundyli <543950155@qq.com>
Date: Wed, 9 Oct 2024 08:57:57 +0800
Subject: [PATCH] feat: add lazy value

---
 src/functions.rs   |  2 +-
 src/lazy_value.rs  | 74 ++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs         |  3 ++
 src/parser.rs      | 11 +++++++
 tests/it/encode.rs | 12 ++++++--
 5 files changed, 99 insertions(+), 3 deletions(-)
 create mode 100644 src/lazy_value.rs

diff --git a/src/functions.rs b/src/functions.rs
index 847d283..cae96cf 100644
--- a/src/functions.rs
+++ b/src/functions.rs
@@ -2619,7 +2619,7 @@ pub fn type_of(value: &[u8]) -> Result<&'static str, Error> {
 
 // Check whether the value is `JSONB` format,
 // for compatibility with previous `JSON` string.
-fn is_jsonb(value: &[u8]) -> bool {
+pub(crate) fn is_jsonb(value: &[u8]) -> bool {
     if let Some(v) = value.first() {
         if matches!(*v, ARRAY_PREFIX | OBJECT_PREFIX | SCALAR_PREFIX) {
             return true;
diff --git a/src/lazy_value.rs b/src/lazy_value.rs
new file mode 100644
index 0000000..a7f6148
--- /dev/null
+++ b/src/lazy_value.rs
@@ -0,0 +1,74 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::borrow::Cow;
+use std::fmt::Debug;
+
+use crate::array_length;
+use crate::ser::Encoder;
+use crate::Value;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum LazyValue<'a> {
+    Value(Value<'a>),
+    // Raw JSONB bytes
+    Raw(Cow<'a, [u8]>),
+}
+
+impl<'a> LazyValue<'a> {
+    /// Serialize the JSONB Value into a byte stream.
+    pub fn write_to_vec(&self, buf: &mut Vec<u8>) {
+        match self {
+            LazyValue::Value(v) => {
+                let mut encoder = Encoder::new(buf);
+                encoder.encode(v)
+            }
+            LazyValue::Raw(v) => buf.extend_from_slice(v),
+        };
+    }
+
+    /// Serialize the JSONB Value into a byte stream.
+    pub fn to_vec(&self) -> Vec<u8> {
+        match self {
+            LazyValue::Value(value) => {
+                let mut buf = Vec::new();
+                value.write_to_vec(&mut buf);
+                buf
+            }
+            LazyValue::Raw(cow) => cow.to_vec(),
+        }
+    }
+
+    // TODO migrate more functions to be methods of LazyValue
+    pub fn array_length(&self) -> Option<usize> {
+        match self {
+            LazyValue::Value(Value::Array(arr)) => Some(arr.len()),
+            LazyValue::Raw(cow) => array_length(cow.as_ref()),
+            _ => None,
+        }
+    }
+
+    pub fn to_value(&'a self) -> Cow<Value<'a>> {
+        match self {
+            LazyValue::Value(v) => Cow::Borrowed(v),
+            LazyValue::Raw(v) => Cow::Owned(crate::from_slice(v.as_ref()).unwrap()),
+        }
+    }
+}
+
+impl<'a> From<Value<'a>> for LazyValue<'a> {
+    fn from(value: Value<'a>) -> Self {
+        LazyValue::Value(value)
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 6e26b26..22a8799 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -73,6 +73,7 @@ mod iterator;
 mod jentry;
 pub mod jsonpath;
 pub mod keypath;
+mod lazy_value;
 mod number;
 mod parser;
 mod ser;
@@ -84,6 +85,8 @@ pub use error::Error;
 #[allow(unused_imports)]
 pub use from::*;
 pub use functions::*;
+pub use lazy_value::*;
 pub use number::Number;
+pub use parser::parse_lazy_value;
 pub use parser::parse_value;
 pub use value::*;
diff --git a/src/parser.rs b/src/parser.rs
index 03067ae..5c382c4 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -14,6 +14,9 @@
 
 use std::borrow::Cow;
 
+use crate::is_jsonb;
+use crate::lazy_value::LazyValue;
+
 use super::constants::*;
 use super::error::Error;
 use super::error::ParseErrorCode;
@@ -30,6 +33,14 @@ pub fn parse_value(buf: &[u8]) -> Result<Value<'_>, Error> {
     parser.parse()
 }
 
+pub fn parse_lazy_value(buf: &[u8]) -> Result<LazyValue<'_>, Error> {
+    if !is_jsonb(buf) {
+        parse_value(buf).map(LazyValue::Value)
+    } else {
+        Ok(LazyValue::Raw(Cow::Borrowed(buf)))
+    }
+}
+
 struct Parser<'a> {
     buf: &'a [u8],
     idx: usize,
diff --git a/tests/it/encode.rs b/tests/it/encode.rs
index 841e6e2..3d202f8 100644
--- a/tests/it/encode.rs
+++ b/tests/it/encode.rs
@@ -14,7 +14,7 @@
 
 use std::borrow::Cow;
 
-use jsonb::{Number, Object, Value};
+use jsonb::{parse_lazy_value, Number, Object, Value};
 
 #[test]
 fn test_encode_null() {
@@ -133,9 +133,17 @@ fn test_encode_float64() {
 
 #[test]
 fn test_encode_array() {
+    let raw = b"\x80\0\0\x02\x30\0\0\0\x40\0\0\0";
     assert_eq!(
         &Value::Array(vec![Value::Bool(false), Value::Bool(true)]).to_vec(),
-        b"\x80\0\0\x02\x30\0\0\0\x40\0\0\0"
+        raw
+    );
+
+    let lazy_value = parse_lazy_value(raw).unwrap();
+    assert_eq!(lazy_value.array_length(), Some(2));
+    assert_eq!(
+        lazy_value.to_value().as_ref(),
+        &Value::Array(vec![Value::Bool(false), Value::Bool(true)])
     );
 }