tenferro_einsum/planning/

tree.rs

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::mem::{size_of, size_of_val};

use omeco::{
    CodeOptimizer, EinCode as OmecoEinCode, Initializer, NestedEinsum, ScoreFunction, TreeSA,
};

use crate::cache::{saturating_sum, vec_of_vec_retained_bytes, vec_retained_bytes};
use crate::planning::plan::{compile_step_plans, DiagPlan, GemmPlan, ReducePlan, StepPlan};
use crate::syntax::subscripts::Subscripts;
use crate::util::{build_size_dict, intermediate_subs};
use crate::{Error, Result};

/// A single step in the contraction sequence.
pub(crate) struct ContractionStep {
    pub(crate) left: usize,
    pub(crate) right: usize,
}

/// Public options for automatic contraction-path optimization.
///
/// The default planner uses TreeSA with a greedy initializer and zero annealing
/// iterations. This keeps the public API on a single optimizer family while
/// making the default behavior effectively "greedy-only".
///
/// # Determinism
///
/// While annealing is disabled (`niters == 0` or empty `betas`, as in the
/// default), the planned path is a function of the subscripts and shapes
/// alone: the same spec yields the same path in every process and on every
/// call. With an annealing schedule the path comes from omeco's TreeSA, whose
/// annealing is seeded but whose greedy initializer (omeco 0.2.6) breaks ties
/// in `HashMap` order, so annealed paths are not guaranteed to agree across
/// processes.
#[derive(Debug, Clone)]
pub struct ContractionOptimizerOptions {
    /// Inverse-temperature schedule for TreeSA.
    pub betas: Vec<f64>,
    /// Number of independent TreeSA trials.
    pub ntrials: usize,
    /// Annealing iterations per temperature level.
    pub niters: usize,
    /// Score function used by TreeSA.
    pub score: ScoreFunction,
}

impl Default for ContractionOptimizerOptions {
    fn default() -> Self {
        Self {
            betas: Vec::new(),
            ntrials: 1,
            niters: 0,
            score: ScoreFunction::default(),
        }
    }
}

impl ContractionOptimizerOptions {
    fn to_treesa(&self) -> TreeSA {
        TreeSA::new(
            self.betas.clone(),
            self.ntrials,
            self.niters,
            Initializer::Greedy,
            self.score.clone(),
        )
    }

    /// Whether TreeSA would run any annealing iteration.
    fn anneals(&self) -> bool {
        self.niters > 0 && !self.betas.is_empty()
    }

    pub(crate) fn validate(&self) -> Result<()> {
        if self.ntrials == 0 {
            return Err(Error::planning(
                "contraction optimizer ntrials must be at least 1",
            ));
        }
        if self.betas.iter().any(|value| value.is_nan()) {
            return Err(Error::planning(
                "contraction optimizer betas must not contain NaN",
            ));
        }
        if self.score.tc_weight.is_nan()
            || self.score.sc_weight.is_nan()
            || self.score.rw_weight.is_nan()
            || self.score.sc_target.is_nan()
        {
            return Err(Error::planning(
                "contraction optimizer score fields must not contain NaN",
            ));
        }
        Ok(())
    }
}

/// Contraction tree determining pairwise contraction order for N-ary einsum.
///
/// When contracting more than two tensors, the order in which pairwise
/// contractions are performed significantly affects performance.
/// `ContractionTree` encodes this order as a binary tree.
///
/// # Optimization
///
/// Use [`ContractionTree::optimize`] for automatic cost-based optimization
/// (e.g., greedy algorithm based on tensor sizes), or
/// [`ContractionTree::from_pairs`] for manual specification.
pub struct ContractionTree {
    /// Original subscripts.
    pub(crate) subscripts: Subscripts,
    /// Steps in the contraction (empty for single-tensor case).
    pub(crate) steps: Vec<ContractionStep>,
    /// Label → dimension size mapping.
    pub(crate) size_dict: HashMap<u32, usize>,
    /// Subscripts for each operand (0..input_count from input, then intermediates).
    pub(crate) operand_subs: Vec<Vec<u32>>,
    /// Pre-compiled step plans (cached to avoid recomputation per execute call).
    pub(crate) step_plans: Vec<StepPlan>,
}

