spiral-sdk 0.1.8

//! Bounded serialization of unbound Vortex expressions.
//!
//! An [`Expression`] is the contract's base64 protobuf text, so it travels unchanged through
//! HTTP and the C ABI. Decoding checks size, node and depth bounds before any recursive
//! function deserialization or type binding.

use std::fmt::Display;

use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use prost::Message as _;
use vortex_array::expr::proto::ExprSerializeProtoExt as _;
use vortex_array::expr::{Expression as VortexExpression, nested_case_when};
use vortex_array::proto::expr::Expr;
use vortex_array::scalar_fn::{ForeignScalarFnVTable, ScalarFnId};
use vortex_error::{VortexResult, vortex_ensure};
use vortex_session::VortexSession;

pub use super::api::{Expression, NamedExpression};
use super::error::ScandalError;

const MAX_NODES: usize = 4096;
const MAX_DEPTH: usize = 64;
const MAX_METADATA_BYTES: usize = 256 * 1024;
const MAX_ENCODED_BYTES: usize = 512 * 1024;

/// Build a `CASE` expression from `(condition, value)` branches and an optional `ELSE`.
///
/// Returns an error if there are no branches or their child count exceeds the wire format's limit.
pub fn case_when(
    branches: Vec<(VortexExpression, VortexExpression)>,
    otherwise: Option<VortexExpression>,
) -> VortexResult<VortexExpression> {
    vortex_ensure!(!branches.is_empty(), "CASE needs at least one WHEN branch");
    vortex_ensure!(
        branches.len() <= (u32::MAX / 2) as usize,
        "CASE has too many branches"
    );
    Ok(nested_case_when(branches, otherwise))
}

/// An unbound SPQL call as a Vortex expression node, for composing with built-in expressions.
///
/// Serialize the composed tree with [`Expression::new`]. The client does not link SPQL
/// implementations, so the node has no return type until the execution host binds it.
pub fn spql_call(
    id: &str,
    arguments: Vec<VortexExpression>,
) -> Result<VortexExpression, ScandalError> {
    let call =
        ForeignScalarFnVTable::make_scalar_fn(ScalarFnId::new(id), Vec::new(), arguments.len());
    VortexExpression::try_new(call, arguments).map_err(invalid)
}

impl Expression {
    /// Construct an unbound SPQL call without linking its private implementation.
    ///
    /// Pass literal arguments as literal expressions. The execution host resolves the canonical
    /// SPQL ID, selects the overload that accepts the argument types and binds it.
    pub fn call(id: impl Into<String>, arguments: Vec<Self>) -> Result<Self, ScandalError> {
        let children = arguments
            .iter()
            .map(Self::proto)
            .collect::<Result<Vec<_>, _>>()?;
        Self::from_proto(Expr {
            id: id.into(),
            metadata: Some(Vec::new()),
            children,
        })
    }

    /// The conjunction of this expression and `other`: rows both accept.
    pub fn and(&self, other: &Self) -> Result<Self, ScandalError> {
        let mut conjunction =
            vortex_array::expr::and(vortex_array::expr::lit(true), vortex_array::expr::lit(true))
                .serialize_proto()
                .map_err(invalid)?;
        conjunction.children = vec![self.proto()?, other.proto()?];
        Self::from_proto(conjunction)
    }

    /// Serialize a Vortex expression after checking its structural work bounds.
    pub fn new(expression: VortexExpression) -> Result<Self, ScandalError> {
        let mut pending = vec![(&expression, 1usize)];
        let mut nodes = 0;
        while let Some((node, depth)) = pending.pop() {
            nodes += 1;
            validate_shape(nodes, pending.len(), node.children().len(), depth)?;
            pending.extend(node.children().iter().map(|child| (child, depth + 1)));
        }
        let proto = expression.serialize_proto().map_err(invalid)?;
        Self::from_proto(proto)
    }

    /// Decode using the execution host's function and data-type registries.
    ///
    /// This does not coerce or bind types. Unknown function IDs fail unless the
    /// supplied session explicitly permits them. Callers must separately check
    /// execution capabilities and bind against the exact input schema.
    pub fn to_vortex(&self, session: &VortexSession) -> Result<VortexExpression, ScandalError> {
        VortexExpression::from_proto(&self.proto()?, session).map_err(invalid)
    }

    /// The unbound Vortex protobuf, for a binary transport.
    pub fn to_bytes(&self) -> Result<Vec<u8>, ScandalError> {
        Ok(self.proto()?.encode_to_vec())
    }

    /// Size of the encoded protobuf, without decoding it.
    pub fn encoded_len(&self) -> usize {
        let padding = self
            .0
            .bytes()
            .rev()
            .take_while(|byte| *byte == b'=')
            .count();
        (self.0.len() / 4 * 3).saturating_sub(padding)
    }

    /// Wrap Vortex protobuf bytes, rejecting unbounded expression work.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, ScandalError> {
        if bytes.len() > MAX_ENCODED_BYTES {
            return Err(invalid("expression protobuf exceeds 512 KiB"));
        }
        Self::from_proto(Expr::decode(bytes).map_err(invalid)?)
    }

    /// The bounded protobuf this text encodes.
    fn proto(&self) -> Result<Expr, ScandalError> {
        if self.0.len() > MAX_ENCODED_BYTES.div_ceil(3) * 4 {
            return Err(invalid("expression protobuf exceeds 512 KiB"));
        }
        let bytes = STANDARD.decode(&self.0).map_err(invalid)?;
        let proto = Expr::decode(bytes.as_slice()).map_err(invalid)?;
        check_bounds(&proto)?;
        Ok(proto)
    }

    fn from_proto(proto: Expr) -> Result<Self, ScandalError> {
        check_bounds(&proto)?;
        Ok(Self(STANDARD.encode(proto.encode_to_vec())))
    }
}

