Add OR (disjunction) to rule bodies

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## [Unreleased] - ReleaseDate
 
+- **Disjunction (`or`) in rule bodies.** A rule body may now contain `(or (branch...) (branch...) ...)` (either spelling `or`/`OR`), where each branch is a conjunction of facts (or a single bare fact); the disjunction matches when at least one branch matches. A branch may be *correlated*: it may reference variables bound by the surrounding conjunction (e.g. `(v a b c)` outside, `(stale a al)` inside a branch). Variables usable outside the `or` are those bound by the surrounding conjunction plus those common to every branch; a branch-local variable that escapes its `or`, and an empty branch, are type errors. `or` is only allowed in rule bodies (including `rewrite` conditions), not in query-shaped commands like `check`, and is allowed under the term encoding (without proofs). It compiles to a single backend rule with a fused, **deduplicating** union node in the free-join engine: every branch is enumerated additively and its output tuple is materialized and deduplicated on the union's output variables, so the shared action fires exactly once per distinct output tuple. For a correlated `or`, the surrounding conjunction is prepended into each branch so the branches bind and dedup on the shared row (disjunctive-semijoin single-rebuild: a row matched via several branches rebuilds once), and it runs **seminaive / delta-driven** — a new tuple in a branch atom drives an index probe of the row rather than a re-scan (semi-naive of a union of conjunctions = the union of the per-disjunct expansions). An independent `or` scans the conjunction once as a continuation joined against the deduplicated branch outputs and runs naive. See `docs/disjunction-design.md`.
 - Add typed `EGraph` extension state that clones with `EGraph` and is restored by `push`/`pop`.
 - Report full source file paths in egglog span and error messages.
 - Fix seminaive matching after nested containers rebuild in place by propagating dirty container ids through parent containers.

core-relations/src/free_join/execute.rs

@@ -1800,6 +1800,58 @@ impl<'a> JoinState<'a> {
                     binding_info.move_back(spec.0.to_index.atom, prober);
                 }
             }
+            JoinStage::Union { branches } => {
+                // Fused disjunction. Enumerate every branch independently — each
+                // over its own atoms with a fresh `BindingInfo` — handing each
+                // complete branch match to the same `action`/`action_buf` as the
+                // enclosing plan. In practice a `Union` is the leading stage of a
+                // materialization block, so `action`/`action_buf` accumulate the
+                // branches' `output_vars` into one materialization (deduplicated
+                // on the key). The surrounding conjunction is then joined once,
+                // in the plan's continuation, against that materialization.
+                for branch in branches {
+                    if self.exec_state.should_stop() {
+                        return;
+                    }
+                    let mut branch_bi = BindingInfo::default();
+                    for (id, info) in branch.atoms.iter() {
+                        let table = self.db.get_table(info.table);
+                        branch_bi.insert_subset(id, table.all());
+                    }
+                    // Copy over any bindings the enclosing plan has already
+                    // established (a `Union` is a leading stage today, so this
+                    // is usually empty, but keep it correct for future use).
+                    for (var, val) in binding_info.bindings.iter() {
+                        branch_bi.bindings.insert(var, *val);
+                    }
+                    let mut skip = false;
+                    for JoinHeader { atom, subset, .. } in &branch.header {
+                        if subset.is_empty() {
+                            skip = true;
+                            break;
+                        }
+                        let mut node = Arc::try_unwrap(branch_bi.unwrap_val(*atom)).unwrap();
+                        node.cached_subsets.take();
+                        node.cached_children.take();
+                        node.subset.intersect(subset.as_ref(), &self.pool);
+                        if node.subset.is_empty() {
+                            skip = true;
+                            break;
+                        }
+                        branch_bi.move_back_node(*atom, Arc::new(node));
+                    }
+                    if skip {
+                        continue;
+                    }
+                    self.run_join_stages(
+                        &branch.stages,
+                        &branch.atoms,
+                        action,
+                        &mut branch_bi,
+                        action_buf,
+                    );
+                }
+            }
         }
     }
 }
@@ -2272,6 +2324,9 @@ fn estimate_size(join_stage: &JoinStage, binding_info: &BindingInfo) -> usize {
             .unwrap_or(0),
         JoinStage::FusedIntersect { cover, .. } => binding_info.subsets[cover.to_index.atom].size(),
         JoinStage::FusedIntersectMat { cover, .. } => binding_info.materializations[*cover].len(), // TODO: len() might be expensive.
+        // A `Union` is always a lone leading stage in its block, so its size
+        // estimate is never used to reorder anything.
+        JoinStage::Union { .. } => usize::MAX,
     }
 }
 
@@ -2280,6 +2335,7 @@ fn num_intersected_rels(join_stage: &JoinStage) -> i32 {
         JoinStage::Intersect { scans, .. } => scans.len() as i32,
         JoinStage::FusedIntersect { to_intersect, .. } => to_intersect.len() as i32 + 1,
         JoinStage::FusedIntersectMat { to_intersect, .. } => to_intersect.len() as i32,
+        JoinStage::Union { branches, .. } => branches.len() as i32,
     }
 }
 