impl fmt::Debug for ContractionTree {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ContractionTree")
            .field("input_count", &self.subscripts.inputs.len())
            .field("output_rank", &self.subscripts.output.len())
            .field("steps_len", &self.steps.len())
            .field("size_dict_len", &self.size_dict.len())
            .field("operand_subs_len", &self.operand_subs.len())
            .field("step_plans_len", &self.step_plans.len())
            .finish_non_exhaustive()
    }
}

impl ContractionTree {
    /// Automatically compute an optimized contraction order.
    ///
    /// Uses a cost-based heuristic (greedy algorithm) to determine
    /// the pairwise contraction sequence that minimizes total operation count.
    /// The path is deterministic: for fixed subscripts and shapes it is the
    /// same in every process (see [`ContractionOptimizerOptions`]).
    ///
    /// # Arguments
    ///
    /// * `subscripts` — Einsum subscripts for all tensors
    /// * `shapes` — Shape of each input tensor
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_einsum::{ContractionTree, Subscripts};
    ///
    /// let subs = Subscripts::parse("abcdef,bf,cf,df,ef->f").unwrap();
    /// let shapes = [&[2, 3, 2, 4, 3, 9][..], &[3, 9], &[2, 9], &[4, 9], &[3, 9]];
    /// let tree = ContractionTree::optimize(&subs, &shapes).unwrap();
    /// assert_eq!(tree.step_count(), 4);
    /// // Equal-cost candidates are broken by operand index, never by hashing.
    /// let again = ContractionTree::optimize(&subs, &shapes).unwrap();
    /// for step in 0..tree.step_count() {
    ///     assert_eq!(tree.step_pair(step), again.step_pair(step));
    /// }
    /// ```
    ///
    /// # Errors
    ///
    /// Returns [`Error::Validation`] when subscripts and shapes have different
    /// ranks or incompatible dimensions, or [`Error::Planning`] when no valid
    /// contraction order can be constructed.
    pub fn optimize(subscripts: &Subscripts, shapes: &[&[usize]]) -> Result<Self> {
        Self::optimize_with_options(subscripts, shapes, &ContractionOptimizerOptions::default())
    }

    /// Automatically compute an optimized contraction order with explicit
    /// planner options.
    ///
    /// For three or more operands with an annealing schedule (`niters > 0`
    /// and non-empty `betas`), this routes planning through TreeSA using the
    /// provided configuration. Without annealing, including the default
    /// options, TreeSA would return its greedy initializer unchanged, so the
    /// deterministic greedy planner runs directly and the path is identical in
    /// every process. One or two operands need no ordering search; their trees
    /// are built directly after validating the options.
    ///
    /// # Errors
    ///
    /// Returns [`Error::Validation`] for rank, shape, or dimension mismatches,
    /// or [`Error::Planning`] when planner options such as `ntrials` are
    /// invalid or no contraction order can be constructed.
    pub fn optimize_with_options(
        subscripts: &Subscripts,
        shapes: &[&[usize]],
        options: &ContractionOptimizerOptions,
    ) -> Result<Self> {
        options.validate()?;
        let input_count = subscripts.inputs.len();
        if input_count <= 1 {
            return Self::from_pairs(subscripts, shapes, &[]);
        }
        if input_count == 2 {
            return Self::from_pairs(subscripts, shapes, &[(0, 1)]);
        }

        let size_dict = build_size_dict(subscripts, shapes, None)?;
        // Without annealing TreeSA returns its greedy initializer unchanged,
        // and omeco 0.2.6's greedy breaks ties in per-process `HashMap` order
        // (#1963), so the deterministic greedy plans this case directly.
        let pairs = if !options.anneals() {
            optimize_self_greedy_pairs(subscripts, &size_dict)?
        } else if let Some(omeco_pairs) = optimize_omeco_pairs(subscripts, &size_dict, options)? {
            omeco_pairs
        } else {
            optimize_self_greedy_pairs(subscripts, &size_dict)?
        };
        Self::from_pairs(subscripts, shapes, &pairs)
    }

