partial eval concat

src/algebra/parser.rs

@@ -1,17 +1,17 @@
 use crate::data::eval::{EvalError, PartialEvalContext};
 use crate::data::expr::{Expr, StaticExpr};
-use crate::data::parser::{ExprParseError, parse_scoped_dict};
+use crate::data::parser::{parse_scoped_dict, ExprParseError};
 use crate::data::tuple::{DataKind, OwnTuple};
 use crate::data::tuple_set::{BindingMap, BindingMapEvalContext, TableId, TupleSet, TupleSetIdx};
 use crate::data::value::{StaticValue, Value};
 use crate::ddl::reify::{AssocInfo, EdgeInfo, IndexInfo, TableInfo};
 use crate::parser::{CozoParser, Pair, Pairs, Rule};
 use crate::runtime::session::Definable;
+use pest::error::Error;
+use pest::Parser;
 use std::collections::BTreeMap;
 use std::result;
 use std::sync::Arc;
-use pest::error::Error;
-use pest::Parser;
 
 #[derive(thiserror::Error, Debug)]
 pub(crate) enum AlgebraParseError {
@@ -43,7 +43,7 @@ pub(crate) enum AlgebraParseError {
     ValueError(StaticValue),
 
     #[error("Parse error {0}")]
-    Parse(String)
+    Parse(String),
 }
 
 impl From<pest::error::Error<Rule>> for AlgebraParseError {
@@ -92,7 +92,7 @@ impl InterpretContext for () {
 pub(crate) trait RelationalAlgebra {
     fn name(&self) -> &str;
     fn binding_map(&self) -> &BindingMap;
-    fn iter<'a>(&'a self) -> Box<dyn Iterator<Item=TupleSet> + 'a>;
+    fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = TupleSet> + 'a>;
 }
 
 const NAME_RA_FROM_VALUES: &str = "Values";
@@ -192,7 +192,7 @@ impl RelationalAlgebra for RaFromValues {
         &self.binding
     }
 
-    fn iter<'a>(&'a self) -> Box<dyn Iterator<Item=TupleSet> + 'a> {
+    fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = TupleSet> + 'a> {
         let it = self.values.iter().map(|vs| {
             let mut tuple = OwnTuple::with_data_prefix(DataKind::Data);
             for v in vs {
@@ -212,7 +212,7 @@ struct RaInsert {
     source: Box<dyn RelationalAlgebra>,
     target: TableInfo,
     binding_map: BindingMap,
-    conversion_map: BTreeMap<String, StaticExpr>
+    conversion_map: BTreeMap<String, StaticExpr>,
 }
 
 impl RaInsert {
@@ -221,23 +221,36 @@ impl RaInsert {
         prev: Option<Arc<dyn RelationalAlgebra>>,
         mut args: Pairs,
     ) -> Result<Self> {
-        let not_enough_args =
-            || AlgebraParseError::NotEnoughArguments(NAME_RA_INSERT.to_string());
+        let not_enough_args = || AlgebraParseError::NotEnoughArguments(NAME_RA_INSERT.to_string());
         let source = match prev {
             Some(v) => v,
-            None => build_ra_expr(ctx, args.next().ok_or_else(not_enough_args)?)?
+            None => build_ra_expr(ctx, args.next().ok_or_else(not_enough_args)?)?,
         };
         let table_name = args.next().ok_or_else(not_enough_args)?;
-        let table_name = CozoParser::parse(Rule::ident_all, table_name.as_str())?.next().unwrap().as_str();
-        let pair = args.next().ok_or_else(not_enough_args)?.into_inner().next().unwrap();
+        let table_name = CozoParser::parse(Rule::ident_all, table_name.as_str())?
+            .next()
+            .unwrap()
+            .as_str();
+        let pair = args
+            .next()
+            .ok_or_else(not_enough_args)?
+            .into_inner()
+            .next()
+            .unwrap();
         assert_rule(&pair, Rule::scoped_dict, NAME_RA_INSERT, 2)?;
         let (binding, keys, vals) = parse_scoped_dict(pair)?;
         let source_map = source.binding_map();
-        let binding_ctx = BindingMapEvalContext { map: source_map, parent: ctx };
-        let keys = keys.into_iter().map(|(k, v)| -> Result<(String, Expr)> {
-            let v = v.partial_eval(&binding_ctx)?;
-            Ok((k, v))
-        }).collect::<Result<BTreeMap<_, _>>>()?;
+        let binding_ctx = BindingMapEvalContext {
+            map: source_map,
+            parent: ctx,
+        };
+        let keys = keys
+            .into_iter()
+            .map(|(k, v)| -> Result<(String, Expr)> {
+                let v = v.partial_eval(&binding_ctx)?;
+                Ok((k, v))
+            })
+            .collect::<Result<BTreeMap<_, _>>>()?;
         dbg!(&vals);
         let vals = vals.partial_eval(&binding_ctx)?;
         dbg!(&vals, &binding, &keys, &table_name);
@@ -254,7 +267,7 @@ impl RelationalAlgebra for RaInsert {
         &self.binding_map
     }
 
-    fn iter<'a>(&'a self) -> Box<dyn Iterator<Item=TupleSet> + 'a> {
+    fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = TupleSet> + 'a> {
         self.source.iter()
     }
 }
@@ -267,9 +280,7 @@ pub(crate) fn build_ra_expr(
     for pair in pair.into_inner() {
         let mut pairs = pair.into_inner();
         match pairs.next().unwrap().as_str() {
-            NAME_RA_INSERT => {
-                built = Some(Arc::new(RaInsert::build(ctx, built, pairs)?))
-            },
+            NAME_RA_INSERT => built = Some(Arc::new(RaInsert::build(ctx, built, pairs)?)),
             NAME_RA_FROM_VALUES => {
                 built = Some(Arc::new(RaFromValues::build(ctx, built, pairs)?));
             }

src/data/op/combine.rs

@@ -1,11 +1,11 @@
+use crate::data::eval::EvalError::OpTypeMismatch;
 use crate::data::eval::{EvalError, PartialEvalContext};
+use crate::data::expr::Expr;
 use crate::data::op::Op;
 use crate::data::value::Value;
 use std::collections::BTreeMap;
 use std::result;
 use std::sync::Arc;
-use crate::data::eval::EvalError::OpTypeMismatch;
-use crate::data::expr::Expr;
 
 type Result<T> = result::Result<T, EvalError>;
 
@@ -46,12 +46,58 @@ impl Op for OpConcat {
     }
 }
 
-
 pub(crate) fn partial_eval_concat_expr<'a, T: PartialEvalContext>(
     ctx: &'a T,
     args: Vec<Expr<'a>>,
 ) -> Result<Expr<'a>> {
-    todo!()
+    let mut can_concat = true;
+    let mut all_const = true;
+    let args = args
+        .into_iter()
+        .map(|a| -> Result<Expr> {
+            let a = a.partial_eval(ctx)?;
+            all_const = all_const && a.is_const();
+            can_concat = all_const && (a.is_const() || matches!(a, Expr::List(_)));
+            Ok(a)
+        })
+        .collect::<Result<Vec<_>>>()?;
+    if all_const {
+        let mut result = vec![];
+        for arg in args.into_iter() {
+            match arg.extract_const().unwrap() {
+                Value::List(l) => result.extend(l),
+                v => {
+                    return Err(OpTypeMismatch(
+                        NAME_OP_CONCAT.to_string(),
+                        vec![v.to_static()],
+                    ))
+                }
+            }
+        }
+        Ok(Expr::Const(Value::List(result)))
+    } else if can_concat {
+        let mut result = vec![];
+        for arg in args.into_iter() {
+            match arg {
+                Expr::Const(Value::Null) => {}
+                Expr::Const(Value::List(l)) => {
+                    for a in l {
+                        result.push(Expr::Const(a))
+                    }
+                }
+                Expr::List(l) => result.extend(l),
+                v => {
+                    return Err(OpTypeMismatch(
+                        NAME_OP_CONCAT.to_string(),
+                        vec![Value::from(v).to_static()],
+                    ))
+                }
+            }
+        }
+        Ok(Expr::List(result))
+    } else {
+        Ok(Expr::Apply(Arc::new(OpConcat), args))
+    }
 }
 
 pub(crate) struct OpMerge;
@@ -91,26 +137,33 @@ impl Op for OpMerge {
     }
 }
 
-
 pub(crate) fn partial_eval_merge_expr<'a, T: PartialEvalContext>(
     ctx: &'a T,
     args: Vec<Expr<'a>>,
 ) -> Result<Expr<'a>> {
     let mut can_merge = true;
     let mut all_const = true;
-    let args = args.into_iter().map(|ex| -> Result<Expr> {
-        let ex = ex.partial_eval(ctx)?;
-        all_const = all_const && ex.is_const();
-        can_merge = can_merge && ex.is_const() || matches!(ex, Expr::Dict(_));
-        Ok(ex)
-    }).collect::<Result<Vec<_>>>()?;
+    let args = args
+        .into_iter()
+        .map(|ex| -> Result<Expr> {
+            let ex = ex.partial_eval(ctx)?;
+            all_const = all_const && ex.is_const();
+            can_merge = can_merge && (ex.is_const() || matches!(ex, Expr::Dict(_)));
+            Ok(ex)
+        })
+        .collect::<Result<Vec<_>>>()?;
     if all_const {
         let mut result = BTreeMap::new();
         for arg in args.into_iter() {
             match arg.extract_const().unwrap() {
                 Value::Null => {}
                 Value::Dict(d) => result.extend(d),
-                v => return Err(OpTypeMismatch(NAME_OP_MERGE.to_string(), vec![v.to_static()]))
+                v => {
+                    return Err(OpTypeMismatch(
+                        NAME_OP_MERGE.to_string(),
+                        vec![v.to_static()],
+                    ))
+                }
             }
         }
         Ok(Expr::Const(Value::Dict(result)))
@@ -124,15 +177,17 @@ pub(crate) fn partial_eval_merge_expr<'a, T: PartialEvalContext>(
                         result.insert(k.to_string(), Expr::Const(v));
                     }
                 }