impl NamedExpression {
    /// Name the output of `expression`.
    pub fn new(name: impl Into<String>, expression: Expression) -> Self {
        Self {
            name: name.into(),
            expression,
        }
    }
}

fn check_bounds(proto: &Expr) -> Result<(), ScandalError> {
    let mut pending = vec![(proto, 1usize)];
    let mut nodes = 0;
    let mut metadata_bytes = 0usize;
    while let Some((node, depth)) = pending.pop() {
        nodes += 1;
        validate_shape(nodes, pending.len(), node.children.len(), depth)?;
        metadata_bytes = metadata_bytes
            .saturating_add(node.id.len())
            .saturating_add(node.metadata.as_ref().map_or(0, Vec::len));
        if metadata_bytes > MAX_METADATA_BYTES {
            return Err(invalid("expression metadata exceeds 256 KiB"));
        }
        pending.extend(node.children.iter().map(|child| (child, depth + 1)));
    }
    if proto.encoded_len() > MAX_ENCODED_BYTES {
        return Err(invalid("expression protobuf exceeds 512 KiB"));
    }
    Ok(())
}

fn validate_shape(
    nodes: usize,
    pending: usize,
    children: usize,
    depth: usize,
) -> Result<(), ScandalError> {
    if nodes.saturating_add(pending).saturating_add(children) > MAX_NODES || depth > MAX_DEPTH {
        return Err(invalid("expression exceeds 4096 nodes or 64 levels"));
    }
    Ok(())
}

fn invalid(error: impl Display) -> ScandalError {
    ScandalError::InvalidArgument {
        message: error.to_string(),
    }
}

#[cfg(test)]
mod tests {
    use vortex_array::array_session;
    use vortex_array::expr::{and, col, eq, gt, lit, not};
    use vortex_array::proto::expr::Expr;

    use super::{Expression, case_when, spql_call};

    #[test]
    fn case_when_round_trips_with_the_default_vortex_session() {
        let branches = vec![
            (eq(col("a"), lit(1i64)), lit("one")),
            (eq(col("a"), lit(2i64)), lit("two")),
        ];
        for otherwise in [None, Some(lit("other"))] {
            let original = case_when(branches.clone(), otherwise).unwrap();
            let expression = Expression::new(original.clone()).unwrap();
            assert_eq!(expression.to_vortex(&array_session()).unwrap(), original);
        }
        assert!(case_when(Vec::new(), Some(lit(1i64))).is_err());
    }

    #[test]
    fn a_conjunction_decodes_to_both_terms() {
        let left = eq(col("episode"), lit(7u64));
        let right = gt(col("log_time"), lit(100i64));
        let conjunction = Expression::new(left.clone())
            .unwrap()
            .and(&Expression::new(right.clone()).unwrap())
            .unwrap();
        assert_eq!(
            conjunction.to_vortex(&array_session()).unwrap(),
            and(left, right)
        );
    }

    #[test]
    fn vortex_expression_round_trips_and_rejects_unbounded_or_invalid_wire() {
        let original = eq(col("value"), lit(u64::MAX));
        let expression = Expression::new(original.clone()).unwrap();
        let decoded: Expression =
            serde_json::from_value(serde_json::to_value(&expression).unwrap()).unwrap();
        assert_eq!(decoded.to_vortex(&array_session()).unwrap(), original);
        assert_eq!(
            Expression::from_bytes(&expression.to_bytes().unwrap()).unwrap(),
            expression
        );
        assert!(Expression::from_bytes(&[0xff]).is_err());
        let unknown = Expression::from_proto(Expr {
            id: "unregistered.function".into(),
            ..Default::default()
        })
        .unwrap();
        assert!(unknown.to_vortex(&array_session()).is_err());
        assert!(
            Expression("not base64".into())
                .to_vortex(&array_session())
                .is_err()
        );
        let mut deep = lit(true);
        for _ in 0..64 {
            deep = not(deep);
        }
        assert!(Expression::new(deep).is_err());
        let wide = Expr {
            id: "vortex.root".into(),
            children: vec![Expr::default(); 4096],
            ..Default::default()
        };
        assert!(Expression::from_proto(wide).is_err());
        assert!(Expression::new(lit("x".repeat(256 * 1024))).is_err());
    }

    #[test]
    fn composed_spql_call_serializes_like_a_direct_call() {
        let composed = spql_call("spql.str.upper", vec![col("label")]).unwrap();
        let direct = Expression::call(
            "spql.str.upper",
            vec![Expression::new(col("label")).unwrap()],
        )
        .unwrap();
        assert_eq!(Expression::new(composed.clone()).unwrap(), direct);
        let nested = Expression::new(not(eq(composed, lit("X")))).unwrap();
        assert_eq!(nested.to_bytes().unwrap().len(), nested.encoded_len());
    }
}