    /// Manually build a contraction tree from a pairwise contraction sequence.
    ///
    /// Each pair `(i, j)` specifies which two tensors (or intermediate results)
    /// to contract next. Intermediate results are assigned indices starting
    /// from the number of input tensors.
    ///
    /// # Arguments
    ///
    /// * `subscripts` — Einsum subscripts for all tensors
    /// * `shapes` — Shape of each input tensor
    /// * `pairs` — Ordered list of pairwise contractions
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_einsum::{ContractionTree, Subscripts};
    ///
    /// // Three tensors: A[ij] B[jk] C[kl] -> D[il]
    /// // Contract B and C first, then A with the result:
    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
    /// let shapes = [&[3, 4][..], &[4, 5], &[5, 6]];
    /// let tree = ContractionTree::from_pairs(
    ///     &subs,
    ///     &shapes,
    ///     &[(1, 2), (0, 3)],  // B*C -> T(index=3), then A*T -> D
    /// ).unwrap();
    /// ```
    ///
    /// # Errors
    ///
    /// Returns [`Error::Planning`] when the pair count, operand indices, or
    /// intermediate sequence is invalid, or [`Error::Validation`] when the
    /// supplied shapes do not match the subscripts.
    pub fn from_pairs(
        subscripts: &Subscripts,
        shapes: &[&[usize]],
        pairs: &[(usize, usize)],
    ) -> Result<Self> {
        let input_count = subscripts.inputs.len();
        let required_steps = input_count.saturating_sub(1);
        if pairs.len() != required_steps {
            return Err(Error::planning(format!(
                "explicit contraction path for {input_count} operands must have {required_steps} steps, got {}",
                pairs.len()
            )));
        }
        let size_dict = build_size_dict(subscripts, shapes, None)?;

        let mut operand_subs: Vec<Vec<u32>> = subscripts.inputs.clone();
        let mut live = vec![false; input_count + pairs.len()];
        for slot in live.iter_mut().take(input_count) {
            *slot = true;
        }
        let mut steps = Vec::new();

        for (step_idx, &(left, right)) in pairs.iter().enumerate() {
            let next_idx = input_count + step_idx;
            if left == right {
                return Err(Error::planning(format!(
                    "pair ({left}, {right}) must reference two distinct live operands"
                )));
            }
            if left >= next_idx || right >= next_idx {
                return Err(Error::planning(format!(
                    "pair ({left}, {right}) references non-existent operand"
                )));
            }
            if !live[left] || !live[right] {
                return Err(Error::planning(format!(
                    "pair ({left}, {right}) references an operand or intermediate that is no longer live"
                )));
            }

            // Labels needed by other live operands + final output
            let mut needed: HashSet<u32> = subscripts.output.iter().copied().collect();
            for (idx, subs) in operand_subs.iter().enumerate() {
                if idx != left && idx != right && live[idx] {
                    needed.extend(subs.iter().copied());
                }
            }

            let new_subs = intermediate_subs(&operand_subs[left], &operand_subs[right], &needed);
            operand_subs.push(new_subs);
            live[left] = false;
            live[right] = false;
            live[next_idx] = true;
            steps.push(ContractionStep { left, right });
        }

        let live_count = live.iter().filter(|&&is_live| is_live).count();
        if live_count != 1 {
            return Err(Error::planning(format!(
                "explicit contraction path must leave exactly one live result, got {live_count}"
            )));
        }

        let mut tree = Self {
            subscripts: subscripts.clone(),
            steps,
            size_dict,
            operand_subs,
            step_plans: Vec::new(),
        };
        tree.step_plans = compile_step_plans(&tree)?;
        Ok(tree)
    }

    /// Return the number of pairwise contraction steps in this tree.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_einsum::{ContractionTree, Subscripts};
    ///
    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
    /// let tree = ContractionTree::from_pairs(
    ///     &subs,
    ///     &[&[2, 2], &[2, 2], &[2, 2]],
    ///     &[(1, 2), (0, 3)],
    /// )
    /// .unwrap();
    /// assert_eq!(tree.step_count(), 2);
    /// ```
    #[must_use]
    pub fn step_count(&self) -> usize {
        self.steps.len()
    }

    pub(crate) fn label_size(&self, label: u32) -> Option<usize> {
        self.size_dict.get(&label).copied()
    }