-                Expr::Dict(d) => {
-                    result.extend(d)
+                Expr::Dict(d) => result.extend(d),
+                v => {
+                    return Err(OpTypeMismatch(
+                        NAME_OP_MERGE.to_string(),
+                        vec![Value::from(v).to_static()],
+                    ))
                 }
-                v => return Err(OpTypeMismatch(NAME_OP_MERGE.to_string(),
-                                               vec![Value::from(v).to_static()]))
             }
         }
         Ok(Expr::Dict(result))
     } else {
         Ok(Expr::Apply(Arc::new(OpMerge), args))
     }
-}
+}

src/data/parser.rs

@@ -327,7 +327,6 @@ fn build_expr_infix<'a>(
     Ok(Expr::Apply(op, vec![lhs, rhs]))
 }
 
-
 pub(crate) fn parse_scoped_dict(pair: Pair) -> Result<(String, BTreeMap<String, Expr>, Expr)> {
     let mut pairs = pair.into_inner();
     let binding = pairs.next().unwrap().as_str().to_string();
@@ -351,21 +350,20 @@ pub(crate) fn parse_scoped_dict(pair: Pair) -> Result<(String, BTreeMap<String,
             }
             Rule::scoped_accessor => {
                 let name = parse_string(p.into_inner().next().unwrap())?;
-                let val =
-                    Expr::FieldAcc(name.clone().into(), Expr::Variable("_".into()).into());
+                let val = Expr::FieldAcc(name.clone().into(), Expr::Variable("_".into()).into());
                 collected.insert(name.into(), val);
             }
             Rule::spreading => {
                 let el = p.into_inner().next().unwrap();
                 let to_concat = build_expr_primary(el)?;
                 if !matches!(
-                            to_concat,
-                            Expr::Dict(_)
-                                | Expr::Variable(_)
-                                | Expr::IdxAcc(_, _)
-                                | Expr::FieldAcc(_, _)
-                                | Expr::Apply(_, _)
-                        ) {
+                    to_concat,
+                    Expr::Dict(_)
+                        | Expr::Variable(_)
+                        | Expr::IdxAcc(_, _)
+                        | Expr::FieldAcc(_, _)
+                        | Expr::Apply(_, _)
+                ) {
                     return Err(ExprParseError::SpreadingError(format!("{:?}", to_concat)));
                 }
                 if !collected.is_empty() {
@@ -389,7 +387,6 @@ pub(crate) fn parse_scoped_dict(pair: Pair) -> Result<(String, BTreeMap<String,
     Ok((binding, keys, vals))
 }
 
-
 #[cfg(test)]
 pub(crate) mod tests {
     use super::*;