pub fn intern_hop_fields(per_query_hops: &[Vec<String>]) -> Vec<String> {
let mut out: Vec<String> = Vec::new();
for hops in per_query_hops {
for h in hops {
if !out.iter().any(|x| x == h) {
out.push(h.clone());
}
}
}
out
}
pub struct RefWalkEdges {
edges: Vec<(u32, u32, u32)>, cap: usize,
truncated: bool,
}
impl RefWalkEdges {
pub fn new(cap: usize) -> Self {
Self {
edges: Vec::new(),
cap,
truncated: false,
}
}
pub fn truncated(&self) -> bool {
self.truncated
}
#[cfg(test)]
pub fn len(&self) -> usize {
self.edges.len()
}
#[cfg(test)]
pub fn is_empty(&self) -> bool {
self.edges.is_empty()
}
pub fn push(&mut self, src: u32, field_id: u32, dst: u32) {
if self.edges.len() >= self.cap {
self.truncated = true;
return;
}
self.edges.push((src, field_id, dst));
}
pub fn into_csr(mut self, n: usize) -> (Vec<u32>, Vec<u32>, Vec<u32>) {
let mut off = vec![0u32; n + 1];
for &(s, _, _) in &self.edges {
off[s as usize + 1] += 1;
}
for i in 0..n {
off[i + 1] += off[i];
}
let total = self.edges.len();
let mut tgt = vec![0u32; total];
let mut fid = vec![0u32; total];
let mut cursor: Vec<u32> = off[..n].to_vec();
for (s, f, d) in self.edges.drain(..) {
let p = cursor[s as usize] as usize;
tgt[p] = d;
fid[p] = f;
cursor[s as usize] += 1;
}
(off, tgt, fid)
}
}
pub const REFWALK_EDGE_CAP: usize = 5_000_000;
pub struct RefWalkTails {
values: std::collections::HashMap<u32, crate::query::model::QueryValue>,
cap: usize,
truncated: bool,
}
impl RefWalkTails {
pub fn new(cap: usize) -> Self {
Self {
values: std::collections::HashMap::new(),
cap,
truncated: false,
}
}
pub fn truncated(&self) -> bool {
self.truncated
}
#[cfg(test)]
pub fn len(&self) -> usize {
self.values.len()
}
#[cfg(test)]
#[allow(dead_code)]
pub fn is_empty(&self) -> bool {
self.values.is_empty()
}
pub fn insert(&mut self, dense_idx: u32, value: crate::query::model::QueryValue) {
if self.values.len() >= self.cap && !self.values.contains_key(&dense_idx) {
self.truncated = true;
return;
}
self.values.insert(dense_idx, value);
}
#[cfg(test)]
pub fn get(&self, dense_idx: u32) -> Option<&crate::query::model::QueryValue> {
self.values.get(&dense_idx)
}
pub fn into_map(self) -> std::collections::HashMap<u32, crate::query::model::QueryValue> {
self.values
}
}
pub fn decode_primitive_tail(
off: u32,
ty: crate::types::HprofType,
blob: &[u8],
) -> Option<crate::query::model::QueryValue> {
use crate::query::model::QueryValue;
use crate::types::HprofType;
let o = off as usize;
let read_be = |o: usize, n: usize| -> Option<u64> {
let end = o + n;
if end > blob.len() {
return None;
}
let mut v: u64 = 0;
for &b in &blob[o..end] {
v = (v << 8) | b as u64;
}
Some(v)
};
match ty {
HprofType::Boolean => blob.get(o).map(|&b| QueryValue::Bool(b != 0)),
HprofType::Byte => blob.get(o).map(|&b| QueryValue::Int(b as i8 as i64)),
HprofType::Short => read_be(o, 2).map(|v| QueryValue::Int(v as i16 as i64)),
HprofType::Char => read_be(o, 2).map(|v| QueryValue::Int(v as i64)),
HprofType::Int => read_be(o, 4).map(|v| QueryValue::Int(v as i32 as i64)),
HprofType::Long => read_be(o, 8).map(|v| QueryValue::Int(v as i64)),
HprofType::Float => {
read_be(o, 4).map(|v| QueryValue::Float(f32::from_bits(v as u32) as f64))
}
HprofType::Double => read_be(o, 8).map(|v| QueryValue::Float(f64::from_bits(v))),
HprofType::Object => None,
}
}
pub fn refwalk_field_names(q: &crate::query::ast::Query) -> Vec<String> {
use crate::query::ast::{Attr, Expr, Predicate, SelectItem};
fn collect_attr(a: &Attr, out: &mut Vec<String>) {
if let Attr::RefPath { hops, tail, .. } = a {
for h in hops {
if !out.iter().any(|x| x == h) {
out.push(h.clone());
}
}
collect_attr(tail, out);
}
}
fn collect_pred(p: &Predicate, out: &mut Vec<String>) {
match p {
Predicate::And(a, b) | Predicate::Or(a, b) => {
collect_pred(a, out);
collect_pred(b, out);
}
Predicate::Not(a) => collect_pred(a, out),
Predicate::Compare { lhs, .. } => {
if let Expr::Attr(a) = lhs {
collect_attr(a, out);
}
}
Predicate::InSubquery { .. } | Predicate::InstanceOf(_) => {}
Predicate::Exists { .. } => {}
}
}
let mut out = Vec::new();
for item in &q.select {
match item {
SelectItem::Attr(a) => collect_attr(a, &mut out),
SelectItem::Aggregate { arg, .. } => {
if let SelectItem::Attr(a) = arg.as_ref() {
collect_attr(a, &mut out);
}
}
SelectItem::Star => {}
SelectItem::Path { .. } => {}
SelectItem::ToString(_) => {}
SelectItem::Expr(_) => {
unreachable!("Expr select item reached before arithmetic wiring")
}
}
}
if let Some(pred) = &q.where_ {
collect_pred(pred, &mut out);
}
out
}
pub fn refwalk_tail_field_names(q: &crate::query::ast::Query) -> Vec<String> {
use crate::query::ast::{Attr, Expr, Predicate, SelectItem};
fn collect_attr(a: &Attr, out: &mut Vec<String>) {
if let Attr::RefPath { tail, .. } = a {
match tail.as_ref() {
Attr::Field(name) => {
if !out.iter().any(|x| x == name) {
out.push(name.clone());
}
}
other => collect_attr(other, out),
}
}
}
fn collect_pred(p: &Predicate, out: &mut Vec<String>) {
match p {
Predicate::And(a, b) | Predicate::Or(a, b) => {
collect_pred(a, out);
collect_pred(b, out);
}
Predicate::Not(a) => collect_pred(a, out),
Predicate::Compare { lhs, .. } => {
if let Expr::Attr(a) = lhs {
collect_attr(a, out);
}
}
Predicate::InSubquery { .. } | Predicate::InstanceOf(_) => {}
Predicate::Exists { .. } => {}
}
}
let mut out = Vec::new();
for item in &q.select {
match item {
SelectItem::Attr(a) => collect_attr(a, &mut out),
SelectItem::Aggregate { arg, .. } => {
if let SelectItem::Attr(a) = arg.as_ref() {
collect_attr(a, &mut out);
}
}
SelectItem::Star => {}
SelectItem::Path { .. } => {}
SelectItem::ToString(_) => {}
SelectItem::Expr(_) => {
unreachable!("Expr select item reached before arithmetic wiring")
}
}
}
if let Some(pred) = &q.where_ {
collect_pred(pred, &mut out);
}
out
}
pub fn refwalk_has_length_tail(q: &crate::query::ast::Query) -> bool {
use crate::query::ast::{Attr, Expr, Predicate, SelectItem};
fn attr_has(a: &Attr) -> bool {
match a {
Attr::RefPath { tail, .. } => matches!(tail.as_ref(), Attr::Length) || attr_has(tail),
Attr::ToHex(inner) => expr_has(inner),
_ => false,
}
}
fn expr_has(e: &Expr) -> bool {
match e {
Expr::Attr(a) => attr_has(a),
Expr::Lit(_) => false,
Expr::Binary { lhs, rhs, .. } => expr_has(lhs) || expr_has(rhs),
Expr::Unary { arg, .. } => expr_has(arg),
Expr::Method { receiver, args, .. } => expr_has(receiver) || args.iter().any(expr_has),
Expr::Aggregate { .. } => false,
Expr::Case { branches, else_ } => {
branches
.iter()
.any(|(pred, then_e)| pred_has(pred) || expr_has(then_e))
|| else_.as_ref().is_some_and(|e| expr_has(e))
}
Expr::Coalesce(args) => args.iter().any(expr_has),
Expr::NullIf { lhs, rhs } => expr_has(lhs) || expr_has(rhs),
}
}
fn pred_has(p: &Predicate) -> bool {
match p {
Predicate::And(a, b) | Predicate::Or(a, b) => pred_has(a) || pred_has(b),
Predicate::Not(a) => pred_has(a),
Predicate::Compare { lhs, rhs, .. } => expr_has(lhs) || expr_has(rhs),
Predicate::InSubquery { .. } | Predicate::InstanceOf(_) => false,
Predicate::Exists { .. } => false,
}
}
let select_has = q.select.iter().any(|item| match item {
SelectItem::Attr(a) => attr_has(a),
SelectItem::Aggregate { arg, .. } => {
matches!(arg.as_ref(), SelectItem::Attr(a) if attr_has(a))
}
SelectItem::Expr(e) => expr_has(e),
_ => false,
});
select_has || q.where_.as_ref().is_some_and(pred_has)
}
pub fn refwalk_has_address_tail(q: &crate::query::ast::Query) -> bool {
use crate::query::ast::{Attr, Expr, Predicate, SelectItem};
fn attr_has(a: &Attr) -> bool {
match a {
Attr::RefPath { tail, .. } => {
matches!(tail.as_ref(), Attr::ObjectAddress) || attr_has(tail)
}
Attr::ToHex(inner) => expr_has(inner),
_ => false,
}
}
fn expr_has(e: &Expr) -> bool {
match e {
Expr::Attr(a) => attr_has(a),
Expr::Lit(_) => false,
Expr::Binary { lhs, rhs, .. } => expr_has(lhs) || expr_has(rhs),
Expr::Unary { arg, .. } => expr_has(arg),
Expr::Method { receiver, args, .. } => expr_has(receiver) || args.iter().any(expr_has),
Expr::Aggregate { .. } => false,
Expr::Case { branches, else_ } => {
branches
.iter()
.any(|(pred, then_e)| pred_has(pred) || expr_has(then_e))
|| else_.as_ref().is_some_and(|e| expr_has(e))
}
Expr::Coalesce(args) => args.iter().any(expr_has),
Expr::NullIf { lhs, rhs } => expr_has(lhs) || expr_has(rhs),
}
}
fn pred_has(p: &Predicate) -> bool {
match p {
Predicate::And(a, b) | Predicate::Or(a, b) => pred_has(a) || pred_has(b),
Predicate::Not(a) => pred_has(a),
Predicate::Compare { lhs, rhs, .. } => expr_has(lhs) || expr_has(rhs),
Predicate::InSubquery { .. } | Predicate::InstanceOf(_) => false,
Predicate::Exists { .. } => false,
}
}
let select_has = q.select.iter().any(|item| match item {
SelectItem::Attr(a) => attr_has(a),
SelectItem::Aggregate { arg, .. } => {
matches!(arg.as_ref(), SelectItem::Attr(a) if attr_has(a))
}
SelectItem::Expr(e) => expr_has(e),
_ => false,
});
select_has || q.where_.as_ref().is_some_and(pred_has)
}
pub struct RefWalkCsr {
pub fwd_off: Vec<u32>,
pub fwd_tgt: Vec<u32>,
pub fwd_field: Vec<u32>,
pub field_names: Vec<String>,
pub tails: std::collections::HashMap<u32, crate::query::model::QueryValue>,
pub truncated: bool,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn intern_dedups_hop_fields_across_queries() {
let names =
intern_hop_fields(&[vec!["parent".into()], vec!["parent".into(), "next".into()]]);
assert_eq!(names, vec!["parent".to_string(), "next".to_string()]);
}
#[test]
fn intern_preserves_first_seen_order() {
let names = intern_hop_fields(&[
vec!["b".into(), "a".into()],
vec!["a".into(), "c".into(), "b".into()],
]);
assert_eq!(
names,
vec!["b".to_string(), "a".to_string(), "c".to_string()]
);
}
#[test]
fn edges_into_csr_sorts_by_src_and_offsets() {
let mut e = RefWalkEdges::new(100);
e.push(2, 0, 9);
e.push(0, 0, 5);
e.push(0, 1, 7);
let (off, tgt, fid) = e.into_csr(3);
assert_eq!(off, vec![0, 2, 2, 3]);
assert_eq!(tgt, vec![5, 7, 9]);
assert_eq!(fid, vec![0, 1, 0]);
}
#[test]
fn edges_cap_sets_truncated() {
let mut e = RefWalkEdges::new(1);
e.push(0, 0, 1);
assert!(!e.truncated());
e.push(0, 0, 2);
assert!(e.truncated());
assert_eq!(e.len(), 1);
}
#[test]
fn empty_edges_into_csr_all_zero_offsets() {
let e = RefWalkEdges::new(10);
assert!(e.is_empty());
let (off, tgt, fid) = e.into_csr(4);
assert_eq!(off, vec![0, 0, 0, 0, 0]);
assert!(tgt.is_empty());
assert!(fid.is_empty());
}
#[test]
fn edge_on_last_node_boundary() {
let mut e = RefWalkEdges::new(10);
e.push(3, 0, 42);
let (off, tgt, fid) = e.into_csr(4);
assert_eq!(off, vec![0, 0, 0, 0, 1]);
assert_eq!(tgt, vec![42]);
assert_eq!(fid, vec![0]);
}
#[test]
fn multiple_fields_on_one_src_preserve_dst_pairing() {
let mut e = RefWalkEdges::new(10);
e.push(0, 0, 100);
e.push(0, 2, 200);
e.push(0, 1, 300);
let (off, tgt, fid) = e.into_csr(1);
assert_eq!(off, vec![0, 3]);
assert_eq!(fid, vec![0, 2, 1]);
assert_eq!(tgt, vec![100, 200, 300]);
}
#[test]
fn refwalk_field_names_gathers_select_and_where_hops() {
let q = crate::query::parse::parse("SELECT x.parent.name FROM C x WHERE x.next.hash > 0")
.unwrap();
let names = refwalk_field_names(&q);
assert!(names.contains(&"parent".to_string()));
assert!(names.contains(&"next".to_string()));
assert!(!names.contains(&"name".to_string()));
assert!(!names.contains(&"hash".to_string()));
}
#[test]
fn refwalk_field_names_empty_when_no_refpath() {
let q = crate::query::parse::parse("SELECT x.count FROM C x").unwrap();
assert!(refwalk_field_names(&q).is_empty());
}
#[test]
fn refwalk_tail_field_names_gathers_field_tails() {
let q = crate::query::parse::parse("SELECT x.parent.name FROM C x WHERE x.next.hash > 0")
.unwrap();
let tails = refwalk_tail_field_names(&q);
assert!(tails.contains(&"name".to_string()));
assert!(tails.contains(&"hash".to_string()));
assert!(!tails.contains(&"parent".to_string()));
assert!(!tails.contains(&"next".to_string()));
}
#[test]
fn refwalk_has_length_tail_detects_select_and_where() {
let q =
crate::query::parse::parse("SELECT s.value.@length FROM java.lang.String s").unwrap();
assert!(refwalk_has_length_tail(&q));
let q = crate::query::parse::parse(
"SELECT s FROM java.lang.String s WHERE s.value.@length > 3",
)
.unwrap();
assert!(refwalk_has_length_tail(&q));
let q = crate::query::parse::parse("SELECT x.parent.name FROM C x").unwrap();
assert!(!refwalk_has_length_tail(&q));
let q = crate::query::parse::parse("SELECT @length FROM char[]").unwrap();
assert!(!refwalk_has_length_tail(&q));
}
#[test]
fn refwalk_has_length_tail_detects_rhs_and_wrapped() {
let q = crate::query::parse::parse(
"SELECT s FROM java.lang.String s WHERE 3 < s.value.@length",
)
.unwrap();
assert!(
refwalk_has_length_tail(&q),
"RHS @length tail must arm capture"
);
let q = crate::query::parse::parse(
"SELECT s FROM java.lang.String s WHERE s.value.@length + 1 > 4",
)
.unwrap();
assert!(
refwalk_has_length_tail(&q),
"wrapped @length tail must arm capture"
);
let q = crate::query::parse::parse("SELECT s.value.@length + 1 FROM java.lang.String s")
.unwrap();
assert!(
refwalk_has_length_tail(&q),
"SELECT-expr @length tail must arm capture"
);
}
#[test]
fn refwalk_has_address_tail_detects_rhs_and_wrapped() {
let q = crate::query::parse::parse(
"SELECT s FROM java.util.HashMap$Node s WHERE 0 < s.key.@objectAddress",
)
.unwrap();
assert!(
refwalk_has_address_tail(&q),
"RHS @objectAddress tail must arm capture"
);
let q = crate::query::parse::parse(
"SELECT s.value.@objectAddress + 0 FROM java.util.HashMap$Node s",
)
.unwrap();
assert!(
refwalk_has_address_tail(&q),
"SELECT-expr @objectAddress tail must arm capture"
);
}
#[test]
fn refwalk_tails_capping_and_lookup() {
use crate::query::model::QueryValue;
let mut t = RefWalkTails::new(1);
t.insert(3, QueryValue::Int(42));
assert!(!t.truncated());
assert_eq!(t.get(3), Some(&QueryValue::Int(42)));
t.insert(4, QueryValue::Int(99));
assert!(t.truncated());
assert_eq!(t.len(), 1);
assert_eq!(t.get(4), None);
t.insert(3, QueryValue::Int(7));
assert_eq!(t.get(3), Some(&QueryValue::Int(7)));
}
}