    pub(crate) fn output_shape(&self) -> Vec<usize> {
        self.subscripts
            .output
            .iter()
            .filter_map(|label| self.label_size(*label))
            .collect()
    }

    /// Return the operand indices for a pairwise contraction step.
    ///
    /// The returned indices refer to the original inputs (`0..input_count`) and
    /// then to intermediates (`input_count..`) produced by earlier steps.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_einsum::{ContractionTree, Subscripts};
    ///
    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
    /// let tree = ContractionTree::from_pairs(
    ///     &subs,
    ///     &[&[2, 2], &[2, 2], &[2, 2]],
    ///     &[(1, 2), (0, 3)],
    /// )
    /// .unwrap();
    /// assert_eq!(tree.step_pair(0), Some((1, 2)));
    /// ```
    #[must_use]
    pub fn step_pair(&self, step_idx: usize) -> Option<(usize, usize)> {
        self.steps.get(step_idx).map(|step| (step.left, step.right))
    }

    /// Return the `(lhs, rhs, output)` subscripts for a pairwise step.
    ///
    /// The output subscripts are the intermediate labels preserved after the
    /// contraction, or the final output labels on the last step.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_einsum::{ContractionTree, Subscripts};
    ///
    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
    /// let tree = ContractionTree::from_pairs(
    ///     &subs,
    ///     &[&[2, 2], &[2, 2], &[2, 2]],
    ///     &[(1, 2), (0, 3)],
    /// )
    /// .unwrap();
    /// let (lhs, rhs, out) = tree.step_subscripts(0).unwrap();
    /// assert_eq!(lhs, &[1, 2]);
    /// assert_eq!(rhs, &[2, 3]);
    /// assert_eq!(out, &[1, 3]);
    /// ```
    #[must_use]
    pub fn step_subscripts(&self, step_idx: usize) -> Option<(&[u32], &[u32], &[u32])> {
        let input_count = self.subscripts.inputs.len();
        let step = self.steps.get(step_idx)?;
        let result_idx = input_count + step_idx;
        let output_subs = if step_idx + 1 == self.steps.len() {
            &self.subscripts.output
        } else {
            &self.operand_subs[result_idx]
        };
        Some((
            &self.operand_subs[step.left],
            &self.operand_subs[step.right],
            output_subs,
        ))
    }

    /// Return the precomputed lowering plan for one pairwise contraction step.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use tenferro_einsum::{ContractionTree, Subscripts};
    ///
    /// let subs = Subscripts::new(&[&[0, 1], &[1, 2]], &[0, 2]);
    /// let tree = ContractionTree::from_pairs(&subs, &[&[2, 3], &[3, 4]], &[(0, 1)]).unwrap();
    ///
    /// assert_eq!(tree.step_plan(0).unwrap().gemm().m(), 2);
    /// ```
    #[must_use]
    pub fn step_plan(&self, step_idx: usize) -> Option<crate::lowering::PairwiseStepPlan<'_>> {
        self.step_plans
            .get(step_idx)
            .map(crate::lowering::PairwiseStepPlan::new)
    }

    #[doc(hidden)]
    #[must_use]
    pub(crate) fn retained_bytes_for_cache_stats(&self) -> usize {
        saturating_sum([
            size_of::<Self>(),
            subscripts_retained_bytes(&self.subscripts),
            self.steps
                .capacity()
                .saturating_mul(size_of::<ContractionStep>()),
            self.size_dict
                .capacity()
                .saturating_mul(size_of::<u32>().saturating_add(size_of::<usize>())),
            vec_of_vec_retained_bytes(&self.operand_subs),
            self.step_plans
                .capacity()
                .saturating_mul(size_of::<StepPlan>()),
            saturating_sum(self.step_plans.iter().map(step_plan_retained_bytes)),
        ])
    }
}

fn subscripts_retained_bytes(subscripts: &Subscripts) -> usize {
    saturating_sum([
        vec_of_vec_retained_bytes(&subscripts.inputs),
        vec_retained_bytes(&subscripts.output),
    ])
}

fn reduce_plan_retained_bytes(plan: &ReducePlan) -> usize {
    saturating_sum([
        vec_retained_bytes(&plan.original_subs),
        vec_retained_bytes(&plan.kept_subs),
        vec_retained_bytes(&plan.out_shape),
    ])
}

