From 551371fd28f1eb1afe246d12135f268fb6633ca7 Mon Sep 17 00:00:00 2001 From: Vytautas Astrauskas Date: Sat, 11 Jan 2020 17:18:45 +0100 Subject: [PATCH] Implement support for transitive closures. --- .../tests/pass/transitive_closure.rs | 16 ++ src/generator_new/ast.rs | 95 ++++++-- src/generator_new/encode.rs | 206 ++++++++++++++---- src/generator_new/to_tokens.rs | 33 ++- 4 files changed, 284 insertions(+), 66 deletions(-) create mode 100644 datapond-derive/tests/pass/transitive_closure.rs diff --git a/datapond-derive/tests/pass/transitive_closure.rs b/datapond-derive/tests/pass/transitive_closure.rs new file mode 100644 index 0000000..ae2b3a8 --- /dev/null +++ b/datapond-derive/tests/pass/transitive_closure.rs @@ -0,0 +1,16 @@ +use datapond_derive::datapond; + +fn main() { + let inp = vec![(1, 2), (2, 3)]; + let out; + datapond! { + input inp(x: u32, y: u32) + output out(x: u32, y: u32) + out(x, y) :- inp(x, y). + out(x, y) :- out(x, z), out(z, y). + }; + assert!(out.len() == 3); + assert!(out[0] == (1, 2)); + assert!(out[1] == (1, 3)); + assert!(out[2] == (2, 3)); +} diff --git a/src/generator_new/ast.rs b/src/generator_new/ast.rs index fe3f440..38e91b8 100644 --- a/src/generator_new/ast.rs +++ b/src/generator_new/ast.rs @@ -42,7 +42,7 @@ use crate::data_structures::OrderedMap; /// ```ignore /// r_1.from_map(&r, |(x, z)| {(z, x)}); /// ``` -#[derive(Debug)] +#[derive(Debug, Clone)] pub(crate) struct DVar { pub name: syn::Ident, } @@ -59,6 +59,14 @@ pub(crate) struct DVarTuple { pub vars: Vec, } +impl DVarTuple { + pub fn new(args: Vec) -> Self { + Self { + vars: args.into_iter().map(|ident| DVar::new(ident)).collect(), + } + } +} + /// A (key, value) representation of `DVar`s. It is used for joins. #[derive(Debug)] pub(crate) struct DVarKeyVal { @@ -75,8 +83,12 @@ pub(crate) enum DVars { impl DVars { pub fn new_tuple(args: Vec) -> Self { - DVars::Tuple(DVarTuple { - vars: args.into_iter().map(|ident| DVar::new(ident)).collect(), + DVars::Tuple(DVarTuple::new(args)) + } + pub fn new_key_val(key: Vec, value: Vec) -> Self { + DVars::KeyVal(DVarKeyVal { + key: key.into_iter().map(|ident| DVar::new(ident)).collect(), + value: value.into_iter().map(|ident| DVar::new(ident)).collect(), }) } } @@ -125,14 +137,6 @@ pub(crate) struct Variable { } impl Variable { - pub fn with_suffix(&self, suffix: &str) -> Self { - Self { - name: syn::Ident::new( - &format!("{}{}", self.name, suffix), - proc_macro2::Span::call_site(), - ), - } - } pub fn with_counter(&self, counter: usize) -> Self { Self { name: syn::Ident::new( @@ -160,40 +164,40 @@ pub(crate) struct ReorderOp { #[derive(Debug)] pub(crate) struct BindVarOp { /// A variable into which we write the result. - output: Variable, + pub output: Variable, /// A variable from which we read the input. - input: Variable, + pub input: Variable, /// Input variables that are copied to output and potentially used for evaluating `expr`. - vars: DVarTuple, + pub vars: DVarTuple, /// The expression whose result is bound to a new variable. - expr: syn::Expr, + pub expr: syn::Expr, } /// An operation that joins two variables. #[derive(Debug)] pub(crate) struct JoinOp { /// A variable into which we write the result. - output: Variable, + pub output: Variable, /// The first variable, which we use in join. - input_first: Variable, + pub input_first: Variable, /// The second variable, which we use in join. - input_second: Variable, + pub input_second: Variable, /// Datalog variables used for joining. - key: Vec, + pub key: DVarTuple, /// Datalog value variables from the first variable. - value_first: Vec, + pub value_first: DVarTuple, /// Datalog value variables from the second variable. - value_second: Vec, + pub value_second: DVarTuple, } /// An operation that filters out facts. #[derive(Debug)] pub(crate) struct FilterOp { /// A variable which we want to filter. - variable: Variable, - vars: DVars, + pub variable: Variable, + pub vars: DVars, /// A boolean expression used for filtering. - expr: syn::Expr, + pub expr: syn::Expr, } /// An operation that inserts the relation into a variable. @@ -260,6 +264,16 @@ impl Iteration { })); new_variable } + /// Get Datafrog variable that corresponds to the given predicate name. If + /// we have only a relation, then convert it into a variable. + pub fn get_or_convert_variable(&mut self, predicate: &syn::Ident) -> Variable { + if let Some(variable) = self.get_relation_var(predicate) { + // TODO: Avoid converting the same relation multiple times. + self.convert_relation_to_variable(&variable) + } else { + self.get_variable(predicate) + } + } pub fn get_relation_var(&self, variable_name: &syn::Ident) -> Option { self.relations .get(variable_name) @@ -271,4 +285,37 @@ impl Iteration { pub fn add_operation(&mut self, operation: Operation) { self.body_operations.push(operation); } + pub fn get_variable_tuple_types(&self, variable: &Variable) -> Vec { + let decl = &self.variables[&variable.name]; + match &decl.typ { + DVarTypes::Tuple(types) => types.clone(), + DVarTypes::KeyVal { .. } => unreachable!(), + } + } + pub fn create_key_val_variable( + &mut self, + variable: &Variable, + key: Vec, + value: Vec, + ) -> Variable { + self.create_variable(variable, DVarTypes::KeyVal { key, value }) + } + pub fn create_tuple_variable( + &mut self, + variable: &Variable, + types: Vec, + ) -> Variable { + self.create_variable(variable, DVarTypes::Tuple(types)) + } + pub fn create_variable(&mut self, variable: &Variable, typ: DVarTypes) -> Variable { + let variable_decl = VariableDecl { + var: variable.with_counter(self.variables.len()), + typ: typ, + is_output: false, + }; + let new_variable = variable_decl.var.clone(); + self.variables + .insert(new_variable.name.clone(), variable_decl); + new_variable + } } diff --git a/src/generator_new/encode.rs b/src/generator_new/encode.rs index f60e3cd..2eba3c6 100644 --- a/src/generator_new/encode.rs +++ b/src/generator_new/encode.rs @@ -1,6 +1,60 @@ use crate::ast; use crate::generator_new::ast as gen; +/// Divide the arguments into three sets: +/// +/// 1. `key` – the arguments that are common in `first` and `second`. +/// 2. `first_remainder` – the arguments that are unique in `first`. +/// 3. `second_remainder` – the arguments that are unique in `second`. +fn common_args( + first: &Vec, + first_types: &Vec, + second: &Vec, + second_types: &Vec, +) -> ( + (Vec, Vec), + (Vec, Vec), + (Vec, Vec), +) { + assert!(first.len() == first_types.len()); + assert!(second.len() == second_types.len()); + + let mut key = Vec::new(); + let mut key_types = Vec::new(); + let mut first_remainder = Vec::new(); + let mut first_remainder_types = Vec::new(); + + for (arg1, arg1_type) in first.iter().zip(first_types) { + let mut found = false; + for arg2 in second { + if arg1 == arg2 { + key.push(arg1.clone()); + key_types.push(arg1_type.clone()); + found = true; + break; + } + } + if !found { + first_remainder.push(arg1.clone()); + first_remainder_types.push(arg1_type.clone()); + } + } + let mut second_remainder = Vec::new(); + let mut second_remainder_types = Vec::new(); + for (arg2, arg2_type) in second.iter().zip(second_types) { + if !key.contains(arg2) { + second_remainder.push(arg2.clone()); + second_remainder_types.push(arg2_type.clone()); + } + } + + ( + (key, key_types), + (first_remainder, first_remainder_types), + (second_remainder, second_remainder_types), + ) +} + pub(crate) fn encode(program: ast::Program) -> gen::Iteration { let mut relations = Vec::new(); let mut variables = Vec::new(); @@ -35,39 +89,69 @@ pub(crate) fn encode(program: ast::Program) -> gen::Iteration { } let mut iteration = gen::Iteration::new(relations, variables); for rule in &program.rules { + let head_variable = iteration.get_variable(&rule.head.predicate); let mut iter = rule.body.iter(); let literal1 = iter.next().unwrap(); if literal1.is_negated { unimplemented!(); } - let variable = if let Some(variable) = iteration.get_relation_var(&literal1.predicate) { - iteration.convert_relation_to_variable(&variable) - } else { - iteration.get_variable(&literal1.predicate) - }; - let args = literal1.args.clone(); - - // Retrieve the main variable for the head. - // Create a new variable `res` for the rule result. - // { - // if rule.body.len() == 1 { - // res.from_map(rule.body[0], ...) - // } else if rule.body.len() == 2 { - // let first = rule.body[0]; - // let second = rule.body[1]; - // let (key, first_remainder, second_remainder) = common_args(first, second); - // let new_first = variable::<(key, first_remainder)>; - // let new_second = variable::<(key, second_remainder)>; - // new_first.from_map(first); - // new_second.from_map(second); - // res.from_join(new_first, new_second, ...) - // } - // } + let mut variable = iteration.get_or_convert_variable(&literal1.predicate); + let mut args = literal1.args.clone(); - // - // Extend the main variable with the new variable. - // rule.head.extend(res); - let head_variable = iteration.get_variable(&rule.head.predicate); + while let Some(literal) = iter.next() { + let joined_variable = iteration.get_or_convert_variable(&literal.predicate); + // TODO: Check during the typechecking phase that no literal has two + // arguments with the same name. + let arg_types = iteration.get_variable_tuple_types(&variable); + let literal_arg_types = iteration.get_variable_tuple_types(&joined_variable); + let ((key, key_types), (remainder1, remainder1_types), (remainder2, remainder2_types)) = + common_args(&args, &arg_types, &literal.args, &literal_arg_types); + let first_variable = iteration.create_key_val_variable( + &variable, + key_types.clone(), + remainder1_types.clone(), + ); + let reorder_first_op = gen::ReorderOp { + output: first_variable.clone(), + input: variable, + input_vars: args.into(), + output_vars: (key.clone(), remainder1.clone()).into(), + }; + iteration.add_operation(gen::Operation::Reorder(reorder_first_op)); + let second_variable = iteration.create_key_val_variable( + &joined_variable, + key_types.clone(), + remainder2_types.clone(), + ); + let reorder_second_op = gen::ReorderOp { + output: second_variable.clone(), + input: joined_variable, + input_vars: literal.args.clone().into(), + output_vars: (key.clone(), remainder2.clone()).into(), + }; + iteration.add_operation(gen::Operation::Reorder(reorder_second_op)); + let result_types = key_types + .into_iter() + .chain(remainder1_types) + .chain(remainder2_types) + .collect(); + args = key + .clone() + .into_iter() + .chain(remainder1.clone()) + .chain(remainder2.clone()) + .collect(); + variable = iteration.create_tuple_variable(&head_variable, result_types); + let join_op = gen::JoinOp { + output: variable.clone(), + input_first: first_variable, + input_second: second_variable, + key: key.into(), + value_first: remainder1.into(), + value_second: remainder2.into(), + }; + iteration.add_operation(gen::Operation::Join(join_op)); + } let reorder_op = gen::ReorderOp { output: head_variable, input: variable, @@ -85,6 +169,21 @@ impl std::convert::From> for gen::DVars { } } +impl std::convert::From> for gen::DVarTuple { + fn from(args: Vec) -> Self { + gen::DVarTuple::new(args.into_iter().map(|arg| arg.to_ident()).collect()) + } +} + +impl std::convert::From<(Vec, Vec)> for gen::DVars { + fn from((key, value): (Vec, Vec)) -> Self { + gen::DVars::new_key_val( + key.into_iter().map(|arg| arg.to_ident()).collect(), + value.into_iter().map(|arg| arg.to_ident()).collect(), + ) + } +} + impl std::convert::From> for gen::DVars { fn from(args: Vec) -> Self { gen::DVars::new_tuple(args) @@ -100,20 +199,24 @@ mod tests { use quote::ToTokens; use std::str::FromStr; + fn compare(datalog_source: &str, exptected_encoding: &str) { + let parsed_program = parse(datalog_source); + let typechecked_program = typecheck(parsed_program).unwrap(); + let iteration = encode(typechecked_program); + let tokens = iteration.to_token_stream().to_string(); + eprintln!("{}", tokens); + let expected_tokens = TokenStream::from_str(exptected_encoding).unwrap(); + assert_eq!(tokens.to_string(), expected_tokens.to_string()); + } + #[test] fn encode_simple1() { - let program = parse( + compare( " input inp(x: u32, y: u32) output out(x: u32, y: u32) out(x, y) :- inp(y, x). ", - ); - let program = typecheck(program).unwrap(); - let iteration = encode(program); - let tokens = iteration.to_token_stream().to_string(); - eprintln!("{}", tokens); - let expected_tokens = TokenStream::from_str( r##" { let mut iteration = datafrog::Iteration::new(); @@ -127,8 +230,37 @@ mod tests { out = var_out.complete(); } "##, - ) - .unwrap(); - assert_eq!(tokens.to_string(), expected_tokens.to_string()); + ); + } + #[test] + fn encode_transitive_closure() { + compare( + " + input inp(x: u32, y: u32) + output out(x: u32, y: u32) + out(x, y) :- inp(x, y). + out(x, y) :- out(x, z), out(z, y). + ", + r##" + { + let mut iteration = datafrog::Iteration::new(); + let var_inp = datafrog::Relation::from_vec(inp); + let var_out = iteration.variable:: <(u32, u32,)>("out"); + let var_inp1 = iteration.variable:: <(u32, u32,)>("inp1"); + let var_out2 = iteration.variable:: <((u32,), (u32,))>("out2"); + let var_out3 = iteration.variable:: <((u32,), (u32,))>("out3"); + let var_out4 = iteration.variable:: <(u32, u32, u32,)>("out4"); + var_inp1.insert(var_inp); + while iteration.changed() { + var_out.from_map(&var_inp1, | &(x, y,)| (x, y,)); + var_out2.from_map(&var_out, | &(x, z,)| ((z,), (x,))); + var_out3.from_map(&var_out, | &(z, y,)| ((z,), (y,))); + var_out4.from_join(&var_out2, &var_out3, | &(z,), &(x,), &(y,)| (z, x, y,)); + var_out.from_map(&var_out4, | &(z, x, y,)| (x, y,)); + } + out = var_out.complete(); + } + "##, + ); } } diff --git a/src/generator_new/to_tokens.rs b/src/generator_new/to_tokens.rs index df75364..0f7c832 100644 --- a/src/generator_new/to_tokens.rs +++ b/src/generator_new/to_tokens.rs @@ -28,16 +28,16 @@ impl ToTokens for DVar { impl ToTokens for DVarTuple { fn to_tokens(&self, tokens: &mut TokenStream) { - let mut vars = var_vec_to_tokens(&self.vars); + let vars = var_vec_to_tokens(&self.vars); tokens.extend(quote! {(#vars)}); } } impl ToTokens for DVarKeyVal { fn to_tokens(&self, tokens: &mut TokenStream) { - let mut key = var_vec_to_tokens(&self.key); - let mut value = var_vec_to_tokens(&self.value); - tokens.extend(quote! {((#key), (value))}); + let key = var_vec_to_tokens(&self.key); + let value = var_vec_to_tokens(&self.value); + tokens.extend(quote! {((#key), (#value))}); } } @@ -96,7 +96,30 @@ impl ToTokens for BindVarOp { impl ToTokens for JoinOp { fn to_tokens(&self, tokens: &mut TokenStream) { - unimplemented!(); + let JoinOp { + output, + input_first, + input_second, + key, + value_first, + value_second, + } = self; + let flattened = DVarTuple { + vars: key + .vars + .iter() + .chain(&value_first.vars) + .chain(&value_second.vars) + .cloned() + .collect(), + }; + + tokens.extend(quote! { + #output.from_join( + &#input_first, + &#input_second, + |&#key, &#value_first, &#value_second| #flattened); + }); } }