Rename accesses to annotations (#3803)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-expr/src/transform/access_analysis.rs

@@ -0,0 +1,95 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::hash::Hash;
+
+use vortex_error::{VortexExpect, VortexResult};
+use vortex_utils::aliases::hash_map::HashMap;
+use vortex_utils::aliases::hash_set::HashSet;
+
+use crate::traversal::{Node, NodeVisitor, TraversalOrder};
+use crate::{ExprRef, Identifier, VarVTable};
+
+pub trait Annotation: Clone + Hash + Eq {}
+
+impl<A> Annotation for A where A: Clone + Hash + Eq {}
+
+pub trait AnnotationFn: Fn(&ExprRef) -> Vec<Self::Annotation> {
+    type Annotation: Annotation;
+}
+
+impl<A, F> AnnotationFn for F
+where
+    A: Annotation,
+    F: Fn(&ExprRef) -> Vec<A>,
+{
+    type Annotation = A;
+}
+
+pub type Annotations<'a, A> = HashMap<&'a ExprRef, HashSet<A>>;
+
+/// Walk the expression tree and annotate each expression with zero or more annotations.
+///
+/// Returns a map of each expression to all annotations that any of its descendent (child)
+/// expressions are annotated with.
+pub fn descendent_annotations<A: AnnotationFn>(
+    expr: &ExprRef,
+    annotate: A,
+) -> Annotations<A::Annotation> {
+    let mut visitor = AnnotationVisitor {
+        annotations: Default::default(),
+        annotate,
+    };
+    expr.accept(&mut visitor).vortex_expect("Infallible");
+    visitor.annotations
+}
+
+pub fn variable_scope_annotations<A: Annotation>(
+    expr: &ExprRef,
+    f: impl Fn(&Identifier) -> A,
+) -> Annotations<A> {
+    descendent_annotations(expr, move |node| {
+        if let Some(variable) = node.as_opt::<VarVTable>() {
+            return vec![f(variable.var())];
+        }
+        vec![]
+    })
+}
+
+struct AnnotationVisitor<'a, A: AnnotationFn> {
+    annotations: Annotations<'a, A::Annotation>,
+    annotate: A,
+}
+
+impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> {
+    type NodeTy = ExprRef;
+
+    fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
+        let annotations = (self.annotate)(node);
+        if annotations.is_empty() {
+            // If the annotate fn returns empty, we do not annotate this node.
+            Ok(TraversalOrder::Continue)
+        } else {
+            self.annotations
+                .entry(node)
+                .or_default()
+                .extend(annotations);
+            Ok(TraversalOrder::Skip)
+        }
+    }
+
+    fn visit_up(&mut self, node: &'a ExprRef) -> VortexResult<TraversalOrder> {
+        let child_annotations = node
+            .children()
+            .iter()
+            .filter_map(|c| self.annotations.get(c).cloned())
+            .collect::<Vec<_>>();
+
+        let annotations = self.annotations.entry(node).or_default();
+        child_annotations
+            .into_iter()
+            .for_each(|ps| annotations.extend(ps.iter().cloned()));
+
+        Ok(TraversalOrder::Continue)
+    }
+}

vortex-expr/src/transform/annotations.rs

@@ -2,15 +2,33 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_dtype::{FieldName, StructFields};
-use vortex_error::{VortexResult, vortex_err};
-use vortex_utils::aliases::hash_map::HashMap;
+use vortex_error::VortexExpect;
 use vortex_utils::aliases::hash_set::HashSet;
 
-use crate::transform::access_analysis::AccessesAnalysis;
-use crate::traversal::TraversalOrder;
+use crate::transform::annotations::{AnnotationFn, Annotations, descendent_annotations};
 use crate::{ExprRef, GetItemVTable, SelectVTable, is_root};
 
-pub type FieldAccesses<'a> = HashMap<&'a ExprRef, HashSet<FieldName>>;
+pub type FieldAccesses<'a> = Annotations<'a, FieldName>;
+
+/// An [`AnnotationFn`] for annotating scope accesses.
+pub fn annotate_scope_access(scope: &StructFields) -> impl AnnotationFn<Annotation = FieldName> {
+    move |expr: &ExprRef| {
+        assert!(
+            !expr.is::<SelectVTable>(),
+            "cannot analyse select, simplify the expression"
+        );
+
+        if let Some(get_item) = expr.as_opt::<GetItemVTable>() {
+            if is_root(get_item.child()) {
+                return vec![get_item.field().clone()];
+            }
+        } else if is_root(expr) {
+            return scope.names().iter().cloned().collect();
+        }
+
+        vec![]
+    }
+}
 
 /// For all subexpressions in an expression, find the fields that are accessed directly from the
 /// scope, but not any fields in those fields
@@ -21,38 +39,18 @@ pub type FieldAccesses<'a> = HashMap<&'a ExprRef, HashSet<FieldName>>;
 /// by an expression.
 pub fn immediate_scope_accesses<'a>(
     expr: &'a ExprRef,
-    scope_dtype: &'a StructFields,
-) -> VortexResult<FieldAccesses<'a>> {
-    AccessesAnalysis::analyze(expr, move |node| {
-        assert!(
-            !node.is::<SelectVTable>(),
-            "cannot analyse select, simplify the expression"
-        );
-        if let Some(get_item) = node.as_opt::<GetItemVTable>() {
-            if is_root(get_item.child()) {
-                return (TraversalOrder::Skip, vec![get_item.field().clone()]);
-            }
-        } else if is_root(node) {
-            let st_dtype = &scope_dtype;
-            return (
-                TraversalOrder::Skip,
-                st_dtype.names().iter().cloned().collect(),
-            );
-        }
-
-        (TraversalOrder::Continue, vec![])
-    })
+    scope: &'a StructFields,
+) -> FieldAccesses<'a> {
+    descendent_annotations(expr, annotate_scope_access(scope))
 }
 
 /// This returns the immediate scope_access (as explained `immediate_scope_accesses`) for `expr`.
 pub fn immediate_scope_access<'a>(
     expr: &'a ExprRef,
-    scope_dtype: &'a StructFields,
-) -> VortexResult<HashSet<FieldName>> {
-    immediate_scope_accesses(expr, scope_dtype)?
+    scope: &'a StructFields,
+) -> HashSet<FieldName> {
+    immediate_scope_accesses(expr, scope)
         .get(expr)
-        .ok_or_else(|| {
-            vortex_err!("Expression missing from scope accesses, this is a internal bug")
-        })
-        .cloned()
+        .vortex_expect("Expression missing from scope accesses, this is a internal bug")
+        .clone()
 }

vortex-expr/src/transform/immediate_access.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 //! A collection of transformations that can be applied to a [`crate::ExprRef`].
-mod access_analysis;
+pub mod annotations;
 pub mod field_mask;
 pub mod immediate_access;
 pub(crate) mod match_between;