fn diag_plan_retained_bytes(plan: &DiagPlan) -> usize {
    saturating_sum([
        vec_retained_bytes(&plan.stages),
        saturating_sum(plan.stages.iter().map(|stage| {
            saturating_sum([
                vec_retained_bytes(&stage.axis_pairs),
                vec_retained_bytes(&stage.result_subs),
            ])
        })),
        vec_retained_bytes(&plan.result_subs),
    ])
}

fn gemm_plan_retained_bytes(plan: &GemmPlan) -> usize {
    saturating_sum([
        plan.reduce_a.as_ref().map_or(0, reduce_plan_retained_bytes),
        plan.reduce_b.as_ref().map_or(0, reduce_plan_retained_bytes),
        vec_retained_bytes(&plan.subs_a),
        vec_retained_bytes(&plan.subs_b),
        vec_retained_bytes(&plan.lo_modes),
        vec_retained_bytes(&plan.ro_modes),
        vec_retained_bytes(&plan.sum_modes),
        vec_retained_bytes(&plan.lo_sizes),
        vec_retained_bytes(&plan.ro_sizes),
        vec_retained_bytes(&plan.sum_sizes),
        vec_retained_bytes(&plan.batch_sizes),
        vec_retained_bytes(&plan.target_a),
        vec_retained_bytes(&plan.target_b),
        vec_retained_bytes(&plan.c_gemm_shape),
        vec_retained_bytes(&plan.expanded_shape),
        vec_retained_bytes(&plan.canonical_modes),
        vec_retained_bytes(&plan.a_gemm_shape),
        vec_retained_bytes(&plan.b_gemm_shape),
    ])
}

fn step_plan_retained_bytes(plan: &StepPlan) -> usize {
    saturating_sum([
        plan.diag_a.as_ref().map_or(0, diag_plan_retained_bytes),
        plan.diag_b.as_ref().map_or(0, diag_plan_retained_bytes),
        plan.strict_binary.as_ref().map_or(0, size_of_val),
        gemm_plan_retained_bytes(&plan.gemm),
    ])
}

fn optimize_omeco_pairs(
    subscripts: &Subscripts,
    size_dict: &HashMap<u32, usize>,
    options: &ContractionOptimizerOptions,
) -> Result<Option<Vec<(usize, usize)>>> {
    #[cfg(test)]
    tests::OMECO_CALLS.with(|count| count.set(count.get() + 1));
    let code = OmecoEinCode::new(subscripts.inputs.clone(), subscripts.output.clone());
    let optimizer = options.to_treesa();
    let Some(nested) = optimizer.optimize(&code, size_dict) else {
        return Ok(None);
    };

    let mut next_operand = subscripts.inputs.len();
    let mut pairs = Vec::with_capacity(subscripts.inputs.len().saturating_sub(1));
    nested_to_pairs(&nested, &mut next_operand, &mut pairs)?;
    Ok(Some(pairs))
}

fn nested_to_pairs(
    nested: &NestedEinsum<u32>,
    next_operand: &mut usize,
    pairs: &mut Vec<(usize, usize)>,
) -> Result<usize> {
    match nested {
        NestedEinsum::Leaf { tensor_index } => Ok(*tensor_index),
        NestedEinsum::Node { args, .. } => {
            if args.len() != 2 {
                return Err(Error::planning(format!(
                    "omeco returned non-binary contraction node with {} children",
                    args.len()
                )));
            }
            let left = nested_to_pairs(&args[0], next_operand, pairs)?;
            let right = nested_to_pairs(&args[1], next_operand, pairs)?;
            pairs.push((left, right));
            let result_idx = *next_operand;
            *next_operand += 1;
            Ok(result_idx)
        }
    }
}

fn build_operand_label_sets(operand_subs: &[Vec<u32>]) -> Vec<HashSet<u32>> {
    operand_subs
        .iter()
        .map(|subs| subs.iter().copied().collect())
        .collect()
}

