make from_join take a Relation, add Relation::from_join

Also `Relation::from_vec` public

Author: nikomatsakis

src/join.rs

@@ -34,6 +34,20 @@ pub(crate) fn join_into<'me, Key: Ord, Val1: Ord, Val2: Ord, Result: Ord>(
     output.insert(Relation::from_vec(results));
 }
 
+pub(crate) fn join_into_relation<'me, Key: Ord, Val1: Ord, Val2: Ord, Result: Ord>(
+    input1: &Relation<(Key, Val1)>,
+    input2: &Relation<(Key, Val2)>,
+    mut logic: impl FnMut(&Key, &Val1, &Val2) -> Result,
+) -> Relation<Result> {
+    let mut results = Vec::new();
+
+    join_helper(&input1.elements, &input2.elements, |k, v1, v2| {
+        results.push(logic(k, v1, v2));
+    });
+
+    Relation::from_vec(results)
+}
+
 /// Moves all recent tuples from `input1` that are not present in `input2` into `output`.
 pub(crate) fn antijoin_into<Key: Ord, Val: Ord, Result: Ord>(
     input1: &Variable<(Key, Val)>,
@@ -116,11 +130,22 @@ pub(crate) fn gallop<T>(mut slice: &[T], mut cmp: impl FnMut(&T) -> bool) -> &[T
     slice
 }
 
+/// An input that can be used with `from_join`; either a `Variable` or a `Relation`.
 pub trait JoinInput<'me, Tuple: Ord>: Copy {
+    /// If we are on iteration N of the loop, these are the tuples
+    /// added on iteration N-1. (For a `Relation`, this is always an
+    /// empty slice.)
     type RecentTuples: Deref<Target = [Tuple]>;
+
+    /// If we are on iteration N of the loop, these are the tuples
+    /// added on iteration N - 2 or before. (For a `Relation`, this is
+    /// just `self`.)
     type StableTuples: Deref<Target = [Relation<Tuple>]>;
 
+    /// Get the set of recent tuples.
     fn recent(self) -> Self::RecentTuples;
+
+    /// Get the set of stable tuples.
     fn stable(self) -> Self::StableTuples;
 }
 

src/lib.rs

@@ -19,6 +19,7 @@ mod join;
 mod map;
 mod test;
 mod treefrog;
