refactor: Stage 1 → Stage 2 boundary types in asap_types::stages (#1047 P5a)

crates/devtools/src/bin/stage_pipeline.rs

@@ -278,13 +278,9 @@ use asap_frontend_promql::{counter_function_over_gauge_warnings, promql_metric_n
 use asap_logical_asap_optimization::pass1::per_computation_candidates::{
     HowToComputeAggregateMeasure, PerComputationCandidates, PruneReason,
 };
-use asap_logical_asap_optimization::pass2::memo::candidates::{
-    CandidateMemo, OriginOfCandidateForMemoGroup,
-};
 use asap_logical_asap_optimization::pass2::memo::enumeration::{
     build_logical_asap_dags_from_chosen_candidates, enumerate, ChosenCandidatePerMemoGroup,
 };
-use asap_logical_asap_optimization::pass2::memo::{queries_rooted_at, MemoGroupId};
 use asap_logical_asap_optimization::DeploymentCapabilities;
 use asap_plan_selection::{
     large_peak_memory_warning, PlanningModels, Rejection, Selection, SelectionError,
@@ -295,6 +291,8 @@ use asap_types::ir::flat::{flatten, FlatDag, FlatNode, NodeId};
 use asap_types::ir::operator::node::OperatorNode;
 use asap_types::ir::properties::{default_metric_for_root, AccuracyMetric, AccuracyTarget};
 use asap_types::ir::QueryRoot;
+use asap_types::stages::{queries_rooted_at, MemoGroupId};
+use asap_types::stages::{CandidateMemo, OriginOfCandidateForMemoGroup};
 use asap_types::workload::{
     BatchEntry, DataArrival, DataWorkload, DurationMs, PlanningWorkload, Query, QueryLanguage,
     QueryRecurrence, QueryRequirements, QueryWorkload, QueryWorkloadEntry, RepeatedDemand,

crates/devtools/src/memo_pricing.rs

@@ -9,15 +9,15 @@
 use std::collections::HashMap;
 use std::rc::Rc;
 
-use asap_logical_asap_optimization::pass2::memo::candidates::CandidateMemo;
 use asap_logical_asap_optimization::pass2::memo::enumeration::{
     build_logical_asap_dags_from_chosen_candidates, ChosenCandidatePerMemoGroup,
 };
-use asap_logical_asap_optimization::pass2::memo::MemoGroupId;
 use asap_physical_asap_optimization::optimize_physical_asap;
 use asap_plan_selection::{select_plan, PlanningModels, Selection, SelectionError};
 use asap_planner::{installable_plan_of, InstallablePlanError};
 use asap_types::ir::{OperatorNode, QueryRoot};
+use asap_types::stages::CandidateMemo;
+use asap_types::stages::MemoGroupId;
 use asap_types::workload::{DataWorkload, QueryWorkloadEntry};
 
 /// One choice's result: valid with its cost, or why it is not valid.

crates/devtools/tests/memo_pricing.rs

@@ -7,12 +7,13 @@ use asap_devtools::planner_layering_examples::example1;
 use asap_devtools::planning_stages::{lower_promql_roots, run_default_search_over};
 use asap_logical_asap_optimization::deployment_capabilities::DeploymentCapabilities;
 use asap_logical_asap_optimization::pass1::per_computation_candidates::generate_per_computation_candidates;
-use asap_logical_asap_optimization::pass2::memo::candidates::{try_populate_pass1, CandidateMemo};
+use asap_logical_asap_optimization::pass2::memo::candidates::try_populate_pass1;
 use asap_logical_asap_optimization::pass2::memo::enumeration::enumerate;
 use asap_logical_asap_optimization::pass2::memo::DagOfLogicalMemoGroups;
 use asap_plan_selection::PlanningModels;
 use asap_types::ir::properties::AccuracyTarget;
 use asap_types::ir::QueryRoot;
+use asap_types::stages::CandidateMemo;
 use asap_types::workload::{
     DataWorkload, DurationMs, PlanningWorkload, Predictability, Query, QueryLanguage,
     QueryRequirements, QueryWorkload, QueryWorkloadEntry, RepeatedDemand, RepeatingEntry,

crates/integration-tests/src/lib.rs

@@ -307,7 +307,7 @@ pub mod planning {
     pub fn expanded_memo(
         roots: Vec<(usize, QueryRoot)>,
         entries: &[asap_types::workload::QueryWorkloadEntry],
-    ) -> asap_logical_asap_optimization::pass2::memo::candidates::CandidateMemo {
+    ) -> asap_types::stages::CandidateMemo {
         use asap_logical_asap_optimization::pass1::per_computation_candidates::generate_per_computation_candidates;
         use asap_logical_asap_optimization::pass2::memo::{
             candidates::try_populate_pass1, shared_configurations::PaneWidthDomain,
@@ -328,13 +328,12 @@ pub mod planning {
     /// each query's first root group's DAG, every group `choice` does not
     /// name implemented by its first candidate.
     pub fn compose_memo_plan(
-        memo: &asap_logical_asap_optimization::pass2::memo::candidates::CandidateMemo,
+        memo: &asap_types::stages::CandidateMemo,
         queries: usize,
         choice: &std::collections::BTreeMap<usize, usize>,
     ) -> Vec<Rc<OperatorNode>> {
-        use asap_logical_asap_optimization::pass2::memo::{
-            enumeration::build_logical_asap_dags_from_chosen_candidates, MemoGroupId,
-        };
+        use asap_logical_asap_optimization::pass2::memo::enumeration::build_logical_asap_dags_from_chosen_candidates;
+        use asap_types::stages::MemoGroupId;
         let choice = choice
             .iter()
             .map(|(&group, &index)| (MemoGroupId(group as u32), index))
@@ -402,11 +401,9 @@ pub mod planning {
     /// alone) candidates `keep` accepts; those of the topmost group (the
     /// fewest child edges below `q`'s root) first.
     pub fn own_choices(
-        memo: &asap_logical_asap_optimization::pass2::memo::candidates::CandidateMemo,
+        memo: &asap_types::stages::CandidateMemo,
         q: usize,
-        keep: impl Fn(
-            &asap_logical_asap_optimization::pass2::memo::candidates::CandidateForMemoGroup,
-        ) -> bool,
+        keep: impl Fn(&asap_types::stages::CandidateForMemoGroup) -> bool,
     ) -> Vec<(usize, usize)> {
         let mut found: Vec<_> = memo
             .candidates
@@ -430,9 +427,9 @@ pub mod planning {
             .iter()
             .enumerate()
             .filter(|(group, _)| {
-                memo.query_roots[q].operator_refs().contains(
-                    &&asap_logical_asap_optimization::pass2::memo::MemoGroupId(*group as u32),
-                )
+                memo.query_roots[q]
+                    .operator_refs()
+                    .contains(&&asap_types::stages::MemoGroupId(*group as u32))
             })
             .map(|(group, _)| (group, 0))
             .collect();
@@ -460,7 +457,7 @@ pub mod planning {
     /// - `shared{g}`: every query reads the summary state group `g`, its own
     ///   estimate or a regrouping merge of it, when every query can.
     pub fn select_rollup(
-        memo: &asap_logical_asap_optimization::pass2::memo::candidates::CandidateMemo,
+        memo: &asap_types::stages::CandidateMemo,
         queries: usize,
         entries: &[asap_types::workload::QueryWorkloadEntry],
         data: &DataWorkload,
@@ -469,8 +466,7 @@ pub mod planning {
         let own_builds: Vec<Vec<(usize, usize)>> = (0..queries)
             .map(|q| {
                 let found = own_choices(memo, q, |c| {
-                    c.origin
-                        == asap_logical_asap_optimization::pass2::memo::candidates::OriginOfCandidateForMemoGroup::Pass1
+                    c.origin == asap_types::stages::OriginOfCandidateForMemoGroup::Pass1
                         && builds(&c.root_node)
                 });
                 let top = found

crates/integration-tests/tests/finer_aggregate_reuse.rs

@@ -6,14 +6,12 @@ use asap_executor::values::Value;
 use asap_logical_asap_optimization::{
     pass1::per_computation_candidates::generate_per_computation_candidates,
     pass2::memo::{
-        candidates::{
-            try_populate_pass1, CandidateForMemoGroup, CandidateMemo, OriginOfCandidateForMemoGroup,
-        },
-        shared_configurations::PaneWidthDomain,
+        candidates::try_populate_pass1, shared_configurations::PaneWidthDomain,
         DagOfLogicalMemoGroups,
     },
 };
 use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode};
+use asap_types::stages::{CandidateForMemoGroup, CandidateMemo, OriginOfCandidateForMemoGroup};
 use serde_json::json;
 use std::path::Path;
 

crates/integration-tests/tests/memo_enumeration.rs

@@ -40,18 +40,16 @@ use asap_integration_tests::planning::expanded_memo;
 use asap_logical_asap_optimization::pass1::per_computation_candidates::{
     compose_logical_candidate, PerComputationCandidates,
 };
-use asap_logical_asap_optimization::pass2::memo::candidates::{
-    CandidateForMemoGroup, CandidateMemo, OriginOfCandidateForMemoGroup,
-};
 use asap_logical_asap_optimization::pass2::memo::enumeration::{
     build_logical_asap_dags_from_chosen_candidates, enumerate, ChosenCandidatePerMemoGroup,
 };
-use asap_logical_asap_optimization::pass2::memo::{queries_rooted_at, MemoGroupId};
 use asap_plan_selection::{
     run_stage1_to_stage3_with, MemoSteps, PlanningModels, Stage1ToStage3Run, Stage1ToStage3Search,
 };
 use asap_types::ir::schema::{FieldDataType, SketchParams};
 use asap_types::ir::{ASAPOp, OperatorNode, QueryRoot};
+use asap_types::stages::{queries_rooted_at, MemoGroupId};
+use asap_types::stages::{CandidateForMemoGroup, CandidateMemo, OriginOfCandidateForMemoGroup};
 use asap_types::workload::{PlanningWorkload, Query, RepeatedDemand, RepetitionInterval};
 
 const RATIO_OF_QUANTILES: &str = concat!(

crates/integration-tests/tests/planner_layering_example2.rs

@@ -291,11 +291,9 @@ fn pipeline_with_distinct_bound(bound: Option<u64>) -> Run {
 /// queries, those queries reading it.
 fn memo_pipeline() -> Run {
     use asap_integration_tests::planning::{compose_memo_plan, expanded_memo, own_choices};
-    use asap_logical_asap_optimization::pass2::memo::candidates::{
-        CandidateForMemoGroup, OriginOfCandidateForMemoGroup,
-    };
     use asap_types::ir::schema::FieldDataType;
     use asap_types::ir::{ASAPOp, OperatorNode};
+    use asap_types::stages::{CandidateForMemoGroup, OriginOfCandidateForMemoGroup};
     let workload = standin_workload();
     let roots = lower_promql(&workload);
     let entries: Vec<_> = workload.query_workload.entries().collect();

crates/integration-tests/tests/reuse_rules.rs

@@ -17,17 +17,15 @@ use asap_integration_tests::fixtures::lower_promql;
 use asap_integration_tests::planning::{
     cheapest_of, compose_memo_plan, expanded_memo, select_among,
 };
-use asap_logical_asap_optimization::pass2::memo::candidates::{
-    CandidateMemo, OriginOfCandidateForMemoGroup,
-};
-use asap_logical_asap_optimization::pass2::memo::{root_groups, MemoGroupId};
 use asap_types::ir::physical_export::compile_physical_asap_dag;
 use asap_types::ir::properties::AccuracyTarget;
 use asap_types::ir::{
     apply_materialization_timings, ASAPOp, MaterializationAssignment, NodeToTimedNodeCache,
     NonASAPOp, OperatorNode, QueryRoot,
 };
 use asap_types::physical::Materialization;
+use asap_types::stages::{root_groups, MemoGroupId};
+use asap_types::stages::{CandidateMemo, OriginOfCandidateForMemoGroup};
 use asap_types::workload::{
     DataArrival, DataWorkload, QueryRecurrence, QueryWorkloadEntry, RepeatedDemand,
     RepetitionInterval,

crates/integration-tests/tests/sharing_issue_regressions.rs

@@ -200,11 +200,11 @@ fn exact_rollup_catalog_crosses_identical_per_series_average() {
     use asap_logical_asap_optimization::{
         pass1::per_computation_candidates::generate_per_computation_candidates,
         pass2::memo::{
-            candidates::{try_populate_pass1, OriginOfCandidateForMemoGroup},
-            shared_configurations::PaneWidthDomain,
+            candidates::try_populate_pass1, shared_configurations::PaneWidthDomain,
             DagOfLogicalMemoGroups,
         },
     };
+    use asap_types::stages::OriginOfCandidateForMemoGroup;
     let roots: Vec<_> = Planner::new(FuzzSettings::default())
         .lower(&workload(923))
         .unwrap()
@@ -239,12 +239,12 @@ fn filtered_group_preserves_sql_empty_and_null_semantics() {
     use asap_logical_asap_optimization::{
         pass1::per_computation_candidates::generate_per_computation_candidates,
         pass2::memo::{
-            candidates::{try_populate_pass1, OriginOfCandidateForMemoGroup},
-            shared_configurations::PaneWidthDomain,
+            candidates::try_populate_pass1, shared_configurations::PaneWidthDomain,
             DagOfLogicalMemoGroups,
         },
     };
     use asap_types::ir::{operator::AggIntent, NonASAPOp};
+    use asap_types::stages::OriginOfCandidateForMemoGroup;
     let roots: Vec<_> = Planner::new(FuzzSettings::default())
         .lower(&workload(902))
         .unwrap()
@@ -400,12 +400,12 @@ fn sql_count_rollup_executes_with_integer_results() {
     use asap_logical_asap_optimization::{
         pass1::per_computation_candidates::generate_per_computation_candidates,
         pass2::memo::{
-            candidates::{try_populate_pass1, OriginOfCandidateForMemoGroup},
-            shared_configurations::PaneWidthDomain,
+            candidates::try_populate_pass1, shared_configurations::PaneWidthDomain,
             DagOfLogicalMemoGroups,
         },
     };
     use asap_types::ir::ASAPOp;
+    use asap_types::stages::OriginOfCandidateForMemoGroup;
     let mut w = workload(902);
     w.queries.truncate(2);
     w.queries[0].entry.query.0 =
@@ -529,11 +529,11 @@ fn sketch_rollup_crosses_inner_range_function() {
     use asap_logical_asap_optimization::{
         pass1::per_computation_candidates::generate_per_computation_candidates,
         pass2::memo::{
-            candidates::{try_populate_pass1, OriginOfCandidateForMemoGroup},
-            shared_configurations::PaneWidthDomain,
+            candidates::try_populate_pass1, shared_configurations::PaneWidthDomain,
             DagOfLogicalMemoGroups,
         },
     };
+    use asap_types::stages::OriginOfCandidateForMemoGroup;
     let mut w = workload(923);
     for (query, keys) in w.queries.iter_mut().zip(["job, region", "job"]) {
         query.entry.query.0 = format!("quantile by ({keys}) (0.95, sum_over_time(m_gauge[5m]))");

crates/logical-asap-optimization/src/pass1/logical_memo_candidates.rs

@@ -16,11 +16,10 @@ use asap_accuracy::propagation::with_resized_summaries;
 use crate::pass1::per_computation_candidates::{
     ComputationTarget, HowToComputeAggregateMeasure, PerComputationCandidates, PruneReason,
 };
-use crate::pass2::memo::candidates::{try_populate_pass1, CandidateMemo, CandidatePopulationError};
-use crate::pass2::memo::{
-    check_generation_deadline, CandidateGenerationTimeout, DagOfLogicalMemoGroups, MemoGroupId,
-};
+use crate::pass2::memo::candidates::{try_populate_pass1, CandidatePopulationError};
+use crate::pass2::memo::{check_generation_deadline, DagOfLogicalMemoGroups};
 use asap_accuracy::estimators::with_failure_probability;
+use asap_types::stages::{CandidateGenerationTimeout, CandidateMemo, MemoGroupId};
 
 /// One candidate a stage removed from group `group`, and why. `C` is the
 /// stage's candidate type.
@@ -37,7 +36,7 @@ pub struct PrunedCandidate<C> {
 pub struct PerComputationLogicalMemoCandidates {
     /// The groups moved from the [`DagOfLogicalMemoGroups`], its query
     /// roots, and every group's candidates (origin
-    /// [`Pass1`](crate::pass2::memo::candidates::OriginOfCandidateForMemoGroup::Pass1)).
+    /// [`Pass1`](asap_types::stages::OriginOfCandidateForMemoGroup::Pass1)).
     pub memo: CandidateMemo,
     /// One per aggregate measure: what Pass 1 chose among.
     pub computation_targets: Vec<ComputationTarget>,

crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs

@@ -661,7 +661,7 @@ pub fn generate_per_computation_candidates_with_memo<Id: std::fmt::Debug>(
 }
 
 type MemoGroupMeasureToComputationChoicesCache =
-    HashMap<(crate::pass2::memo::MemoGroupId, usize), ComputationChoices>;
+    HashMap<(asap_types::stages::MemoGroupId, usize), ComputationChoices>;
 
 fn generate_with_computation_choices_cache<Id: std::fmt::Debug>(
     roots: Vec<(Id, QueryRoot)>,

crates/logical-asap-optimization/src/pass2.rs

@@ -6,7 +6,7 @@
 pub mod asap_aware_cse;
 pub mod cse;
 pub(crate) mod finer_aggregate;
-/// Stage 1, Pass 2's entry and its output, [`logical_memo_candidates::LogicalASAPMemoCandidates`].
+/// Stage 1, Pass 2's entry, producing [`asap_types::stages::LogicalASAPMemoCandidates`].
 pub mod logical_memo_candidates;
 /// Logical group identity available before Pass 1 candidate generation.
 pub mod memo;

crates/logical-asap-optimization/src/pass2/logical_memo_candidates.rs

@@ -1,11 +1,11 @@
 //! Stage 1, Pass 2's entry and output: the Pass 1 memo extended with the
 //! candidates and groups that share work across computations
-//! ([`LogicalASAPMemoCandidates`]), which Stage 2 implements.
+//! ([`asap_types::stages::LogicalASAPMemoCandidates`]), which Stage 2 implements.
 
 use std::time::Instant;
 
 use crate::pass1::logical_memo_candidates::PerComputationLogicalMemoCandidates;
-use crate::pass2::memo::candidates::CandidateMemo;
+use crate::pass2::memo::check_generation_deadline;
 use crate::pass2::memo::rules::{
     finer_aggregate, rollup, scan_labels, shared_configuration_readers, summary_inputs,
     summary_results, topk,
@@ -15,20 +15,10 @@ use crate::pass2::memo::state_groups::{
     add_shared_summary_state_groups, check_unexpanded_memo, StateGroupExpansionError,
 };
 use crate::pass2::memo::window_candidates::extend_local_window_composition;
-use crate::pass2::memo::{check_generation_deadline, CandidateGenerationTimeout};
+use asap_types::stages::CandidateMemo;
+use asap_types::stages::{CandidateGenerationTimeout, LogicalASAPMemoCandidates};
 use asap_types::workload::QueryWorkloadEntry;
 
-/// Stage 1, Pass 2's output: the Pass 1 groups plus the groups Pass 2 adds
-/// (shared producers, shared summary states), each with its Pass 1
-/// candidates plus every Pass 2 rule's (by `origin`), including the
-/// candidates reading a shared summary state.
-#[derive(Debug, Clone)]
-pub struct LogicalASAPMemoCandidates {
-    /// The expanded memo: groups, query roots (passed through from Pass 1)
-    /// and each group's candidates.
-    pub memo: CandidateMemo,
-}
-
 /// What Pass 2 keeps between pane-width rounds: the memo after local
 /// window composition and before the other rules, and the pane widths it
 /// already holds. Each round adds only the newly exposed widths to it,