Skip to content

Commit 754395c

Browse files
davidhewittdavidhewitt
authored
iterative algorithm for take_value (#134)
1 parent dd25fd0 commit 754395c

File tree

1 file changed

+250
-42
lines changed

1 file changed

+250
-42
lines changed

‎crates/jiter/src/value.rs‎

Lines changed: 250 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -99,19 +99,6 @@ impl JsonValue<'static> {
9999
}
100100
}
101101

102-
macro_rules! check_recursion {
103-
($recursion_limit:ident, $index:expr, $($body:tt)*) => {
104-
$recursion_limit = match $recursion_limit.checked_sub(1) {
105-
Some(limit) => limit,
106-
None => return crate::errors::json_err!(RecursionLimitExceeded, $index),
107-
};
108-
109-
$($body)*
110-
111-
$recursion_limit += 1;
112-
};
113-
}
114-
115102
pub(crate) fn take_value_borrowed<'j>(
116103
peek: Peek,
117104
parser: &mut Parser<'j>,
@@ -150,7 +137,7 @@ fn take_value<'j, 's>(
150137
peek: Peek,
151138
parser: &mut Parser<'j>,
152139
tape: &mut Tape,
153-
mut recursion_limit: u8,
140+
recursion_limit: u8,
154141
allow_inf_nan: bool,
155142
create_cow: &impl Fn(StringOutput<'_, 'j>) -> Cow<'s, str>,
156143
) -> JsonResult<JsonValue<'s>> {
@@ -173,42 +160,41 @@ fn take_value<'j, 's>(
173160
}
174161
Peek::Array => {
175162
// we could do something clever about guessing the size of the array
176-
let mut array: SmallVec<[JsonValue<'s>; 8]> = SmallVec::new();
163+
let array = Arc::new(SmallVec::new());
177164
if let Some(peek_first) = parser.array_first()? {
178-
check_recursion!(recursion_limit, parser.index,
179-
let v = take_value(peek_first, parser, tape, recursion_limit, allow_inf_nan, create_cow)?;
180-
);
181-
array.push(v);
182-
while let Some(peek) = parser.array_step()? {
183-
check_recursion!(recursion_limit, parser.index,
184-
let v = take_value(peek, parser, tape, recursion_limit, allow_inf_nan, create_cow)?;
185-
);
186-
array.push(v);
187-
}
165+
take_value_recursive(
166+
peek_first,
167+
RecursedValue::Array(array),
168+
parser,
169+
tape,
170+
recursion_limit,
171+
allow_inf_nan,
172+
create_cow,
173+
)
174+
} else {
175+
Ok(JsonValue::Array(array))
188176
}
189-
Ok(JsonValue::Array(Arc::new(array)))
190177
}
191178
Peek::Object => {
192179
// same for objects
193-
let mut object: LazyIndexMap<Cow<'s, str>, JsonValue<'s>> = LazyIndexMap::new();
180+
let object = Arc::new(LazyIndexMap::new());
194181
if let Some(first_key) = parser.object_first::<StringDecoder>(tape)? {
195182
let first_key = create_cow(first_key);
196-
let peek = parser.peek()?;
197-
check_recursion!(recursion_limit, parser.index,
198-
let first_value = take_value(peek, parser, tape, recursion_limit, allow_inf_nan, create_cow)?;
199-
);
200-
object.insert(first_key, first_value);
201-
while let Some(key) = parser.object_step::<StringDecoder>(tape)? {
202-
let key = create_cow(key);
203-
let peek = parser.peek()?;
204-
check_recursion!(recursion_limit, parser.index,
205-
let value = take_value(peek, parser, tape, recursion_limit, allow_inf_nan, create_cow)?;
206-
);
207-
object.insert(key, value);
208-
}
183+
take_value_recursive(
184+
parser.peek()?,
185+
RecursedValue::Object {
186+
partial: object,
187+
next_key: first_key,
188+
},
189+
parser,
190+
tape,
191+
recursion_limit,
192+
allow_inf_nan,
193+
create_cow,
194+
)
195+
} else {
196+
Ok(JsonValue::Object(object))
209197
}
210-
211-
Ok(JsonValue::Object(Arc::new(object)))
212198
}
213199
_ => {
214200
let n = parser.consume_number::<NumberAny>(peek.into_inner(), allow_inf_nan);
@@ -228,6 +214,228 @@ fn take_value<'j, 's>(
228214
}
229215
}
230216

217+
enum RecursedValue<'s> {
218+
Array(JsonArray<'s>),
219+
Object {
220+
partial: JsonObject<'s>,
221+
next_key: Cow<'s, str>,
222+
},
223+
}
224+
225+
#[inline(never)] // this is an iterative algo called only from take_value, no point in inlining
226+
#[allow(clippy::too_many_lines)] // FIXME?
227+
fn take_value_recursive<'j, 's>(
228+
mut peek: Peek,
229+
mut current_recursion: RecursedValue<'s>,
230+
parser: &mut Parser<'j>,
231+
tape: &mut Tape,
232+
recursion_limit: u8,
233+
allow_inf_nan: bool,
234+
create_cow: &impl Fn(StringOutput<'_, 'j>) -> Cow<'s, str>,
235+
) -> JsonResult<JsonValue<'s>> {
236+
let recursion_limit: usize = recursion_limit.into();
237+
238+
let mut recursion_stack: SmallVec<[RecursedValue; 8]> = SmallVec::new();
239+
240+
macro_rules! push_recursion {
241+
($next_peek:expr, $value:expr) => {
242+
peek = $next_peek;
243+
recursion_stack.push(std::mem::replace(&mut current_recursion, $value));
244+
if recursion_stack.len() >= recursion_limit {
245+
return Err(json_error!(RecursionLimitExceeded, parser.index));
246+
}
247+
};
248+
}
249+
250+
'recursion: loop {
251+
let mut value = match &mut current_recursion {
252+
RecursedValue::Array(array) => {
253+
let array = Arc::get_mut(array).expect("sole writer");
254+
loop {
255+
let value = match peek {
256+
Peek::True => {
257+
parser.consume_true()?;
258+
JsonValue::Bool(true)
259+
}
260+
Peek::False => {
261+
parser.consume_false()?;
262+
JsonValue::Bool(false)
263+
}
264+
Peek::Null => {
265+
parser.consume_null()?;
266+
JsonValue::Null
267+
}
268+
Peek::String => {
269+
let s = parser.consume_string::<StringDecoder>(tape, false)?;
270+
JsonValue::Str(create_cow(s))
271+
}
272+
Peek::Array => {
273+
let array = Arc::new(SmallVec::new());
274+
if let Some(next_peek) = parser.array_first()? {
275+
push_recursion!(next_peek, RecursedValue::Array(array));
276+
// immediately jump to process the first value in the array
277+
continue 'recursion;
278+
}
279+
JsonValue::Array(array)
280+
}
281+
Peek::Object => {
282+
let object = Arc::new(LazyIndexMap::new());
283+
if let Some(next_key) = parser.object_first::<StringDecoder>(tape)? {
284+
push_recursion!(
285+
parser.peek()?,
286+
RecursedValue::Object {
287+
partial: object,
288+
next_key: create_cow(next_key)
289+
}
290+
);
291+
continue 'recursion;
292+
}
293+
JsonValue::Object(object)
294+
}
295+
_ => {
296+
let n = parser
297+
.consume_number::<NumberAny>(peek.into_inner(), allow_inf_nan)
298+
.map_err(|e| {
299+
if !peek.is_num() {
300+
json_error!(ExpectedSomeValue, parser.index)
301+
} else {
302+
e
303+
}
304+
})?;
305+
match n {
306+
NumberAny::Int(NumberInt::Int(int)) => JsonValue::Int(int),
307+
NumberAny::Int(NumberInt::BigInt(big_int)) => JsonValue::BigInt(big_int),
308+
NumberAny::Float(float) => JsonValue::Float(float),
309+
}
310+
}
311+
};
312+
313+
// now try to advance position in the current array
314+
if let Some(next_peek) = parser.array_step()? {
315+
array.push(value);
316+
peek = next_peek;
317+
// array continuing
318+
continue;
319+
}
320+
321+
let RecursedValue::Array(mut array) = current_recursion else {
322+
unreachable!("known to be in array recursion");
323+
};
324+
325+
Arc::get_mut(&mut array).expect("sole writer to value").push(value);
326+
break JsonValue::Array(array);
327+
}
328+
}
329+
RecursedValue::Object { partial, next_key } => {
330+
let partial = Arc::get_mut(partial).expect("sole writer");
331+
loop {
332+
let value = match peek {
333+
Peek::True => {
334+
parser.consume_true()?;
335+
JsonValue::Bool(true)
336+
}
337+
Peek::False => {
338+
parser.consume_false()?;
339+
JsonValue::Bool(false)
340+
}
341+
Peek::Null => {
342+
parser.consume_null()?;
343+
JsonValue::Null
344+
}
345+
Peek::String => {
346+
let s = parser.consume_string::<StringDecoder>(tape, false)?;
347+
JsonValue::Str(create_cow(s))
348+
}
349+
Peek::Array => {
350+
let array = Arc::new(SmallVec::new());
351+
if let Some(next_peek) = parser.array_first()? {
352+
push_recursion!(next_peek, RecursedValue::Array(array));
353+
// immediately jump to process the first value in the array
354+
continue 'recursion;
355+
}
356+
JsonValue::Array(array)
357+
}
358+
Peek::Object => {
359+
let object = Arc::new(LazyIndexMap::new());
360+
if let Some(yet_another_key) = parser.object_first::<StringDecoder>(tape)? {
361+
push_recursion!(
362+
parser.peek()?,
363+
RecursedValue::Object {
364+
partial: object,
365+
next_key: create_cow(yet_another_key)
366+
}
367+
);
368+
continue 'recursion;
369+
}
370+
JsonValue::Object(object)
371+
}
372+
_ => {
373+
let n = parser
374+
.consume_number::<NumberAny>(peek.into_inner(), allow_inf_nan)
375+
.map_err(|e| {
376+
if !peek.is_num() {
377+
json_error!(ExpectedSomeValue, parser.index)
378+
} else {
379+
e
380+
}
381+
})?;
382+
match n {
383+
NumberAny::Int(NumberInt::Int(int)) => JsonValue::Int(int),
384+
NumberAny::Int(NumberInt::BigInt(big_int)) => JsonValue::BigInt(big_int),
385+
NumberAny::Float(float) => JsonValue::Float(float),
386+
}
387+
}
388+
};
389+
390+
// now try to advance position in the current object
391+
if let Some(yet_another_key) = parser.object_step::<StringDecoder>(tape)?.map(create_cow) {
392+
// object continuing
393+
partial.insert(std::mem::replace(next_key, yet_another_key), value);
394+
peek = parser.peek()?;
395+
continue;
396+
}
397+
398+
let RecursedValue::Object { mut partial, next_key } = current_recursion else {
399+
unreachable!("known to be in object recursion");
400+
};
401+
402+
Arc::get_mut(&mut partial).expect("sole writer").insert(next_key, value);
403+
break JsonValue::Object(partial);
404+
}
405+
}
406+
};
407+
408+
// current array or object has finished;
409+
// try to pop and continue with the parent
410+
peek = loop {
411+
if let Some(next_recursion) = recursion_stack.pop() {
412+
current_recursion = next_recursion;
413+
} else {
414+
return Ok(value);
415+
}
416+
417+
value = match current_recursion {
418+
RecursedValue::Array(mut array) => {
419+
Arc::get_mut(&mut array).expect("sole writer").push(value);
420+
if let Some(next_peek) = parser.array_step()? {
421+
current_recursion = RecursedValue::Array(array);
422+
break next_peek;
423+
}
424+
JsonValue::Array(array)
425+
}
426+
RecursedValue::Object { mut partial, next_key } => {
427+
Arc::get_mut(&mut partial).expect("sole writer").insert(next_key, value);
428+
if let Some(next_key) = parser.object_step::<StringDecoder>(tape)?.map(create_cow) {
429+
current_recursion = RecursedValue::Object { partial, next_key };
430+
break parser.peek()?;
431+
}
432+
JsonValue::Object(partial)
433+
}
434+
}
435+
};
436+
}
437+
}
438+
231439
/// like `take_value`, but nothing is returned, should be faster than `take_value`, useful when you don't care
232440
/// about the value, but just want to consume it
233441
pub(crate) fn take_value_skip(

0 commit comments

Comments
 (0)