@@ -2390,6 +2446,9 @@ fn recompute_leaf_scans(
                         break;
                     }
                 }
+                // A `Union` is always a lone stage in its block, so it never
+                // appears alongside the stages considered here.
+                JoinStage::Union { .. } => {}
             }
         }
         leaf_scans[i] = !blocked;
@@ -2433,6 +2492,7 @@ fn sort_plan_by_size_inner(
                         spec.to_index.vars.len() as i64;
                 });
             }
+            JoinStage::Union { .. } => {}
         }
     }
 
@@ -2459,6 +2519,7 @@ fn sort_plan_by_size_inner(
                 .copied()
                 .unwrap_or_default(),
             JoinStage::FusedIntersectMat { bind, .. } => bind.len() as _,
+            JoinStage::Union { .. } => 0,
         };
         (
             -refine,
@@ -2500,6 +2561,7 @@ fn sort_plan_by_size_inner(
                         spec.to_index.vars.len() as i64;
                 });
             }
+            JoinStage::Union { .. } => {}
         }
     }
 }

core-relations/src/free_join/plan.rs

@@ -60,7 +60,7 @@ use crate::{
     common::{HashMap, HashSet, IndexSet},
     offsets::Subset,
     pool::Pooled,
-    query::{Atom, Query, VarColumnMap},
+    query::{Atom, Query, UnionSpec, VarColumnMap},
     table_spec::Constraint,
 };
 
@@ -154,6 +154,24 @@ pub(crate) enum JoinStage {
         bind: SmallVec<[(ColumnId, Variable); 2]>,
         to_intersect: Vec<(ScanSpec, SmallVec<[ColumnId; 2]>)>,
     },
+    /// A fused disjunction (`or`). Each branch is a self-contained sub-plan over
+    /// its own atoms. At runtime every branch is enumerated, and each branch
+    /// match is handed to the enclosing block's materializer, which keys it on
+    /// the union's output variables (the block's [`MatSpec::msg_vars`]) —
+    /// deduplicating them. The surrounding conjunction is then joined once, in
+    /// the plan's result block, against that materialization via the atoms'
+    /// indexes rather than being re-scanned per branch. `Union` is only ever
+    /// emitted as the lone stage of a materialization block.
+    Union { branches: Vec<UnionBranch> },
+}
+
+/// One branch of a [`JoinStage::Union`]: a self-contained sub-plan over the
+/// branch's own atoms.
+#[derive(Debug, Clone)]
+pub(crate) struct UnionBranch {
+    pub atoms: Arc<DenseIdMap<AtomId, Atom>>,
+    pub header: Vec<JoinHeader>,
+    pub stages: JoinStages,
 }
 
 /// Merge every `FusedIntersect { to_intersect: [] }` into the first earlier such stage on the
@@ -375,6 +393,9 @@ impl SinglePlan {
                 } => {
                     todo!("materialization")
                 }
+                JoinStage::Union { .. } => {
+                    todo!("union")
+                }
             };
             let next = if i == self.stages.instrs.len() - 1 {
                 vec![]
@@ -1180,6 +1201,146 @@ pub(crate) fn tree_decompose_and_plan(
     })
 }
 
