Unify Continuations (#45)

connortsui20 · web-flow · commit 0f6379feaa1a · 2025-03-26T20:59:44.000-04:00
## Problem

There were a lot of type aliases that looked identical, so I removed
them for this:

```rust
pub type Continuation&lt;Input&gt; = Arc&lt;dyn Fn(Input) -&gt; UnitFuture + Send + Sync + 'static&gt;;
```

## Summary of changes

Removed all of the `???Continuation` type aliases

I also lifted the `Generator` trait into the same module as `engine`
diff --git a/optd-core/src/engine/eval/core.rs b/optd-core/src/engine/eval/core.rs
@@ -3,8 +3,8 @@ use super::{
     Engine, Evaluate, Generator,
 };
 use crate::capture;
-use crate::engine::generator::Continuation;
 use crate::engine::utils::evaluate_sequence;
+use crate::engine::Continuation;
 use optd_dsl::analyzer::hir::{CoreData, Expr, Value};
 use std::sync::Arc;
 use CoreData::*;
@@ -22,7 +22,7 @@ use CoreData::*;
 pub(super) async fn evaluate_core_expr<G>(
     data: CoreData<Arc<Expr>>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -75,7 +75,7 @@ async fn evaluate_collection<G, F>(
     items: Vec<Arc<Expr>>,
     constructor: F,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
     F: FnOnce(Vec<Value>) -> CoreData<Value> + Clone + Send + Sync + 'static,
@@ -100,8 +100,11 @@ async fn evaluate_collection<G, F>(
 /// * `items` - The key-value pairs to evaluate.
 /// * `engine` - The evaluation engine.
 /// * `k` - The continuation to receive evaluation results.
-async fn evaluate_map<G>(items: Vec<(Arc<Expr>, Arc<Expr>)>, engine: Engine<G>, k: Continuation)
-where
+async fn evaluate_map<G>(
+    items: Vec<(Arc<Expr>, Arc<Expr>)>,
+    engine: Engine<G>,
+    k: Continuation<Value>,
+) where
     G: Generator,
 {
     // Extract keys and values.
@@ -139,7 +142,7 @@ where
 /// * `msg` - The message expression to evaluate
 /// * `engine` - The evaluation engine
 /// * `k` - The continuation to receive evaluation results
-async fn evaluate_fail<G>(msg: Arc<Expr>, engine: Engine<G>, k: Continuation)
+async fn evaluate_fail<G>(msg: Arc<Expr>, engine: Engine<G>, k: Continuation<Value>)
 where
     G: Generator,
 {
diff --git a/optd-core/src/engine/eval/expr.rs b/optd-core/src/engine/eval/expr.rs
@@ -1,11 +1,7 @@
 use super::{binary::eval_binary_op, unary::eval_unary_op, Evaluate};
 use crate::{
     capture,
-    engine::{
-        generator::{Continuation, Generator},
-        utils::evaluate_sequence,
-        Engine,
-    },
+    engine::{utils::evaluate_sequence, Continuation, Engine, Generator},
 };
 use optd_dsl::analyzer::hir::{
     BinOp, CoreData, Expr, FunKind, Identifier, Literal, UnaryOp, Value,
@@ -31,7 +27,7 @@ pub(super) async fn evaluate_if_then_else<G>(
     then_expr: Arc<Expr>,
     else_expr: Arc<Expr>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -73,7 +69,7 @@ pub(super) async fn evaluate_let_binding<G>(
     assignee: Arc<Expr>,
     after: Arc<Expr>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -111,7 +107,7 @@ pub(super) async fn evaluate_binary_expr<G>(
     op: BinOp,
     right: Arc<Expr>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -121,7 +117,7 @@ pub(super) async fn evaluate_binary_expr<G>(
         right: Arc<Expr>,
         op: BinOp,
         engine: Engine<G>,
-        k: Continuation,
+        k: Continuation<Value>,
     ) where
         G: Generator,
     {
@@ -165,7 +161,7 @@ pub(super) async fn evaluate_unary_expr<G>(
     op: UnaryOp,
     expr: Arc<Expr>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -198,7 +194,7 @@ pub(super) async fn evaluate_function_call<G>(
     fun: Arc<Expr>,
     args: Vec<Arc<Expr>>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -242,7 +238,7 @@ pub(super) async fn evaluate_closure_call<G>(
     body: Arc<Expr>,
     args: Vec<Arc<Expr>>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -282,7 +278,7 @@ pub(super) async fn evaluate_rust_udf_call<G>(
     udf: fn(Vec<Value>) -> Value,
     args: Vec<Arc<Expr>>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -311,7 +307,7 @@ pub(super) async fn evaluate_rust_udf_call<G>(
 /// * `ident` - The identifier to look up
 /// * `engine` - The evaluation engine
 /// * `k` - The continuation to receive the variable value
-pub(super) async fn evaluate_reference<G>(ident: String, engine: Engine<G>, k: Continuation)
+pub(super) async fn evaluate_reference<G>(ident: String, engine: Engine<G>, k: Continuation<Value>)
 where
     G: Generator,
 {
diff --git a/optd-core/src/engine/eval/match.rs b/optd-core/src/engine/eval/match.rs
@@ -1,6 +1,8 @@
 use super::{Engine, Evaluate, Generator};
-use crate::engine::generator::Continuation;
-use crate::{capture, engine::UnitFuture};
+use crate::{
+    capture,
+    engine::{Continuation, UnitFuture},
+};
 use optd_dsl::analyzer::{
     context::Context,
     hir::{
@@ -39,7 +41,7 @@ pub(super) async fn evaluate_pattern_match<G>(
     expr: Arc<Expr>,
     match_arms: Vec<MatchArm>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -74,7 +76,7 @@ fn try_match_arms<G>(
     value: Value,
     match_arms: Vec<MatchArm>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) -> UnitFuture
 where
     G: Generator,
diff --git a/optd-core/src/engine/eval/mod.rs b/optd-core/src/engine/eval/mod.rs
@@ -1,11 +1,10 @@
-use super::generator::{Continuation, Generator};
-use super::{Engine, UnitFuture};
+use super::{Continuation, Engine, Generator, UnitFuture};
 use core::evaluate_core_expr;
 use expr::{
     evaluate_binary_expr, evaluate_function_call, evaluate_if_then_else, evaluate_let_binding,
     evaluate_reference, evaluate_unary_expr,
 };
-use optd_dsl::analyzer::hir::Expr;
+use optd_dsl::analyzer::hir::{Expr, Value};
 use r#match::evaluate_pattern_match;
 use std::sync::Arc;
 use Expr::*;
@@ -26,13 +25,13 @@ pub trait Evaluate {
     /// * `self` - The expression to evaluate
     /// * `engine` - The evaluation engine (owned)
     /// * `k` - The continuation to receive each evaluation result
-    fn evaluate<G>(self, engine: Engine<G>, k: Continuation) -> UnitFuture
+    fn evaluate<G>(self, engine: Engine<G>, k: Continuation<Value>) -> UnitFuture
     where
         G: Generator;
 }
 
 impl Evaluate for Arc<Expr> {
-    fn evaluate<G>(self, engine: Engine<G>, k: Continuation) -> UnitFuture
+    fn evaluate<G>(self, engine: Engine<G>, k: Continuation<Value>) -> UnitFuture
     where
         G: Generator,
     {
diff --git a/optd-core/src/engine/eval/operator.rs b/optd-core/src/engine/eval/operator.rs
@@ -1,7 +1,7 @@
 use super::{Engine, Generator};
-use crate::engine::generator::Continuation;
+use crate::capture;
 use crate::engine::utils::evaluate_sequence;
-use crate::{capture, engine::utils::UnitFuture};
+use crate::engine::Continuation;
 use optd_dsl::analyzer::hir::{
     CoreData, Expr, LogicalOp, Materializable, Operator, PhysicalOp, Value,
 };
@@ -10,7 +10,7 @@ use CoreData::{Logical, Physical};
 use Materializable::*;
 
 /// Specialized continuation type for operator values
-type OperatorContinuation = Arc<dyn Fn(Operator<Value>) -> UnitFuture + Send + Sync + 'static>;
+type OperatorContinuation = Continuation<Operator<Value>>;
 
 /// Evaluates a logical operator by generating all possible combinations of its components.
 ///
@@ -22,7 +22,7 @@ type OperatorContinuation = Arc<dyn Fn(Operator<Value>) -> UnitFuture + Send + S
 pub(super) async fn evaluate_logical_operator<G>(
     op: LogicalOp<Arc<Expr>>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
@@ -62,7 +62,7 @@ pub(super) async fn evaluate_logical_operator<G>(
 pub(super) async fn evaluate_physical_operator<G>(
     op: PhysicalOp<Arc<Expr>>,
     engine: Engine<G>,
-    k: Continuation,
+    k: Continuation<Value>,
 ) where
     G: Generator,
 {
diff --git a/optd-core/src/engine/generator.rs b/optd-core/src/engine/generator.rs
@@ -1,10 +1,5 @@
+use super::Continuation;
 use optd_dsl::analyzer::hir::{Goal, GroupId, Value};
-use std::sync::Arc;
-
-use super::utils::UnitFuture;
-
-/// A continuation function that processes a Value and returns a Future.
-pub type Continuation = Arc<dyn Fn(Value) -> UnitFuture + Send + Sync + 'static>;
 
 /// Defines operations for expanding references in the query plan using CPS.
 ///
@@ -21,7 +16,7 @@ pub trait Generator: Clone + Send + Sync + 'static {
     /// # Parameters
     /// * `group_id` - The ID of the group to expand
     /// * `k` - The continuation to process each expression in the group
-    async fn yield_group(&self, group_id: GroupId, k: Continuation);
+    async fn yield_group(&self, group_id: GroupId, k: Continuation<Value>);
 
     /// Expands a physical goal and passes each implementation to the continuation.
     ///
@@ -31,5 +26,5 @@ pub trait Generator: Clone + Send + Sync + 'static {
     /// # Parameters
     /// * `physical_goal` - The goal describing required properties
     /// * `k` - The continuation to process each implementation
-    async fn yield_goal(&self, physical_goal: &Goal, k: Continuation);
+    async fn yield_goal(&self, physical_goal: &Goal, k: Continuation<Value>);
 }
diff --git a/optd-core/src/engine/mod.rs b/optd-core/src/engine/mod.rs
@@ -11,8 +11,7 @@ use crate::{
     capture,
     cir::{Cost, LogicalProperties, PartialLogicalPlan, PartialPhysicalPlan, PhysicalProperties},
 };
-use eval::Evaluate;
-use generator::Generator;
+
 use optd_dsl::analyzer::{
     context::Context,
     hir::{CoreData, Expr, Literal, Value},
@@ -22,25 +21,22 @@ use utils::UnitFuture;
 use Expr::*;
 
 mod eval;
-pub(crate) mod generator;
-#[cfg(test)]
-pub(super) mod test_utils;
-pub(crate) mod utils;
+use eval::Evaluate;
 
-/// Type alias for a continuation that receives a PartialLogicalPlan
-pub(crate) type LogicalPlanContinuation =
-    Arc<dyn Fn(PartialLogicalPlan) -> UnitFuture + Send + Sync + 'static>;
+mod generator;
+pub use generator::Generator;
 
-/// Type alias for a continuation that receives a PartialPhysicalPlan
-pub(crate) type PhysicalPlanContinuation =
-    Arc<dyn Fn(PartialPhysicalPlan) -> UnitFuture + Send + Sync + 'static>;
+pub(crate) mod utils;
 
-/// Type alias for a continuation that receives a cost value
-pub(crate) type CostContinuation = Arc<dyn Fn(Cost) -> UnitFuture + Send + Sync + 'static>;
+#[cfg(test)]
+pub(super) mod test_utils;
 
-/// Type alias for a continuation that receives LogicalProperties
-pub(crate) type PropertiesContinuation =
-    Arc<dyn Fn(LogicalProperties) -> UnitFuture + Send + Sync + 'static>;
+/// A type alias for continuations used in the rule engine.
+///
+/// The engine uses continuation-passing-style (CPS) since it requires advanced control flow to
+/// expand and iterate over expressions within groups (where each expression requires
+/// plan-dependent state).
+pub type Continuation<Input> = Arc<dyn Fn(Input) -> UnitFuture + Send + Sync + 'static>;
 
 /// The engine for evaluating HIR expressions and applying rules.
 #[derive(Debug, Clone)]
@@ -87,7 +83,7 @@ impl<G: Generator> Engine<G> {
         self,
         rule_name: String,
         plan: &PartialLogicalPlan,
-        k: LogicalPlanContinuation,
+        k: Continuation<PartialLogicalPlan>,
     ) -> UnitFuture {
         let rule_call = self.create_rule_call(&rule_name, vec![partial_logical_to_value(plan)]);
 
@@ -126,7 +122,7 @@ impl<G: Generator> Engine<G> {
         rule_name: String,
         plan: &PartialLogicalPlan,
         props: &PhysicalProperties,
-        k: PhysicalPlanContinuation,
+        k: Continuation<PartialPhysicalPlan>,
     ) -> UnitFuture {
         let plan_value = partial_logical_to_value(plan);
         let props_value = physical_properties_to_value(props);
@@ -164,7 +160,7 @@ impl<G: Generator> Engine<G> {
     pub(crate) fn launch_cost_plan(
         self,
         plan: &PartialPhysicalPlan,
-        k: CostContinuation,
+        k: Continuation<Cost>,
     ) -> UnitFuture {
         // Create a call to the reserved "cost" function
         let rule_call = self.create_rule_call("cost", vec![partial_physical_to_value(plan)]);
@@ -194,7 +190,7 @@ impl<G: Generator> Engine<G> {
     pub(crate) fn launch_derive_properties(
         self,
         plan: &PartialLogicalPlan,
-        k: PropertiesContinuation,
+        k: Continuation<LogicalProperties>,
     ) -> UnitFuture {
         // Create a call to the reserved "derive" function
         let rule_call = self.create_rule_call("derive", vec![partial_logical_to_value(plan)]);
diff --git a/optd-core/src/engine/test_utils.rs b/optd-core/src/engine/test_utils.rs
@@ -1,7 +1,4 @@
-use crate::engine::{
-    generator::{Continuation, Generator},
-    Engine, Evaluate,
-};
+use crate::engine::{generator::Generator, Continuation, Engine, Evaluate};
 use optd_dsl::analyzer::hir::{
     CoreData, Expr, Goal, GroupId, Literal, LogicalOp, MatchArm, Materializable, Operator, Pattern,
     PhysicalOp, Value,
@@ -44,7 +41,7 @@ impl MockGenerator {
 }
 
 impl Generator for MockGenerator {
-    async fn yield_group(&self, group_id: GroupId, k: Continuation) {
+    async fn yield_group(&self, group_id: GroupId, k: Continuation<Value>) {
         let key = format!("{:?}", group_id);
         let values = {
             let mappings = self.group_mappings.lock().unwrap();
@@ -56,7 +53,7 @@ impl Generator for MockGenerator {
         }
     }
 
-    async fn yield_goal(&self, physical_goal: &Goal, k: Continuation) {
+    async fn yield_goal(&self, physical_goal: &Goal, k: Continuation<Value>) {
         let key = format!(
             "{:?}:{:?}",
             physical_goal.group_id, physical_goal.properties
diff --git a/optd-core/src/optimizer/generator.rs b/optd-core/src/optimizer/generator.rs
diff --git a/optd-core/src/optimizer/handlers.rs b/optd-core/src/optimizer/handlers.rs
diff --git a/optd-core/src/optimizer/ingest.rs b/optd-core/src/optimizer/ingest.rs
diff --git a/optd-core/src/optimizer/subscriptions.rs b/optd-core/src/optimizer/subscriptions.rs
