vortex_array/expr/analysis/
immediate_access.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::FieldName;
5use vortex_dtype::StructFields;
6use vortex_error::VortexExpect;
7use vortex_utils::aliases::hash_set::HashSet;
8
9use crate::expr::Expression;
10use crate::expr::analysis::AnnotationFn;
11use crate::expr::analysis::Annotations;
12use crate::expr::descendent_annotations;
13use crate::expr::exprs::get_item::GetItem;
14use crate::expr::exprs::root::Root;
15use crate::expr::exprs::select::Select;
16
17pub type FieldAccesses<'a> = Annotations<'a, FieldName>;
18
19/// An [`AnnotationFn`] for annotating scope accesses.
20pub fn annotate_scope_access(scope: &StructFields) -> impl AnnotationFn<Annotation = FieldName> {
21    move |expr: &Expression| {
22        assert!(
23            !expr.is::<Select>(),
24            "cannot analyse select, simplify the expression"
25        );
26
27        if let Some(get_item) = expr.as_opt::<GetItem>() {
28            if get_item.child(0).is::<Root>() {
29                return vec![get_item.data().clone()];
30            }
31        } else if expr.is::<Root>() {
32            return scope.names().iter().cloned().collect();
33        }
34
35        vec![]
36    }
37}
38
39/// For all subexpressions in an expression, find the fields that are accessed directly from the
40/// scope, but not any fields in those fields
41/// e.g. scope = {a: {b: .., c: ..}, d: ..}, expr = root().a.b + root().d accesses {a,d} (not b).
42///
43/// Note: This is a very naive, but simple analysis to find the fields that are accessed directly on an
44/// identity node. This is combined to provide an over-approximation of the fields that are accessed
45/// by an expression.
46pub fn immediate_scope_accesses<'a>(
47    expr: &'a Expression,
48    scope: &'a StructFields,
49) -> FieldAccesses<'a> {
50    descendent_annotations(expr, annotate_scope_access(scope))
51}
52
53/// This returns the immediate scope_access (as explained `immediate_scope_accesses`) for `expr`.
54pub fn immediate_scope_access<'a>(
55    expr: &'a Expression,
56    scope: &'a StructFields,
57) -> HashSet<FieldName> {
58    immediate_scope_accesses(expr, scope)
59        .get(expr)
60        .vortex_expect("Expression missing from scope accesses, this is a internal bug")
61        .clone()
62}
vortex_array/expr/analysis/immediate_access.rs

vortex_array/expr/analysis/
immediate_access.rs