fn build_needed_label_counts(
    output_subs: &[u32],
    available: &[usize],
    operand_label_sets: &[HashSet<u32>],
) -> HashMap<u32, usize> {
    let mut counts = HashMap::new();
    for &label in output_subs {
        counts.entry(label).or_insert(1);
    }
    for &idx in available {
        add_labels_to_counts(&mut counts, &operand_label_sets[idx]);
    }
    counts
}

fn add_labels_to_counts(counts: &mut HashMap<u32, usize>, labels: &HashSet<u32>) {
    for &label in labels {
        *counts.entry(label).or_insert(0) += 1;
    }
}

fn remove_labels_from_counts(counts: &mut HashMap<u32, usize>, labels: &HashSet<u32>) {
    for &label in labels {
        match counts.get(&label).copied() {
            Some(1) => {
                counts.remove(&label);
            }
            Some(count) => {
                counts.insert(label, count - 1);
            }
            None => {}
        }
    }
}

fn candidate_label_is_needed(
    label: u32,
    left: usize,
    right: usize,
    operand_label_sets: &[HashSet<u32>],
    needed_label_counts: &HashMap<u32, usize>,
) -> bool {
    let mut selected_count = 0;
    if operand_label_sets[left].contains(&label) {
        selected_count += 1;
    }
    if operand_label_sets[right].contains(&label) {
        selected_count += 1;
    }
    needed_label_counts.get(&label).copied().unwrap_or(0) > selected_count
}

fn collect_candidate_intermediate_subs(
    subs_left: &[u32],
    subs_right: &[u32],
    left: usize,
    right: usize,
    operand_label_sets: &[HashSet<u32>],
    needed_label_counts: &HashMap<u32, usize>,
    output: &mut Vec<u32>,
) {
    output.clear();
    for &label in subs_left.iter().chain(subs_right.iter()) {
        if candidate_label_is_needed(label, left, right, operand_label_sets, needed_label_counts)
            && !output.contains(&label)
        {
            output.push(label);
        }
    }
}

#[derive(Clone, Copy)]
struct CandidateCostContext<'a> {
    operand_label_sets: &'a [HashSet<u32>],
    needed_label_counts: &'a HashMap<u32, usize>,
    size_dict: &'a HashMap<u32, usize>,
}

fn candidate_contraction_cost(
    subs_left: &[u32],
    subs_right: &[u32],
    left: usize,
    right: usize,
    context: CandidateCostContext<'_>,
    candidate_subs: &mut Vec<u32>,
) -> Result<usize> {
    collect_candidate_intermediate_subs(
        subs_left,
        subs_right,
        left,
        right,
        context.operand_label_sets,
        context.needed_label_counts,
        candidate_subs,
    );
    let mut cost = 1usize;
    for &label in candidate_subs.iter() {
        let size = context.size_dict.get(&label).copied().ok_or_else(|| {
            Error::planning(format!(
                "unknown size for label {label} in contraction cost"
            ))
        })?;
        cost = cost.saturating_mul(size);
    }
    Ok(cost.max(1))
}