+pub use crate::join::JoinInput;
 pub use crate::treefrog::{
     extend_anti::ExtendAnti,
     extend_with::ExtendWith,
@@ -98,7 +99,20 @@ impl<Tuple: Ord> Relation<Tuple> {
         Relation { elements }
     }
 
-    fn from_vec(mut elements: Vec<Tuple>) -> Self {
+    /// Creates a `Relation` by joining the values from `input1` and
+    /// `input2` and then applying `logic`. Like
+    /// [`Variable::from_join`] except for use where the inputs are
+    /// not varying across iterations.
+    pub fn from_join<Key: Ord, Val1: Ord, Val2: Ord>(
+        input1: &Relation<(Key, Val1)>,
+        input2: &Relation<(Key, Val2)>,
+        logic: impl FnMut(&Key, &Val1, &Val2) -> Tuple,
+    ) -> Self {
+        join::join_into_relation(input1, input2, logic)
+    }
+
+    /// Creates a `Relation` from a vector of tuples.
+    pub fn from_vec(mut elements: Vec<Tuple>) -> Self {
         elements.sort();
         elements.dedup();
         Relation { elements }
@@ -201,7 +215,22 @@ pub struct Variable<Tuple: Ord> {
 
 // Operator implementations.
 impl<Tuple: Ord> Variable<Tuple> {
-    /// Adds tuples that result from joining `input1` and `input2`.
+    /// Adds tuples that result from joining `input1` and `input2` --
+    /// each of the inputs must be a set of (Key, Value) tuples. Both
+    /// `input1` and `input2` must have the same type of key (`K`) but
+    /// they can have distinct value types (`V1` and `V2`
+    /// respectively). The `logic` closure will be invoked for each
+    /// key that appears in both inputs; it is also given the two
+    /// values, and from those it should construct the resulting
+    /// value.
+    ///
+    /// Note that `input1` must be a variable, but `input2` can be a
+    /// relation or a variable. Therefore, you cannot join two
+    /// relations with this method. This is not because the result
+    /// would be wrong, but because it would be inefficient: the
+    /// result from such a join cannot vary across iterations (as
+    /// relations are fixed), so you should prefer to invoke `insert`
+    /// on a relation created by `Relation::from_join` instead.
     ///
     /// # Examples
     ///
@@ -224,10 +253,10 @@ impl<Tuple: Ord> Variable<Tuple> {
     /// let result = variable.complete();
     /// assert_eq!(result.len(), 121);
     /// ```
-    pub fn from_join<K: Ord, V1: Ord, V2: Ord>(
+    pub fn from_join<'me, K: Ord, V1: Ord, V2: Ord>(
         &self,
-        input1: &Variable<(K, V1)>,
-        input2: &Variable<(K, V2)>,
+        input1: &'me Variable<(K, V1)>,
+        input2: impl JoinInput<'me, (K, V2)>,
         logic: impl FnMut(&K, &V1, &V2) -> Tuple,
     ) {
         join::join_into(input1, input2, self, logic)

src/test.rs

@@ -1,16 +1,17 @@
 #![cfg(test)]
 
-use crate::Relation;
 use crate::Iteration;
+use crate::Relation;
 use crate::RelationLeaper;
 use proptest::prelude::*;
 use proptest::{proptest, proptest_helper};
 
 fn inputs() -> impl Strategy<Value = Vec<(u32, u32)>> {
-    prop::collection::vec((0_u32..100, 0_u32..100), 1..100)
+    prop::collection::vec((0_u32..100, 0_u32..100), 1..500)
 }
 
-fn reachable_with_join(edges: &[(u32, u32)]) -> Relation<(u32, u32)> {
+/// The original way to use datafrog -- computes reachable nodes from a set of edges
+fn reachable_with_var_join(edges: &[(u32, u32)]) -> Relation<(u32, u32)> {
     let edges = Relation::from(edges.iter().cloned());
     let mut iteration = Iteration::new();
 
@@ -28,6 +29,25 @@ fn reachable_with_join(edges: &[(u32, u32)]) -> Relation<(u32, u32)> {
     reachable.complete()
 }
 
+/// Like `reachable`, but using a relation as an input to `from_join`
+fn reachable_with_relation_join(edges: &[(u32, u32)]) -> Relation<(u32, u32)> {
+    let edges = Relation::from(edges.iter().cloned());
+    let mut iteration = Iteration::new();
+
+    // NB. Changed from `reachable_with_var_join`:
+    let edges_by_successor = Relation::from(edges.iter().map(|&(n1, n2)| (n2, n1)));
+
+    let reachable = iteration.variable::<(u32, u32)>("reachable");
+    reachable.insert(edges);
+
+    while iteration.changed() {
+        // reachable(N1, N3) :- edges(N1, N2), reachable(N2, N3).
+        reachable.from_join(&reachable, &edges_by_successor, |&_, &n3, &n1| (n1, n3));
+    }
+
+    reachable.complete()
+}
+
 fn reachable_with_leapfrog(edges: &[(u32, u32)]) -> Relation<(u32, u32)> {
     let edges = Relation::from(edges.iter().cloned());
     let mut iteration = Iteration::new();
@@ -41,21 +61,68 @@ fn reachable_with_leapfrog(edges: &[(u32, u32)]) -> Relation<(u32, u32)> {
         // reachable(N1, N3) :- edges(N1, N2), reachable(N2, N3).
         reachable.from_leapjoin(
             &reachable,
-            &mut [
-                &mut edges_by_successor.extend_with(|&(n2, _)| n2),
-            ],
+            &mut [&mut edges_by_successor.extend_with(|&(n2, _)| n2)],
             |&(_, n3), &n1| (n1, n3),
         );
     }
 
     reachable.complete()
 }
 
+/// Computes a join where the values are summed -- uses iteration
+/// variables (the original datafrog technique).
+fn sum_join_via_var(
+    input1_slice: &[(u32, u32)],
+    input2_slice: &[(u32, u32)],
+) -> Relation<(u32, u32)> {
+    let mut iteration = Iteration::new();
+
+    let input1 = iteration.variable::<(u32, u32)>("input1");
+    input1.insert(Relation::from(input1_slice.iter().cloned()));
+
+    let input2 = iteration.variable::<(u32, u32)>("input1");
+    input2.insert(Relation::from(input2_slice.iter().cloned()));
+
+    let output = iteration.variable::<(u32, u32)>("output");
+
+    while iteration.changed() {
+        // output(K1, V1 * 100 + V2) :- input1(K1, V1), input2(K1, V2).
+        output.from_join(&input1, &input2, |&k1, &v1, &v2| (k1, v1 * 100 + v2));
+    }
+
+    output.complete()
+}
+
+/// Computes a join where the values are summed -- uses iteration
+/// variables (the original datafrog technique).
+fn sum_join_via_relation(
+    input1_slice: &[(u32, u32)],
+    input2_slice: &[(u32, u32)],
+) -> Relation<(u32, u32)> {
+    let input1 = Relation::from(input1_slice.iter().cloned());
+    let input2 = Relation::from(input2_slice.iter().cloned());
+    Relation::from_join(&input1, &input2, |&k1, &v1, &v2| (k1, v1 * 100 + v2))
+}
+
 proptest! {
     #[test]
-    fn reachable(edges in inputs()) {
-        let reachable1 = reachable_with_join(&edges);
+    fn reachable_leapfrog_vs_var_join(edges in inputs()) {
+        let reachable1 = reachable_with_var_join(&edges);
         let reachable2 = reachable_with_leapfrog(&edges);
         assert_eq!(reachable1.elements, reachable2.elements);
     }
+
+    #[test]
+    fn reachable_rel_join_vs_var_join(edges in inputs()) {
+        let reachable1 = reachable_with_var_join(&edges);
+        let reachable2 = reachable_with_relation_join(&edges);
+        assert_eq!(reachable1.elements, reachable2.elements);
+    }
+
+    #[test]
+    fn sum_join_from_var_vs_rel((set1, set2) in (inputs(), inputs())) {
+        let output1 = sum_join_via_var(&set1, &set2);
+        let output2 = sum_join_via_relation(&set1, &set2);
+        assert_eq!(output1.elements, output2.elements);
+    }
 }