DDIR: slots written as corgi lambdas

interactive/Cargo.toml

@@ -14,7 +14,7 @@ workspace = true
 [dependencies]
 columnar = { workspace = true }
 # The columnar kernels for the interpreted backend, pinned by git rev.
-corgi = { git = "https://github.com/frankmcsherry/WIP", rev = "9737792dd5db4b99e8299898ff917b174d37778b", features = ["serde"] }
+corgi = { git = "https://github.com/frankmcsherry/WIP", rev = "91df1f0ba2a16b48e3ebee5c6e5c77dbb3cf527e", features = ["serde"] }
 differential-dataflow = { workspace = true }
 serde = { version = "1.0", features = ["derive"] }
 smallvec = "1.15.1"

interactive/src/corgi/logic.rs

@@ -188,7 +188,7 @@ fn mentions_env(t: &Term, depth: usize) -> bool {
         Term::Var(_) => true,
         Term::Bound(k) => *k >= depth,
         Term::Int(_) => false,
-        Term::Tuple(fs) | Term::List(fs) | Term::Hash(fs) | Term::Call(_, fs) => fs.iter().any(|f| mentions_env(f, depth)),
+        Term::Tuple(fs) | Term::List(fs) | Term::Hash(fs) | Term::Call(_, fs) | Term::Corgi(_, fs) => fs.iter().any(|f| mentions_env(f, depth)),
         Term::Spread(inner) | Term::Proj(inner, _) | Term::Unary(_, inner) => mentions_env(inner, depth),
         Term::Inject { tag, payload, .. } => mentions_env(tag, depth) || mentions_env(payload, depth),
         Term::Case { scrutinee, arms, default } => {
@@ -278,6 +278,17 @@ pub fn compile(
             };
             Ok(b.add(NumOp::Host(corgi::HostOp(kernel)), vec![input]))
         }
+        Term::Corgi(src, args) => {
+            // The program's input is the tuple of its arguments, a `Unit` over the anchor for none.
+            let input = if args.is_empty() {
+                b.add(Op::Unit, vec![anchor])
+            } else {
+                let fields = args.iter().map(|a| compile(a, b, env, env_shapes, anchor, None)).collect::<Res<Vec<_>>>()?;
+                b.tuple(fields)
+            };
+            let program = parsed(src)?;
+            Ok(b.graft(&program, input))
+        }
         Term::Var(i) => env.get(*i).copied().ok_or_else(|| format!("`${i}` is not in scope here")),
         Term::Bound(k) => {
             env.len().checked_sub(1 + *k).map(|i| env[i]).ok_or_else(|| format!("binder `^{k}` is not in scope here"))
@@ -763,7 +774,8 @@ fn lower(terms: &[&Term], shapes: &[Shape]) -> Res<(Graph<NumOp>, Shape)> {
     let env: Vec<usize> = (0..shapes.len()).map(|i| b.add(Op::Field(i), vec![input])).collect();
     let outs = terms.iter().map(|t| compile(t, &mut b, &env, shapes, input, None)).collect::<Res<Vec<_>>>()?;
     let out = if let [one] = outs[..] { one } else { b.tuple(outs) };
-    let g = b.finish(out);
+    // CSE, so a term compiled twice runs once: a lambda's key and val are two fields of one program.
+    let g = corgi::cse(&b.finish(out));
     let output = corgi::shape_of(&g, &Shape::Prod(shapes.to_vec()))?;
     Ok((g, output))
 }
@@ -807,6 +819,31 @@ pub fn compile_projection(key: &Term, val: &Term, kshape: &Shape, vshape: &Shape
     Ok(lower(&[key, val], &[kshape.clone(), vshape.clone()])?.0)
 }
 
+/// Evaluate the corgi program `src` on one row, as `ir::eval` does for `Term::Corgi`. The row's
+/// shape is its own, so a row whose shape it cannot tell alone (an empty list, a variant) panics.
+pub fn eval_corgi_row(src: &str, input: DValue) -> DValue {
+    let shape = shape_of_row(&input).unwrap_or_else(|e| panic!("corgi program `{src}`: {e}"));
+    let graph = parsed(src).unwrap_or_else(|e| panic!("corgi program `{src}`: {e}"));
+    let out = corgi::eval_graph(&graph, transcode_owned(vec![input], &shape));
+    let out_shape = corgi::shape_of_value(&out);
+    untranscode(out, &out_shape).pop().expect("one row in, one row out")
+}
+
+/// The graph of the corgi program `src` (`Term::Corgi`), parsed once for the life of the process.
+/// `Err` is a parse error.
+fn parsed(src: &str) -> Res<std::sync::Arc<Graph<NumOp>>> {
+    use std::collections::HashMap;
+    use std::sync::{Arc, Mutex, OnceLock};
+    static GRAPHS: OnceLock<Mutex<HashMap<String, Arc<Graph<NumOp>>>>> = OnceLock::new();
+    let mut map = GRAPHS.get_or_init(Default::default).lock().unwrap();
+    if let Some(g) = map.get(src) {
+        return Ok(Arc::clone(g));
+    }
+    let g = Arc::new(corgi::parse_ml(src)?);
+    map.insert(src.to_string(), Arc::clone(&g));
+    Ok(g)
+}
+
 /// The Float functions without a composition of corgi ops, as host kernels over the float
 /// column: apply the function as `ir::eval` does. One kernel per function, shared by every call
 /// site, so CSE merges equal calls.

interactive/src/explain/mod.rs

@@ -264,6 +264,7 @@ fn with_rows(t: &Term, rows: &[Term]) -> Term {
         Term::List(ts) => Term::List(ts.iter().map(go).collect()),
         Term::Hash(ts) => Term::Hash(ts.iter().map(go).collect()),
         Term::Call(name, ts) => Term::Call(name.clone(), ts.iter().map(go).collect()),
+        Term::Corgi(src, ts) => Term::Corgi(src.clone(), ts.iter().map(go).collect()),
         Term::Spread(t) => Term::Spread(bx(t)),
         Term::Proj(t, i) => Term::Proj(bx(t), *i),
         Term::Unary(op, t) => Term::Unary(*op, bx(t)),

interactive/src/ir.rs

@@ -131,6 +131,9 @@ pub enum Term {
     /// columnar kernel if one is registered (`register_kernel`), otherwise its
     /// row body over just the call's arguments.
     Call(String, Vec<Term>),
+    /// A corgi ML program, run with the tuple of the arguments as its `input`. A slot written
+    /// `pat -> body` is `let pat = input in body` over the slot's input rows.
+    Corgi(String, Vec<Term>),
 }
 
 /// The sum type an `Inject` builds into. `Declared` carries the full lane shapes of a `type`
@@ -327,6 +330,10 @@ pub fn eval(term: &Term, env: &mut Vec<Value>) -> Value {
             }
             (f.body)(&vals)
         }
+        Term::Corgi(src, args) => {
+            let input = Value::Tuple(args.iter().map(|a| eval(a, env)).collect());
+            crate::corgi::logic::eval_corgi_row(src, input)
+        }
         Term::Var(i) => env[*i].clone(),
         Term::Bound(k) => env[env.len() - 1 - *k].clone(),
         Term::Int(n) => Value::Int(*n),

interactive/src/parse/pipe.rs

@@ -92,6 +92,14 @@
 //! - Binders: `^k` refers to the k-th enclosing `case`/`fold` binder (de
 //!   Bruijn, innermost = 0); `case` patterns also bind payload fields by name.
 //! - Conditional: `if(cond, then, els)`.
+//!
+//! # Corgi lambdas
+//!
+//! Any slot may instead be a corgi lambda, `pat -> body`, in corgi's ML syntax: it runs
+//! `let pat = input in body` with `input` the tuple of the slot's rows, `(key, val)` for a linear
+//! op and `(key, left val, right val)` for a join. In `key`, `map`, and `join` the body gives the
+//! pair `(key, val)`; in `value`, `filter`, `flatmap`, and `enter_at` it gives the one value. For
+//! example, `key(((src, dst), ()) -> ((src), (dst)))` is `key($0[0] ; $0[1])`.
 
 use super::*;
 use crate::ir::{BinOp, Projection, Reducer, SumTy, Term, UnOp};
@@ -101,21 +109,83 @@ enum Token {
     Let, Var, Export, Type, Case, Fold,
     Input, Import, Key, Map, Join, Min, Distinct, Count, Collect, Arrange, Negate, Filter, EnterAt, LiftIter, FlatMap, Inspect,
     Ident(String), Int(i64), Str(String),
+    /// A slot written as a corgi lambda, `pat -> body`, kept as corgi source.
+    Lambda { pat: String, body: String },
     Dollar, Caret, LParen, RParen, LBrace, RBrace, LBracket, RBracket,
     Comma, Semi, Colon, ColonColon, Eq, EqEq, NotEq, Lt, LtEq, Gt, GtEq, AndAnd, FatArrow,
     Pipe, Plus, Minus, Star, Eof,
 }
 
+/// The source as a stream of chars, with the byte offset of the next one.
+#[derive(Clone)]
+struct Chars<'a> { src: &'a str, at: usize, peeked: Option<char> }
+
+impl Chars<'_> {
+    fn peek(&mut self) -> Option<&char> { self.peeked = self.src[self.at..].chars().next(); self.peeked.as_ref() }
+}
+
+impl Iterator for Chars<'_> {
+    type Item = char;
+    fn next(&mut self) -> Option<char> { let c = self.src[self.at..].chars().next()?; self.at += c.len_utf8(); Some(c) }
+}
+
+/// A corgi lambda that ends the parenthesized group starting at byte `from` of `src` (just after
+/// its `(`): the group holds a `->` outside any nested bracket. Returns the lambda's start (after
+/// the group's last comma before the arrow, or `from`), the arrow, and the group's `)`.
+///
+/// Corgi's text is not DDIR's (it has `.2`, float literals, `#` comments), so the lexer takes the
+/// lambda whole rather than lexing it. A pattern holds no comma outside parentheses, so the last
+/// comma before the arrow separates the lambda from the arguments before it, as in `join`.
+fn lambda_in(src: &str, from: usize) -> Option<(usize, usize, usize)> {
+    let bytes = src.as_bytes();
+    let (mut depth, mut start, mut arrow) = (0usize, from, None);
+    let mut i = from;
+    while i < bytes.len() {
+        match bytes[i] {
+            b'"' => {
+                i += 1;
+                while i < bytes.len() && bytes[i] != b'"' { if bytes[i] == b'\\' { i += 1; } i += 1; }
+            }
+            b'#' => { while i < bytes.len() && bytes[i] != b'\n' { i += 1; } }
+            b'-' if bytes.get(i + 1) == Some(&b'-') => { while i < bytes.len() && bytes[i] != b'\n' { i += 1; } }
+            b'-' if bytes.get(i + 1) == Some(&b'>') && depth == 0 && arrow.is_none() => { arrow = Some(i); i += 1; }
+            b'(' | b'[' | b'{' => depth += 1,
+            b')' | b']' | b'}' if depth == 0 => return arrow.map(|a| (start, a, i)),
+            b')' | b']' | b'}' => depth -= 1,
+            b',' if depth == 0 && arrow.is_none() => start = i + 1,
+            _ => {}
+        }
+        i += 1;
+    }
+    None
+}
+
 fn tokenize(input: &str) -> Vec<Token> {
+    let mut tokens = lex(input);
+    tokens.push(Token::Eof);
+    tokens
+}
+
+fn lex(input: &str) -> Vec<Token> {
     let mut tokens = Vec::new();
-    let mut chars = input.chars().peekable();
+    let mut chars = Chars { src: input, at: 0, peeked: None };
     while let Some(&ch) = chars.peek() {
         match ch {
             ' ' | '\t' | '\n' | '\r' => { chars.next(); },
             '-' if chars.clone().nth(1).map_or(false, |c| c == '-') => {
                 while let Some(&c) = chars.peek() { chars.next(); if c == '\n' { break; } }
             },
-            '(' => { chars.next(); tokens.push(Token::LParen); },
+            '(' => {
+                chars.next();
+                tokens.push(Token::LParen);
+                if let Some((start, arrow, close)) = lambda_in(input, chars.at) {
+                    tokens.extend(lex(&input[chars.at..start]));
+                    let (pat, body) = (input[start..arrow].trim().to_string(), input[arrow + 2..close].trim().to_string());
+                    tokens.push(Token::Lambda { pat, body });
+                    tokens.push(Token::RParen);
+                    chars.at = close + 1;
+                }
+            },
             ')' => { chars.next(); tokens.push(Token::RParen); },
             '{' => { chars.next(); tokens.push(Token::LBrace); },
             '}' => { chars.next(); tokens.push(Token::RBrace); },
@@ -161,7 +231,6 @@ fn tokenize(input: &str) -> Vec<Token> {
             other => panic!("Unexpected character: {:?}", other),
         }
     }
-    tokens.push(Token::Eof);
     tokens
 }
 
@@ -455,19 +524,19 @@ impl Parser {
             Token::Map => { self.next(); let p = self.parse_projection(); Expr::Map(Box::new(lhs), p) },
             Token::Ident(name) if name == "value" => {
                 self.next(); self.expect(&Token::LParen);
-                let val = self.parse_term(); self.expect(&Token::RParen);
+                let val = self.parse_slot_term(); self.expect(&Token::RParen);
                 Expr::Map(Box::new(lhs), Projection { key: Term::Var(0), val })
             },
-            Token::Join => { self.next(); self.expect(&Token::LParen); let r = self.parse_join_arg(); self.expect(&Token::Comma); let p = self.parse_projection(); self.expect(&Token::RParen); Expr::Join(Box::new(lhs), Box::new(r), p) },
+            Token::Join => { self.next(); self.expect(&Token::LParen); let r = self.parse_join_arg(); self.expect(&Token::Comma); let p = match self.lambda(3) { Some(l) => Self::split(l), None => self.parse_projection() }; self.expect(&Token::RParen); Expr::Join(Box::new(lhs), Box::new(r), p) },
             Token::Min => { self.next(); Expr::Reduce(Box::new(lhs), Reducer::Min) },
             Token::Distinct => { self.next(); Expr::Reduce(Box::new(lhs), Reducer::Distinct) },
             Token::Count => { self.next(); Expr::Reduce(Box::new(lhs), Reducer::Count) },
             Token::Collect => { self.next(); Expr::Reduce(Box::new(lhs), Reducer::Collect) },
-            Token::FlatMap => { self.next(); self.expect(&Token::LParen); let l = self.parse_term(); self.expect(&Token::RParen); Expr::FlatMap(Box::new(lhs), l) },
+            Token::FlatMap => { self.next(); self.expect(&Token::LParen); let l = self.parse_slot_term(); self.expect(&Token::RParen); Expr::FlatMap(Box::new(lhs), l) },
             Token::Arrange => { self.next(); Expr::Arrange(Box::new(lhs)) },
             Token::Negate => { self.next(); Expr::Negate(Box::new(lhs)) },
-            Token::Filter => { self.next(); self.expect(&Token::LParen); let c = self.parse_term(); self.expect(&Token::RParen); Expr::Filter(Box::new(lhs), c) },
-            Token::EnterAt => { self.next(); self.expect(&Token::LParen); let f = self.parse_term(); self.expect(&Token::RParen); Expr::EnterAt(Box::new(lhs), f) },
+            Token::Filter => { self.next(); self.expect(&Token::LParen); let c = self.parse_slot_term(); self.expect(&Token::RParen); Expr::Filter(Box::new(lhs), c) },
+            Token::EnterAt => { self.next(); self.expect(&Token::LParen); let f = self.parse_slot_term(); self.expect(&Token::RParen); Expr::EnterAt(Box::new(lhs), f) },
             Token::LiftIter => { self.next(); Expr::LiftIter(Box::new(lhs)) },
             Token::Inspect => { self.next(); self.expect(&Token::LParen); let l = self.parse_ident(); self.expect(&Token::RParen); Expr::Inspect(Box::new(lhs), l) },
             other => panic!("Expected pipe operator, got {:?}", other),
@@ -476,9 +545,31 @@ impl Parser {
 
     fn parse_projection(&mut self) -> Projection {
         self.expect(&Token::LParen);
+        if let Some(l) = self.lambda(2) {
+            self.expect(&Token::RParen);
+            return Self::split(l);
+        }
         self.parse_projection_inner()
     }
 
+    /// A slot written as a corgi lambda, over the slot's `arity` input rows (two for a linear
+    /// op, three for a join): a term running `let pat = input in body` on the tuple of them.
+    fn lambda(&mut self, arity: usize) -> Option<Term> {
+        let Token::Lambda { pat, body } = self.peek().clone() else { return None };
+        self.next();
+        Some(Term::Corgi(format!("let {pat} = input in {body}"), (0..arity).map(Term::Var).collect()))
+    }
+
+    /// The projection a lambda's `(key, val)` result gives.
+    fn split(lambda: Term) -> Projection {
+        Projection { key: Term::Proj(Box::new(lambda.clone()), 0), val: Term::Proj(Box::new(lambda), 1) }
+    }
+
+    /// A linear op's term slot: a corgi lambda, or a term.
+    fn parse_slot_term(&mut self) -> Term {
+        self.lambda(2).unwrap_or_else(|| self.parse_term())
+    }
+
     // A projection `(k1, k2 ; v1, v2)` builds `key = tuple(k1, k2)` and
     // `val = tuple(v1, v2)`. A bare `$n` field splices the input row's tuple
     // fields (`Spread`), or is one field when the row is not a tuple; any other