use crate::engine::error::Result;
use bumpalo::Bump;
use datalogic_rs::{Engine, Logic};
use datavalue::{DataValue, OwnedDataValue};
use log::error;
use std::cell::RefCell;
use std::sync::Arc;
const ARENA_INITIAL_CAPACITY: usize = 128 * 1024;
thread_local! {
static EVAL_ARENA: RefCell<Bump> = RefCell::new(Bump::with_capacity(ARENA_INITIAL_CAPACITY));
}
#[inline]
pub(crate) fn eval_to_owned(
engine: &Engine,
compiled: &Logic,
context: &OwnedDataValue,
) -> std::result::Result<OwnedDataValue, datalogic_rs::Error> {
EVAL_ARENA.with(|cell| {
let mut arena = cell.borrow_mut();
arena.reset();
let r = engine.evaluate(compiled, context, &arena)?;
Ok(r.to_owned())
})
}
#[inline]
pub(crate) fn with_arena<R>(f: impl FnOnce(&Bump) -> R) -> R {
EVAL_ARENA.with(|cell| {
let mut arena = cell.borrow_mut();
arena.reset();
f(&arena)
})
}
pub(crate) struct ArenaContext<'a> {
arena: &'a Bump,
top_keys: Vec<&'a str>,
top_values: Vec<DataValue<'a>>,
depth2: Vec<Option<Depth2Cache<'a>>>,
}
struct Depth2Cache<'a> {
keys: Vec<&'a str>,
values: Vec<DataValue<'a>>,
}
impl<'a> ArenaContext<'a> {
pub fn from_owned(ctx: &OwnedDataValue, arena: &'a Bump) -> Self {
let mut top_keys: Vec<&'a str> = Vec::with_capacity(4);
let mut top_values: Vec<DataValue<'a>> = Vec::with_capacity(4);
let mut depth2: Vec<Option<Depth2Cache<'a>>> = Vec::with_capacity(4);
if let OwnedDataValue::Object(pairs) = ctx {
for (k, v) in pairs {
top_keys.push(arena.alloc_str(k));
match v {
OwnedDataValue::Object(children) => {
let mut d2_keys: Vec<&'a str> = Vec::with_capacity(children.len());
let mut d2_vals: Vec<DataValue<'a>> = Vec::with_capacity(children.len());
for (ck, cv) in children {
d2_keys.push(arena.alloc_str(ck));
d2_vals.push(cv.to_arena(arena));
}
let slice = build_object_slice(arena, &d2_keys, &d2_vals);
top_values.push(DataValue::Object(slice));
depth2.push(Some(Depth2Cache {
keys: d2_keys,
values: d2_vals,
}));
}
_ => {
top_values.push(v.to_arena(arena));
depth2.push(None);
}
}
}
}
Self {
arena,
top_keys,
top_values,
depth2,
}
}
pub fn as_data_value(&self) -> DataValue<'a> {
let slice = build_object_slice(self.arena, &self.top_keys, &self.top_values);
DataValue::Object(slice)
}
#[inline]
pub fn arena(&self) -> &'a Bump {
self.arena
}
pub fn apply_mutation_parts_write_through(
&mut self,
owned_ctx: &mut OwnedDataValue,
parts: &[Arc<str>],
value_av: DataValue<'a>,
apply: impl FnOnce(&mut OwnedDataValue),
) {
apply(owned_ctx);
if !self.try_splice_write_parts(parts, value_av) {
self.refresh_after_write_parts(owned_ctx, parts);
}
#[cfg(test)]
self.assert_matches_owned(owned_ctx);
}
fn try_splice_write_parts(&mut self, parts: &[Arc<str>], value_av: DataValue<'a>) -> bool {
if parts.len() < 2 {
return false;
}
if parts[2..].iter().any(|p| p.parse::<usize>().is_ok()) {
return false;
}
let top = strip_hash(&parts[0]);
let Some(top_idx) = self.top_keys.iter().position(|k| *k == top) else {
return false;
};
let arena = self.arena;
let Some(d2) = self.depth2[top_idx].as_mut() else {
return false;
};
let d2_key = strip_hash(&parts[1]);
let d2_pos = d2.keys.iter().position(|k| *k == d2_key);
if parts.len() == 2 {
match d2_pos {
Some(pos) => d2.values[pos] = value_av,
None => {
d2.keys.push(arena.alloc_str(d2_key));
d2.values.push(value_av);
}
}
} else {
let Some(pos) = d2_pos else {
return false;
};
let Some(new_subtree) =
splice_object_write(arena, d2.values[pos], &parts[2..], value_av)
else {
return false;
};
d2.values[pos] = new_subtree;
}
let slice = build_object_slice(arena, &d2.keys, &d2.values);
self.top_values[top_idx] = DataValue::Object(slice);
true
}
#[cfg(test)]
fn assert_matches_owned(&self, owned_ctx: &OwnedDataValue) {
let rebuilt = ArenaContext::from_owned(owned_ctx, self.arena);
assert_eq!(
self.as_data_value().to_owned(),
rebuilt.as_data_value().to_owned(),
"arena cache diverged from owned context"
);
}
pub fn refresh_for_path(&mut self, owned_ctx: &OwnedDataValue, path: &str) {
self.refresh_after_write(owned_ctx, path);
}
pub fn refresh_for_path_parts(&mut self, owned_ctx: &OwnedDataValue, parts: &[Arc<str>]) {
self.refresh_after_write_parts(owned_ctx, parts);
}
fn refresh_after_write_parts(&mut self, owned_ctx: &OwnedDataValue, parts: &[Arc<str>]) {
let top_raw: &str = match parts.first() {
Some(p) if !p.is_empty() => p,
_ => {
self.rebuild_all_from(owned_ctx);
return;
}
};
let top = top_raw.strip_prefix('#').unwrap_or(top_raw);
fn strip<'p>(p: &'p Arc<str>) -> &'p str {
let s: &'p str = p;
s.strip_prefix('#').unwrap_or(s)
}
let depth2_key: Option<&str> = parts.get(1).map(strip);
let depth3_key: Option<&str> = parts.get(2).map(strip);
self.refresh_after_write_inner(owned_ctx, top, depth2_key, depth3_key);
}
fn refresh_after_write(&mut self, owned_ctx: &OwnedDataValue, path: &str) {
let mut parts = path.split('.');
let top_raw = match parts.next() {
Some(p) if !p.is_empty() => p,
_ => {
self.rebuild_all_from(owned_ctx);
return;
}
};
let top = top_raw.strip_prefix('#').unwrap_or(top_raw);
let depth2_key = parts.next().map(|p| p.strip_prefix('#').unwrap_or(p));
let depth3_key = parts.next().map(|p| p.strip_prefix('#').unwrap_or(p));
self.refresh_after_write_inner(owned_ctx, top, depth2_key, depth3_key);
}
fn refresh_after_write_inner(
&mut self,
owned_ctx: &OwnedDataValue,
top: &str,
depth2_key: Option<&str>,
_depth3_key: Option<&str>,
) {
let OwnedDataValue::Object(owned_pairs) = owned_ctx else {
self.rebuild_all_from(owned_ctx);
return;
};
let owned_top_val = owned_pairs.iter().find(|(k, _)| k == top).map(|(_, v)| v);
let top_idx = self.top_keys.iter().position(|k| *k == top);
match (owned_top_val, top_idx) {
(None, Some(idx)) => {
self.top_keys.remove(idx);
self.top_values.remove(idx);
self.depth2.remove(idx);
}
(Some(new_val), idx_opt) => {
let idx = match idx_opt {
Some(i) => i,
None => {
self.top_keys.push(self.arena.alloc_str(top));
self.top_values.push(DataValue::Null);
self.depth2.push(None);
self.top_keys.len() - 1
}
};
match (new_val, depth2_key, &mut self.depth2[idx]) {
(OwnedDataValue::Object(new_children), Some(child_key), Some(d2)) => {
if let Some(new_child) = new_children
.iter()
.find(|(k, _)| k == child_key)
.map(|(_, v)| v)
{
let child_arena = new_child.to_arena(self.arena);
if let Some(pos) = d2.keys.iter().position(|k| *k == child_key) {
d2.values[pos] = child_arena;
} else {
d2.keys.push(self.arena.alloc_str(child_key));
d2.values.push(child_arena);
}
if d2.keys.len() != new_children.len() {
self.rebuild_top_slot(owned_top_val.unwrap(), idx);
return;
}
} else {
if let Some(pos) = d2.keys.iter().position(|k| *k == child_key) {
d2.keys.remove(pos);
d2.values.remove(pos);
}
}
let slice = build_object_slice(self.arena, &d2.keys, &d2.values);
self.top_values[idx] = DataValue::Object(slice);
}
_ => {
self.rebuild_top_slot(new_val, idx);
}
}
}
(None, None) => { }
}
}
fn rebuild_top_slot(&mut self, owned: &OwnedDataValue, idx: usize) {
match owned {
OwnedDataValue::Object(children) => {
let mut d2_keys: Vec<&'a str> = Vec::with_capacity(children.len());
let mut d2_vals: Vec<DataValue<'a>> = Vec::with_capacity(children.len());
for (ck, cv) in children {
d2_keys.push(self.arena.alloc_str(ck));
d2_vals.push(cv.to_arena(self.arena));
}
let slice = build_object_slice(self.arena, &d2_keys, &d2_vals);
self.top_values[idx] = DataValue::Object(slice);
self.depth2[idx] = Some(Depth2Cache {
keys: d2_keys,
values: d2_vals,
});
}
_ => {
self.top_values[idx] = owned.to_arena(self.arena);
self.depth2[idx] = None;
}
}
}
fn rebuild_all_from(&mut self, ctx: &OwnedDataValue) {
let rebuilt = ArenaContext::from_owned(ctx, self.arena);
self.top_keys = rebuilt.top_keys;
self.top_values = rebuilt.top_values;
self.depth2 = rebuilt.depth2;
}
}
fn build_object_slice<'a>(
arena: &'a Bump,
keys: &[&'a str],
values: &[DataValue<'a>],
) -> &'a [(&'a str, DataValue<'a>)] {
debug_assert_eq!(keys.len(), values.len());
arena.alloc_slice_fill_iter(keys.iter().zip(values.iter()).map(|(k, v)| (*k, *v)))
}
#[inline]
fn strip_hash(part: &str) -> &str {
part.strip_prefix('#').unwrap_or(part)
}
fn splice_object_write<'a>(
arena: &'a Bump,
current: DataValue<'a>,
parts: &[Arc<str>],
value: DataValue<'a>,
) -> Option<DataValue<'a>> {
let DataValue::Object(pairs) = current else {
return None;
};
let key = strip_hash(&parts[0]);
let pos = pairs.iter().position(|(k, _)| *k == key);
let new_slice = if parts.len() == 1 {
match pos {
Some(p) => alloc_pairs_with_replaced(arena, pairs, p, value),
None => alloc_pairs_with_appended(arena, pairs, arena.alloc_str(key), value),
}
} else {
match pos {
Some(p) => {
let child = splice_object_write(arena, pairs[p].1, &parts[1..], value)?;
alloc_pairs_with_replaced(arena, pairs, p, child)
}
None => {
let chain = build_object_chain(arena, &parts[1..], value);
alloc_pairs_with_appended(arena, pairs, arena.alloc_str(key), chain)
}
}
};
Some(DataValue::Object(new_slice))
}
fn build_object_chain<'a>(
arena: &'a Bump,
parts: &[Arc<str>],
value: DataValue<'a>,
) -> DataValue<'a> {
let mut acc = value;
for part in parts.iter().rev() {
let key: &'a str = arena.alloc_str(strip_hash(part));
let slice = arena.alloc_slice_fill_iter(std::iter::once((key, acc)));
acc = DataValue::Object(slice);
}
acc
}
fn alloc_pairs_with_replaced<'a>(
arena: &'a Bump,
pairs: &'a [(&'a str, DataValue<'a>)],
replace_idx: usize,
new_value: DataValue<'a>,
) -> &'a [(&'a str, DataValue<'a>)] {
arena.alloc_slice_fill_with(pairs.len(), |i| {
if i == replace_idx {
(pairs[i].0, new_value)
} else {
pairs[i]
}
})
}
fn alloc_pairs_with_appended<'a>(
arena: &'a Bump,
pairs: &'a [(&'a str, DataValue<'a>)],
key: &'a str,
value: DataValue<'a>,
) -> &'a [(&'a str, DataValue<'a>)] {
arena.alloc_slice_fill_with(pairs.len() + 1, |i| {
if i < pairs.len() {
pairs[i]
} else {
(key, value)
}
})
}
pub fn evaluate_condition(
engine: &Engine,
compiled_condition: Option<&Arc<Logic>>,
context: &OwnedDataValue,
) -> Result<bool> {
match compiled_condition {
Some(compiled) => match eval_to_owned(engine, compiled, context) {
Ok(value) => Ok(matches!(value, OwnedDataValue::Bool(true))),
Err(e) => {
error!("Failed to evaluate condition: {:?}", e);
Ok(false)
}
},
None => Ok(true),
}
}
pub fn evaluate_condition_in_arena(
engine: &Engine,
compiled_condition: Option<&Arc<Logic>>,
ctx: DataValue<'_>,
arena: &Bump,
) -> Result<bool> {
match compiled_condition {
Some(compiled) => match engine.evaluate(compiled, ctx, arena) {
Ok(value) => Ok(matches!(value, DataValue::Bool(true))),
Err(e) => {
error!("Failed to evaluate condition: {:?}", e);
Ok(false)
}
},
None => Ok(true),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::engine::utils::{set_nested_value, set_nested_value_parts};
use serde_json::json;
fn dv(v: serde_json::Value) -> OwnedDataValue {
OwnedDataValue::from(&v)
}
fn parts_of(path: &str) -> Vec<Arc<str>> {
path.split('.').map(Arc::from).collect()
}
#[test]
fn write_through_differential_corpus() {
let mut owned = dv(json!({
"data": {
"input": {"big": {"nested": [1, 2, 3]}, "flag": true},
"MT103": {"20": "REF", "72": "existing"},
"items": [{"x": 1}, {"x": 2}],
"scalar": "leaf"
},
"metadata": {"processed_at": "t0"},
"temp_data": {}
}));
with_arena(|arena| {
let mut ctx = ArenaContext::from_owned(&owned, arena);
let corpus: &[(&str, serde_json::Value)] = &[
("data.MT103.20", json!("NEWREF")),
("data.MT103.23B", json!("CRED")),
("data.MT103.32A.amount", json!(1500.25)),
("data.MT103.32A.currency", json!("USD")),
("data.MT103.a.b.c", json!({"deep": [1, {"k": "v"}]})),
("data.MT103.32A", json!({"replaced": true})),
("data.MT103", json!({"fresh": {"start": 1}})),
("data.scalar2", json!("plain")),
("data.MT202", json!({"20": "REF2"})),
("data.MT205.20", json!("REF5")),
("data.MT103.#20", json!("HASH20")),
("data.MT103.##", json!("HASHKEY")),
("data.items.0.x", json!(10)),
("data.items.3", json!({"x": 4})),
("data.scalar.sub.key", json!("ignored")),
("metadata.progress.workflow_id", json!("wf1")),
("temp_data.uetr", json!("uuid-here")),
("other_top.x.y", json!(42)),
];
for (path, val) in corpus {
let parts = parts_of(path);
let owned_val = dv(val.clone());
let value_av = owned_val.to_arena(arena);
ctx.apply_mutation_parts_write_through(&mut owned, &parts, value_av, |c| {
set_nested_value_parts(c, &parts, owned_val.clone());
});
}
assert_eq!(owned["data"]["MT103"]["20"], dv(json!("HASH20")));
assert_eq!(owned["data"]["MT103"]["fresh"]["start"], dv(json!(1)));
assert_eq!(owned["data"]["items"][3]["x"], dv(json!(4)));
assert_eq!(owned["data"]["scalar"], dv(json!("leaf")));
assert_eq!(
owned["metadata"]["progress"]["workflow_id"],
dv(json!("wf1"))
);
});
}
#[test]
fn write_through_reads_see_latest_state() {
let mut owned = dv(json!({
"data": {"doc": {"a": 1}},
"metadata": {},
"temp_data": {}
}));
with_arena(|arena| {
let mut ctx = ArenaContext::from_owned(&owned, arena);
for i in 0..10u64 {
let key = format!("data.doc.f{i}");
let parts = parts_of(&key);
let owned_val = OwnedDataValue::from(i);
let value_av = owned_val.to_arena(arena);
ctx.apply_mutation_parts_write_through(&mut owned, &parts, value_av, |c| {
set_nested_value_parts(c, &parts, owned_val.clone());
});
assert_eq!(ctx.as_data_value().to_owned(), owned);
}
});
}
#[test]
fn refresh_and_write_through_interleave() {
let mut owned = dv(json!({
"data": {},
"metadata": {},
"temp_data": {}
}));
with_arena(|arena| {
let mut ctx = ArenaContext::from_owned(&owned, arena);
set_nested_value(&mut owned, "data.input", dv(json!({"payload": {"v": 7}})));
ctx.refresh_for_path(&owned, "data.input");
assert_eq!(ctx.as_data_value().to_owned(), owned);
let parts = parts_of("data.out.v");
let owned_val = dv(json!(7));
let value_av = owned_val.to_arena(arena);
ctx.apply_mutation_parts_write_through(&mut owned, &parts, value_av, |c| {
set_nested_value_parts(c, &parts, owned_val.clone());
});
set_nested_value(
&mut owned,
"metadata.progress",
dv(json!({"task_id": "t1"})),
);
ctx.refresh_for_path(&owned, "metadata.progress");
assert_eq!(ctx.as_data_value().to_owned(), owned);
});
}
}