+/// Builds a `PlanningContext` restricted to `atom_ids`, keeping only the
+/// variables that occur in those atoms and trimming each variable's
+/// occurrences to those atoms.
+fn restrict_context<'a>(ctx: &PlanningContext<'a>, atom_ids: &[AtomId]) -> PlanningContext<'a> {
+    let atoms: DenseIdMap<AtomId, Atom> = atom_ids
+        .iter()
+        .map(|id| (*id, ctx.atoms[*id].clone()))
+        .collect();
+    let mut vars = DenseIdMap::new();
+    for (var, vinfo) in ctx.vars.iter() {
+        let mut vinfo = vinfo.clone();
+        vinfo.occurrences.retain(|occ| atoms.contains_key(occ.atom));
+        if !vinfo.occurrences.is_empty() {
+            vars.insert(var, vinfo);
+        }
+    }
+    PlanningContext {
+        vars,
+        atoms,
+        fun_deps: ctx.fun_deps.clone(),
+        col_est: ctx.col_est.clone(),
+    }
+}
+
+/// Plans a query containing a disjunction (`or`) as a fused union. The result is
+/// a [`DecomposedPlan`] with a single block — a [`JoinStage::Union`] that
+/// enumerates every branch and materializes the `output_vars` (deduplicated) —
+/// followed by a result block that joins the surrounding conjunction against
+/// that materialization exactly once. See [`JoinStage::Union`].
+fn plan_union(
+    ctx: PlanningContext,
+    strat: PlanStrategy,
+    actions: ActionId,
+    spec: UnionSpec,
+) -> Plan {
+    let UnionSpec {
+        branch_atoms,
+        output_vars,
+    } = spec;
+
+    // Plan each branch over its own atoms. Every output variable must be bound
+    // by the branch (so it can be materialized), even if the branch itself does
+    // not otherwise use it, so mark them as used.
+    let branches: Vec<UnionBranch> = branch_atoms
+        .iter()
+        .map(|atom_ids| {
+            let mut branch_ctx = restrict_context(&ctx, atom_ids);
+            for var in output_vars.iter() {
+                if let Some(vinfo) = branch_ctx.vars.get_mut(*var) {
+                    vinfo.used_in_rhs = true;
+                }
+            }
+            let (header, instrs) = plan_stages(&branch_ctx, strat);
+            UnionBranch {
+                atoms: Arc::new(branch_ctx.atoms),
+                header,
+                stages: JoinStages {
+                    instrs: Arc::new(instrs),
+                },
+            }
+        })
+        .collect();
+
+    let output_vars: SmallVec<[Variable; 8]> = output_vars.into_iter().collect();
+
+    // Block 0 is the union: a lone `Union` stage whose "action" (when run as a
+    // materialization block) keys the branch outputs on `output_vars`.
+    let union_stages = JoinStages {
+        instrs: Arc::new(vec![JoinStage::Union { branches }]),
+    };
+    let union_spec = MatSpec {
+        msg_vars: output_vars.iter().copied().collect(),
+        val_vars: smallvec![],
+    };
+    let blocks = vec![(union_stages, union_spec)];
+
+    // The continuation joins the atoms not in any branch, receiving
+    // `output_vars` as message variables from block 0. `plan_single_bag`
+    // produces exactly this: a prologue that scans the materialization and
+    // probes the continuation atoms by `output_vars`, followed by the join of
+    // the remaining continuation variables. Running those stages as the result
+    // block (with the rule's action) fires the action per match.
+    let cont_atom_ids: Vec<AtomId> = ctx
+        .atoms
+        .iter()
+        .map(|(id, _)| id)
+        .filter(|id| !branch_atoms.iter().any(|b| b.contains(id)))
+        .collect();
+    let mut cont_ctx = restrict_context(&ctx, &cont_atom_ids);
+    // The continuation must know about `output_vars` so the message-var prologue
+    // can match, even for output vars that appear only in the branches.
+    for var in output_vars.iter() {
+        if !cont_ctx.vars.contains_key(*var) {
+            cont_ctx.vars.insert(
+                *var,
+                VarInfo {
+                    occurrences: Default::default(),
+                    used_in_rhs: true,
+                    defined_in_rhs: false,
+                    name: None,
+                },
+            );
+        }
+    }
+
+    let mut n_used_in_bag = count_variable_usage_per_bag(std::slice::from_ref(&cont_ctx));
+    // Bump usage for output vars so `plan_single_bag` treats them as coming from
+    // the (single) prior block rather than pruning them.
+    for var in output_vars.iter() {
+        let count = n_used_in_bag.get_or_default(*var);
+        *count += 1;
+    }
+    let mut has_block_contributed = vec![false];
+    let (cont_header, result_block, _cont_spec) = plan_single_bag(
+        &mut cont_ctx,
+        &blocks,
+        &mut has_block_contributed,
+        &mut n_used_in_bag,
+        strat,
+    );
+
+    let blocks = blocks
+        .into_iter()
+        .map(|(stages, mat_spec)| (loop_lifting(stages), mat_spec))
+        .collect::<Vec<_>>();
+    let result_block = loop_lifting(result_block);
+
+    // The plan's atom set is the continuation atoms only; each branch carries
+    // its own atoms inside the `Union` stage. This keeps the execution-time
+    // `BindingInfo` (seeded from these atoms) free of branch atoms, so an empty
+    // branch relation does not abort the whole rule.
+    Plan::DecomposedPlan(DecomposedPlan {
+        atoms: Arc::new(cont_ctx.atoms),
+        header: cont_header,
+        stages: JoinStageBlocks { blocks },
+        result_block,
+        actions,
+    })
+}
+
 pub(crate) fn plan_query<'a>(query: Query, col_est: ColumnCardEst<'a>) -> Plan {
     let atoms = query.atoms;
     let ctx = PlanningContext {
@@ -1188,6 +1349,9 @@ pub(crate) fn plan_query<'a>(query: Query, col_est: ColumnCardEst<'a>) -> Plan {
         fun_deps: Arc::new(query.fun_deps),
         col_est,
     };
+    if let Some(union) = query.union {
+        return plan_union(ctx, query.plan_strategy, query.action, union);
+    }
     tree_decompose_and_plan(ctx, query.plan_strategy, query.action, query.no_decomp)
 }