Add scope variables to expression scope (#3672)

Part of #3671 

Adds a new way of defining scope variables (before we deprecate the old
method) whereby variables are keyed by type. This is typical for e.g.
attaching context to an HTTP request using "middleware", which is
essentially what we're doing here during expression evaluation.

---------

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

Author: gatesn

vortex-expr/src/lib.rs

@@ -15,6 +15,7 @@ pub mod pruning;
 #[cfg(feature = "proto")]
 mod registry;
 mod scope;
+mod scope_vars;
 pub mod transform;
 pub mod traversal;
 

vortex-expr/src/scope.rs

@@ -1,5 +1,5 @@
-use std::any::Any;
-use std::fmt::Display;
+use std::any::{Any, TypeId};
+use std::fmt::{Debug, Display};
 use std::str::FromStr;
 use std::sync::Arc;
 
@@ -9,6 +9,8 @@ use vortex_dtype::{DType, FieldPathSet};
 use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
 use vortex_utils::aliases::hash_map::HashMap;
 
+use crate::scope_vars::{ScopeVar, ScopeVars};
+
 type ExprScope<T> = HashMap<Identifier, T>;
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
@@ -88,6 +90,8 @@ pub struct Scope {
     /// A map identifiers to opaque values used by expressions, but
     /// cannot affect the result type/shape.
     vars: ExprScope<Arc<dyn Any + Send + Sync>>,
+    /// Variables that can be set on the scope during expression evaluation.
+    scope_vars: ScopeVars,
 }
 
 pub type ScopeElement = (Identifier, ArrayRef);
@@ -156,6 +160,26 @@ impl Scope {
         self
     }
 
+    /// Returns a new evaluation scope with the given variable applied.
+    pub fn with_scope_var<V: ScopeVar>(mut self, var: V) -> Self {
+        self.scope_vars.insert(TypeId::of::<V>(), Box::new(var));
+        self
+    }
+
+    /// Returns the scope variable of type `V` if it exists.
+    pub fn scope_var<V: ScopeVar>(&self) -> Option<&V> {
+        self.scope_vars
+            .get(&TypeId::of::<V>())
+            .and_then(|boxed| (**boxed).as_any().downcast_ref::<V>())
+    }
+
+    /// Returns the mutable scope variable of type `V` if it exists.
+    pub fn scope_var_mut<V: ScopeVar>(&mut self) -> Option<&mut V> {
+        self.scope_vars
+            .get_mut(&TypeId::of::<V>())
+            .and_then(|boxed| (**boxed).as_any_mut().downcast_mut::<V>())
+    }
+
     pub fn iter(&self) -> impl Iterator<Item = (&Identifier, &ArrayRef)> {
         let values = self.arrays.iter();
 
@@ -285,3 +309,26 @@ impl ScopeFieldPathSet {
         self.with_set(ident, set)
     }
 }
+
+#[cfg(test)]
+mod test {
+    #[test]
+    fn test_scope_var() {
+        use super::*;
+
+        #[derive(Clone, PartialEq, Eq, Debug)]
+        struct TestVar {
+            value: i32,
+        }
+
+        let scope = Scope::empty(100);
+        assert!(scope.scope_var::<TestVar>().is_none());
+
+        let var = TestVar { value: 42 };
+        let mut scope = scope.with_scope_var(var.clone());
+        assert_eq!(scope.scope_var::<TestVar>(), Some(&var));
+
+        scope.scope_var_mut::<TestVar>().unwrap().value = 43;
+        assert_eq!(scope.scope_var::<TestVar>(), Some(&TestVar { value: 43 }));
+    }
+}

vortex-expr/src/scope_vars.rs

@@ -0,0 +1,56 @@
+use std::any::{Any, TypeId};
+use std::hash::{BuildHasherDefault, Hasher};
+
+use vortex_utils::aliases::hash_map::HashMap;
+
+/// A TypeMap based on `https://docs.rs/http/1.2.0/src/http/extensions.rs.html#41-266`.
+pub(crate) type ScopeVars = HashMap<TypeId, Box<dyn ScopeVar>, BuildHasherDefault<IdHasher>>;
+
+/// With TypeIds as keys, there's no need to hash them. They are already hashes
+/// themselves, coming from the compiler. The IdHasher just holds the u64 of
+/// the TypeId, and then returns it, instead of doing any bit fiddling.
+#[derive(Default)]
+pub(super) struct IdHasher(u64);
+
+impl Hasher for IdHasher {
+    #[inline]
+    fn finish(&self) -> u64 {
+        self.0
+    }
+
+    fn write(&mut self, _: &[u8]) {
+        unreachable!("TypeId calls write_u64");
+    }
+
+    #[inline]
+    fn write_u64(&mut self, id: u64) {
+        self.0 = id;
+    }
+}
+
+/// A trait for scope variables that can be stored in a `ScopeVars` map.
+pub trait ScopeVar: Any + Send + Sync {
+    fn as_any(&self) -> &dyn Any;
+    fn as_any_mut(&mut self) -> &mut dyn Any;
+    fn clone_box(&self) -> Box<dyn ScopeVar>;
+}
+
+impl Clone for Box<dyn ScopeVar> {
+    fn clone(&self) -> Self {
+        (**self).clone_box()
+    }
+}
+
+impl<T: Clone + Send + Sync + 'static> ScopeVar for T {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        self
+    }
+
+    fn clone_box(&self) -> Box<dyn ScopeVar> {
+        Box::new(self.clone())
+    }
+}