/// Deterministic greedy contraction order.
///
/// Each step contracts the connected pair (sharing at least one label) whose
/// result is smallest, the size counting every label still needed after the
/// pair: by another operand or by the output. Ties go to the smallest
/// `(left, right)` operand index pair, where an intermediate takes the next
/// index after all existing operands. When no remaining operands share a
/// label, the two smallest indices form an outer product.
///
/// This is the greedy rule of omeco's `tree_greedy` with `alpha = 0` and
/// `temperature = 0` (omeco, MIT, <https://github.com/GiggleLiu/omeco>, which
/// ports `OMEinsumContractionOrders.jl`), with the deterministic vertex and
/// tie ordering adopted upstream after omeco 0.2.6. omeco 0.2.6 iterates a
/// `HashMap` to break ties, so its path differed between processes (#1963);
/// tenferro plans its default (non-annealing) path here instead, until an
/// omeco release carries the fix (<https://github.com/GiggleLiu/omeco/issues/44>). Only the
/// pairs whose operands changed are re-scored after each step, because the
/// cost of any other pair is unaffected by the merge.
fn optimize_self_greedy_pairs(
    subscripts: &Subscripts,
    size_dict: &HashMap<u32, usize>,
) -> Result<Vec<(usize, usize)>> {
    use std::cmp::Reverse;
    use std::collections::{BTreeMap, BTreeSet, BinaryHeap};

    #[cfg(test)]
    tests::SELF_GREEDY_CALLS.with(|count| count.set(count.get() + 1));
    let input_count = subscripts.inputs.len();
    let available: Vec<usize> = (0..input_count).collect();
    let mut operand_subs: Vec<Vec<u32>> = subscripts.inputs.clone();
    let mut operand_label_sets = build_operand_label_sets(&operand_subs);
    let mut needed_label_counts =
        build_needed_label_counts(&subscripts.output, &available, &operand_label_sets);
    let mut live: BTreeSet<usize> = available.into_iter().collect();
    let mut label_owners: BTreeMap<u32, BTreeSet<usize>> = BTreeMap::new();
    for (operand, labels) in operand_label_sets.iter().enumerate() {
        for &label in labels {
            label_owners.entry(label).or_default().insert(operand);
        }
    }

    let mut candidate_subs = Vec::new();
    let mut heap: BinaryHeap<Reverse<(usize, usize, usize)>> = BinaryHeap::new();
    let mut score = |left: usize,
                     right: usize,
                     operand_subs: &[Vec<u32>],
                     operand_label_sets: &[HashSet<u32>],
                     needed_label_counts: &HashMap<u32, usize>|
     -> Result<Reverse<(usize, usize, usize)>> {
        let cost = candidate_contraction_cost(
            &operand_subs[left],
            &operand_subs[right],
            left,
            right,
            CandidateCostContext {
                operand_label_sets,
                needed_label_counts,
                size_dict,
            },
            &mut candidate_subs,
        )?;
        Ok(Reverse((cost, left, right)))
    };

    let mut initial_pairs = BTreeSet::new();
    for owners in label_owners.values() {
        for &left in owners {
            for &right in owners.range(left + 1..) {
                initial_pairs.insert((left, right));
            }
        }
    }
    for (left, right) in initial_pairs {
        heap.push(score(
            left,
            right,
            &operand_subs,
            &operand_label_sets,
            &needed_label_counts,
        )?);
    }

    let mut pairs: Vec<(usize, usize)> = Vec::with_capacity(input_count.saturating_sub(1));
    while live.len() > 1 {
        let mut chosen = None;
        while let Some(Reverse((_, left, right))) = heap.pop() {
            if live.contains(&left) && live.contains(&right) {
                chosen = Some((left, right));
                break;
            }
        }
        let (left, right) = match chosen {
            Some(pair) => pair,
            None => {
                let mut smallest = live.iter().copied();
                match (smallest.next(), smallest.next()) {
                    (Some(left), Some(right)) => (left, right),
                    _ => break,
                }
            }
        };
        pairs.push((left, right));

        let mut new_subs = Vec::new();
        collect_candidate_intermediate_subs(
            &operand_subs[left],
            &operand_subs[right],
            left,
            right,
            &operand_label_sets,
            &needed_label_counts,
            &mut new_subs,
        );
        let new_idx = operand_subs.len();
        let new_label_set: HashSet<u32> = new_subs.iter().copied().collect();
        remove_labels_from_counts(&mut needed_label_counts, &operand_label_sets[left]);
        remove_labels_from_counts(&mut needed_label_counts, &operand_label_sets[right]);
        add_labels_to_counts(&mut needed_label_counts, &new_label_set);
        for operand in [left, right] {
            for label in &operand_label_sets[operand] {
                if let Some(owners) = label_owners.get_mut(label) {
                    owners.remove(&operand);
                }
            }
        }
        let mut neighbors = BTreeSet::new();
        for &label in &new_subs {
            let owners = label_owners.entry(label).or_default();
            neighbors.extend(owners.iter().copied());
            owners.insert(new_idx);
        }
        live.remove(&left);
        live.remove(&right);
        live.insert(new_idx);
        operand_subs.push(new_subs);
        operand_label_sets.push(new_label_set);
        for neighbor in neighbors {
            heap.push(score(
                neighbor,
                new_idx,
                &operand_subs,
                &operand_label_sets,
                &needed_label_counts,
            )?);
        }
    }

    Ok(pairs)
}

#[cfg(test)]
mod tests;