DDIR: compile join, half-join and reduce lambdas once

interactive/src/backend/corgi.rs

@@ -63,7 +63,7 @@ impl Plan {
     fn new(op: &LinearOp, kshape: Shape, vshape: Shape, level: Level) -> Plan {
         let (k, v) = (&kshape, &vshape);
         let (what, graph) = match op {
-            LinearOp::Project(p) => ("map", Some(compile_projection(p, &[kshape.clone(), vshape.clone()]))),
+            LinearOp::Project(p) => ("map", Some(compile_projection(p, &[kshape.clone(), vshape.clone()]).map(|c| c.graph().clone()))),
             LinearOp::Filter(cond) => ("filter", Some(compile_predicate(cond, k, v))),
             LinearOp::FlatMap(list) => ("flatmap", Some(compile_flatmap(list, k, v))),
             LinearOp::JoinFunction(lambda) => ("join_function", Some(compile_join_function(lambda, k, v, level.width()))),

interactive/src/lambda.rs

@@ -382,27 +382,28 @@ pub(crate) fn slot_shape(lambda: &Lambda, shapes: &[Shape]) -> Res<Shape> {
     corgi::shape_of(lambda.graph()?, &Shape::Prod(shapes.to_vec()))
 }
 
-/// A slot's graph over rows of `shapes`, and the shape it gives. Typechecked once here, so an `Ok`
-/// kernel runs on every batch of these shapes.
-fn slot_graph(lambda: &Lambda, shapes: &[Shape]) -> Res<(Graph<NumOp>, Shape)> {
+/// A slot's graph over rows of `shapes`, compiled at them, and the shape it gives. Typechecked once
+/// here, so an `Ok` kernel runs on every batch of these shapes.
+fn slot_graph(lambda: &Lambda, shapes: &[Shape]) -> Res<(corgi::Compiled, Shape)> {
     let output = slot_shape(lambda, shapes)?;
     // Typed compilation removes proven patterns before normalizing the executable graph.
     let compiled = corgi::compile(lambda.graph()?.as_ref().clone(), &Shape::Prod(shapes.to_vec()))?;
-    Ok((compiled.graph().clone(), output))
+    Ok((compiled, output))
 }
 
 /// Compile a `FlatMap`'s lambda → a corgi `List` column, one list per input row. A lambda that is
 /// not list-shaped is the type error: the backend explodes the column structurally.
 pub(crate) fn compile_flatmap(list: &Lambda, kshape: &Shape, vshape: &Shape) -> Res<Graph<NumOp>> {
     match slot_graph(list, &[kshape.clone(), vshape.clone()])? {
-        (k, Shape::List(_)) => Ok(k),
+        (k, Shape::List(_)) => Ok(k.graph().clone()),
         (_, other) => Err(format!("flatmap over a non-list: {other}")),
     }
 }
 
 /// Compile a `reduce` lambda over `(key, vals)`, `vals` a list of `(val, diff)` (`vshape` the
-/// values') → a corgi `List` column of `(val, diff)`, one list per key.
-pub(crate) fn compile_reduce(lambda: &Lambda, kshape: &Shape, vshape: &Shape) -> Res<Graph<NumOp>> {
+/// values') → a corgi `List` column of `(val, diff)`, one list per key. The result is compiled at
+/// that input, so the backend narrows it without compiling it again.
+pub(crate) fn compile_reduce(lambda: &Lambda, kshape: &Shape, vshape: &Shape) -> Res<corgi::Compiled> {
     let list = Shape::List(Box::new(Shape::Prod(vec![vshape.clone(), Shape::Int])));
     match slot_graph(lambda, &[kshape.clone(), list])? {
         (g, Shape::List(e)) if matches!(&*e, Shape::Prod(fs) if fs.len() == 2 && fs[1] == Shape::Int) => Ok(g),
@@ -419,7 +420,7 @@ pub(crate) fn compile_join_function(lambda: &Lambda, kshape: &Shape, vshape: &Sh
         if let Shape::Prod(fs) = &**elem {
             let time = Shape::Prod(vec![Shape::Int; width]);
             if fs.len() == 3 && matches!(&fs[0], Shape::Prod(kv) if kv.len() == 2) && fs[1] == time && fs[2] == Shape::Int {
-                return Ok(g);
+                return Ok(g.graph().clone());
             }
         }
     }
@@ -430,14 +431,15 @@ pub(crate) fn compile_join_function(lambda: &Lambda, kshape: &Shape, vshape: &Sh
 /// `Int`; any other shape is a type error.
 pub(crate) fn compile_predicate(cond: &Lambda, kshape: &Shape, vshape: &Shape) -> Res<Graph<NumOp>> {
     match slot_graph(cond, &[kshape.clone(), vshape.clone()])? {
-        (k, Shape::Int) => Ok(k),
+        (k, Shape::Int) => Ok(k.graph().clone()),
         (_, other) => Err(format!("a filter predicate must be an Int, got {other}")),
     }
 }
 
 /// Compile a projection over `rows`, the shapes of `(key, val)` or `(key, val0, val1)`: input
-/// `Prod(rows)`, output `Prod([newkey, newval])`.
-pub(crate) fn compile_projection(lambda: &Lambda, rows: &[Shape]) -> Res<Graph<NumOp>> {
+/// `Prod(rows)`, output `Prod([newkey, newval])`. The result is compiled at that input, so the
+/// join backends narrow it without compiling it again.
+pub(crate) fn compile_projection(lambda: &Lambda, rows: &[Shape]) -> Res<corgi::Compiled> {
     match slot_graph(lambda, rows)? {
         (g, Shape::Prod(fs)) if fs.len() == 2 => Ok(g),
         (_, other) => Err(format!("a projection gives (key, val), not {other}")),

interactive/src/reducers.rs

@@ -174,15 +174,15 @@ mod tests {
         assert_eq!(readable.graph().unwrap(), mechanical.graph().unwrap());
         let parsed = Lambda::new(readable.src.clone(), &corgi::Scope::default());
         let val = Shape::Prod(vec![Shape::Int, Shape::Int]);
-        assert_eq!(crate::lambda::compile_reduce(&readable, &Shape::Int, &val).unwrap(),
-            crate::lambda::compile_reduce(&parsed, &Shape::Int, &val).unwrap());
+        assert_eq!(crate::lambda::compile_reduce(&readable, &Shape::Int, &val).unwrap().graph(),
+            crate::lambda::compile_reduce(&parsed, &Shape::Int, &val).unwrap().graph());
         // The known input type also proves authored nested pattern checks.
         let base = Lambda::new("let ((k, _), vals) = input in (vals head.0, 1) enlist".into(), &corgi::Scope::default());
         let Reducer::Expr(kept) = Reducer::Expr(base).project(&projection, &Shape::Prod(vec![Shape::Int, Shape::Int]), &Shape::Prod(vec![Shape::Int, Shape::Int]));
         assert!(!kept.src.contains("_call0"), "{}", kept.src);
         let parsed = Lambda::new(kept.src.clone(), &corgi::Scope::default());
-        assert_eq!(crate::lambda::compile_reduce(&kept, &val, &val).unwrap(),
-            crate::lambda::compile_reduce(&parsed, &val, &val).unwrap());
+        assert_eq!(crate::lambda::compile_reduce(&kept, &val, &val).unwrap().graph(),
+            crate::lambda::compile_reduce(&parsed, &val, &val).unwrap().graph());
     }
 
     #[test]
@@ -198,7 +198,7 @@ mod tests {
             let reparsed = Lambda::new(lambda.src.clone(), &corgi::Scope::default());
             let running = crate::lambda::compile_reduce(&lambda, &Shape::Int, &shape).unwrap();
             let printed = crate::lambda::compile_reduce(&reparsed, &Shape::Int, &shape).unwrap();
-            assert_eq!(running, printed, "{sugar:?}: {}", lambda.src);
+            assert_eq!(running.graph(), printed.graph(), "{sugar:?}: {}", lambda.src);
             // The spelling emitted by ddir-plan is already valid DDP. Replanning it
             // must not add a mask to the output that was narrowed in the expression.
             let text = format!("export \"out\" = input 0 : (int, (int, int)) | reduce(input -> {});", lambda.src);
@@ -207,7 +207,7 @@ mod tests {
             assert_eq!(crate::passes::projection::pushdown(&mut program, &|_| None).unwrap(), 0);
         }
         // Nonempty input with zero total count emits nothing; a negative total is retained.
-        let lambda = crate::lambda::compile_reduce(&Sugar::Count.lambda(), &Shape::Int, &Shape::Int).unwrap();
+        let lambda = crate::lambda::compile_reduce(&Sugar::Count.lambda(), &Shape::Int, &Shape::Int).unwrap().graph().clone();
         let input = Value::Prod(vec![Value::i64(vec![1, 2]), Value::List(Bounds::offsets(vec![2, 3]), Box::new(Value::Prod(vec![
             Value::i64(vec![10, 20, 30]), Value::i64(vec![1, -1, -2]),
         ])))]);

interactive/src/runtime/half_join/backend.rs

@@ -57,10 +57,10 @@ impl CorgiHalfJoinBackend {
     /// type error is a panic with corgi's message: install type-checked the program, so it is not
     /// reached.
     pub(crate) fn new(projection: &Lambda, width: usize, [k, v0, v1]: [Shape; 3]) -> Self {
-        let g = compile_projection(projection, &[k.clone(), v0.clone(), v1.clone()])
+        let compiled = compile_projection(projection, &[k.clone(), v0.clone(), v1.clone()])
             .unwrap_or_else(|e| panic!("half_join: type error at shapes ({k}, {v0}, {v1}): {e}"));
         let input = Shape::Prod(vec![k, v0, v1]);
-        let narrowed = corgi::compile(g, &input).and_then(|compiled| corgi::narrow(compiled, &Demand::All))
+        let narrowed = corgi::narrow(compiled, &Demand::All)
             .unwrap_or_else(|e| panic!("half_join: narrowing a compiled projection at {input}: {e}"));
         let reads = Reads::new(narrowed.input, &input);
         let Shape::Prod(parts) = &input else { unreachable!() };
@@ -645,7 +645,7 @@ mod tests {
                 let lambda = Lambda::new(format!("let (k, v0, v1) = input in let (a, b, c, d) = v0 in let (e, f, g, h) = v1 in {body}"), &corgi::Scope::default());
                 let narrowed = CorgiHalfJoinBackend::new(&lambda, 2, shapes.clone());
                 let original = CorgiHalfJoinBackend {
-                    width: 2, project_stream: false, graph: compile_projection(&lambda, &shapes).unwrap(),
+                    width: 2, project_stream: false, graph: compile_projection(&lambda, &shapes).unwrap().graph().clone(),
                     reads: Reads::new(Demand::All, &Shape::Prod(shapes.to_vec())), colliding: Vec::new(),
                 };
                 let rows = |backend: CorgiHalfJoinBackend| {

interactive/src/runtime/join.rs

@@ -98,10 +98,10 @@ impl<T: ColTime> CorgiJoinBackend<T> {
     /// The projection over inputs of the key's and both values' shapes. A type error is a panic
     /// with corgi's message: install type-checked the program, so it is not reached.
     pub(crate) fn new(projection: &Lambda, [k, v0, v1]: [Shape; 3]) -> Self {
-        let g = compile_projection(projection, &[k.clone(), v0.clone(), v1.clone()])
+        let compiled = compile_projection(projection, &[k.clone(), v0.clone(), v1.clone()])
             .unwrap_or_else(|e| panic!("join projection: type error at shapes ({k}, {v0}, {v1}): {e}"));
         let input = Shape::Prod(vec![k, v0, v1]);
-        let narrowed = corgi::compile(g, &input).and_then(|compiled| corgi::narrow(compiled, &Demand::All))
+        let narrowed = corgi::narrow(compiled, &Demand::All)
             .unwrap_or_else(|e| panic!("join projection: narrowing a compiled graph at {input}: {e}"));
         CorgiJoinBackend { colliding: Vec::new(), graph: narrowed.graph, reads: Reads::new(narrowed.input, &input), _t: PhantomData }
     }
@@ -1097,7 +1097,7 @@ mod tests {
                 let lambda = Lambda::new(format!("let (k, (a, b, c, d), v) = input in {body}"), &corgi::Scope::default());
                 let mut narrowed = CorgiJoinBackend::new(&lambda, shapes.clone());
                 let mut original = CorgiJoinBackend {
-                    colliding: Vec::new(), graph: compile_projection(&lambda, &shapes).unwrap(),
+                    colliding: Vec::new(), graph: compile_projection(&lambda, &shapes).unwrap().graph().clone(),
                     reads: Reads::new(Demand::All, &Shape::Prod(shapes.to_vec())), _t: PhantomData,
                 };
                 let run = |backend: &mut CorgiJoinBackend<u64>| {

interactive/src/runtime/reduce.rs

@@ -213,12 +213,10 @@ impl<T> CorgiReduceBackend<T> {
     /// type check keeps from being reached.
     pub(crate) fn with_window(reducer: Reducer, limit: Option<usize>, (kshape, vshape): (Shape, Shape), window_size: usize) -> Self {
         let Reducer::Expr(lambda) = reducer;
-        let graph = compile_reduce(&lambda, &kshape, &vshape)
+        let compiled = compile_reduce(&lambda, &kshape, &vshape)
             .unwrap_or_else(|e| panic!("reduce: type error at shapes ({kshape}, {vshape}): {e}"));
         // The graph rebuilt to read only what it reads, and the demand saying what that is.
         let input = Shape::Prod(vec![kshape, Shape::List(Box::new(Shape::Prod(vec![vshape, Shape::Int])))]);
-        let compiled = corgi::compile(graph, &input)
-            .unwrap_or_else(|e| panic!("reduce: compiling at {input}: {e}"));
         // A limit observes whole-value order and entry boundaries. Demand alone does not prove
         // that merging those entries is valid; the sum analysis also checks disappearing groups.
         let project_values = limit.is_none() && crate::passes::reduce_demand::projection(&compiled).is_some();