immediate_access.rs

use itertools::Itertools;
use vortex_array::aliases::hash_map::HashMap;
use vortex_array::aliases::hash_set::HashSet;
use vortex_dtype::{FieldName, StructDType};
use vortex_error::VortexResult;

use crate::traversal::{Node, NodeVisitor, TraversalOrder};
use crate::{ExprRef, GetItem, Identity, Select};

pub(crate) type FieldAccesses<'a> = HashMap<&'a ExprRef, HashSet<FieldName>>;

/// For all subexpressions in an expression, find the fields that are accessed directly from the
/// scope, but not any fields in those fields
/// e.g. scope = {a: {b: .., c: ..}, d: ..}, expr = ident().a.b + ident().d accesses {a,d} (not b).
pub(crate) struct ImmediateIdentityAccessesAnalysis<'a> {
    sub_expressions: FieldAccesses<'a>,
    scope_dtype: &'a StructDType,
}

impl<'a> ImmediateIdentityAccessesAnalysis<'a> {
    fn new(scope_dtype: &'a StructDType) -> Self {
        Self {
            sub_expressions: HashMap::new(),
            scope_dtype,
        }
    }

    pub(crate) fn analyze(
        expr: &'a ExprRef,
        scope_dtype: &'a StructDType,
    ) -> VortexResult<FieldAccesses<'a>> {
        let mut analysis = Self::new(scope_dtype);
        expr.accept(&mut analysis)?;
        Ok(analysis.sub_expressions)
    }
}

// This is a very naive, but simple analysis to find the fields that are accessed directly on an
// identity node. This is combined to provide an over-approximation of the fields that are accessed
// by an expression.
impl<'a> NodeVisitor<'a> for ImmediateIdentityAccessesAnalysis<'a> {
    type NodeTy = ExprRef;

    fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
        assert!(
            !node.as_any().is::<Select>(),
            "cannot analyse select, simply the expression"
        );
        if let Some(get_item) = node.as_any().downcast_ref::<GetItem>() {
            if get_item
                .child()
                .as_any()
                .downcast_ref::<Identity>()
                .is_some()
            {
                self.sub_expressions
                    .insert(node, HashSet::from_iter(vec![get_item.field().clone()]));

                return Ok(TraversalOrder::Skip);
            }
        } else if node.as_any().downcast_ref::<Identity>().is_some() {
            let st_dtype = &self.scope_dtype;
            self.sub_expressions
                .insert(node, st_dtype.names().iter().cloned().collect());
        }

        Ok(TraversalOrder::Continue)
    }

    fn visit_up(&mut self, node: &'a ExprRef) -> VortexResult<TraversalOrder> {
        let accesses = node
            .children()
            .iter()
            .filter_map(|c| self.sub_expressions.get(c).cloned())
            .collect_vec();

        let node_accesses = self.sub_expressions.entry(node).or_default();
        accesses
            .into_iter()
            .for_each(|fields| node_accesses.extend(fields.iter().cloned()));

        Ok(TraversalOrder::Continue)
    }
}