fix(Inference): Work harder in variable instantiation

src/extension/infer.rs

@@ -22,6 +22,7 @@ use crate::{
 use super::validate::ExtensionError;
 
 use petgraph::graph as pg;
+use petgraph::{Directed, EdgeType, Undirected};
 
 use std::collections::{HashMap, HashSet};
 
@@ -107,53 +108,65 @@ pub enum InferExtensionError {
     EdgeMismatch(#[from] ExtensionError),
 }
 
-/// A graph of metavariables which we've found equality constraints for. Edges
-/// between nodes represent equality constraints.
-struct EqGraph {
-    equalities: pg::Graph<Meta, (), petgraph::Undirected>,
+/// A graph of metavariables connected by constraints.
+/// The edges represent `Equal` constraints in the undirected graph and `Plus`
+/// constraints in the directed case.
+struct GraphContainer<Dir: EdgeType> {
+    graph: pg::Graph<Meta, (), Dir>,
     node_map: HashMap<Meta, pg::NodeIndex>,
 }
 
-impl EqGraph {
-    /// Create a new `EqGraph`
-    fn new() -> Self {
-        EqGraph {
-            equalities: pg::Graph::new_undirected(),
-            node_map: HashMap::new(),
-        }
-    }
-
+impl<T: EdgeType> GraphContainer<T> {
     /// Add a metavariable to the graph as a node and return the `NodeIndex`.
     /// If it's already there, just return the existing `NodeIndex`
     fn add_or_retrieve(&mut self, m: Meta) -> pg::NodeIndex {
         self.node_map.get(&m).cloned().unwrap_or_else(|| {
-            let ix = self.equalities.add_node(m);
+            let ix = self.graph.add_node(m);
             self.node_map.insert(m, ix);
             ix
         })
     }
 
-    /// Create an edge between two nodes on the graph, declaring that they stand
-    /// for metavariables which should be equal.
-    fn register_eq(&mut self, src: Meta, tgt: Meta) {
+    /// Create an edge between two nodes on the graph
+    fn add_edge(&mut self, src: Meta, tgt: Meta) {
         let src_ix = self.add_or_retrieve(src);
         let tgt_ix = self.add_or_retrieve(tgt);
-        self.equalities.add_edge(src_ix, tgt_ix, ());
+        self.graph.add_edge(src_ix, tgt_ix, ());
     }
 
     /// Return the connected components of the graph in terms of metavariables
     fn ccs(&self) -> Vec<Vec<Meta>> {
-        petgraph::algo::tarjan_scc(&self.equalities)
+        petgraph::algo::tarjan_scc(&self.graph)
             .into_iter()
             .map(|cc| {
                 cc.into_iter()
-                    .map(|n| *self.equalities.node_weight(n).unwrap())
+                    .map(|n| *self.graph.node_weight(n).unwrap())
                     .collect()
             })
             .collect()
     }
 }
 
+impl GraphContainer<Undirected> {
+    fn new() -> Self {
+        GraphContainer {
+            graph: pg::Graph::new_undirected(),
+            node_map: HashMap::new(),
+        }
+    }
+}
+
+impl GraphContainer<Directed> {
+    fn new() -> Self {
+        GraphContainer {
+            graph: pg::Graph::new(),
+            node_map: HashMap::new(),
+        }
+    }
+}
+
+type EqGraph = GraphContainer<Undirected>;
+
 /// Our current knowledge about the extensions of the graph
 struct UnificationContext {
     /// A list of constraints for each metavariable
@@ -504,7 +517,7 @@ impl UnificationContext {
             match c {
                 // Just register the equality in the EqGraph, we'll process it later
                 Constraint::Equal(other_meta) => {
-                    self.eq_graph.register_eq(meta, *other_meta);
+                    self.eq_graph.add_edge(meta, *other_meta);
                 }
                 // N.B. If `meta` is already solved, we can't use that
                 // information to solve `other_meta`. This is because the Plus
@@ -663,31 +676,90 @@ impl UnificationContext {
         self.results()
     }
 
-    /// Instantiate all variables in the graph with the empty extension set.
+    /// Gather all the transitive dependencies (induced by constraints) of the
+    /// variables in the context.
+    fn search_variable_deps(&self) -> HashSet<Meta> {
+        let mut seen = HashSet::new();
+        let mut new_variables: HashSet<Meta> = self.variables.clone();
+        while !new_variables.is_empty() {
+            new_variables = new_variables
+                .into_iter()
+                .filter(|m| seen.insert(*m))
+                .flat_map(|m| self.get_constraints(&m).unwrap())
+                .map(|c| match c {
+                    Constraint::Plus(_, other) => self.resolve(*other),
+                    Constraint::Equal(other) => self.resolve(*other),
+                })
+                .collect();
+        }
+        seen
+    }
+
     /// Instantiate all variables in the graph with the empty extension set, or
     /// the smallest solution possible given their constraints.
     /// This is done to solve metas which depend on variables, which allows
     /// us to come up with a fully concrete solution to pass into validation.
+    ///
+    /// Nodes which loop into themselves must be considered as a "minimum" set
+    /// of requirements. If we have
+    ///   1 = 2 + X, ...
+    ///   2 = 1 + x, ...
+    /// then 1 and 2 both definitely contain X, even if we don't know what else.
+    /// So instead of instantiating to the empty set, we'll instantiate to `{X}`
     pub fn instantiate_variables(&mut self) {
-        for m in self.variables.clone().into_iter() {
+        // A directed graph to keep track of `Plus` constraint relationships
+        let mut relations = GraphContainer::<Directed>::new();
+        let mut solutions: HashMap<Meta, ExtensionSet> = HashMap::new();
+
+        let variable_scope = self.search_variable_deps();
+        for m in variable_scope.into_iter() {
+            // If `m` has been merged, [`self.variables`] entry
+            // will have already been updated to the merged
+            // value by [`self.merge_equal_metas`] so we don't
+            // need to worry about resolving it.
             if !self.solved.contains_key(&m) {
                 // Handle the case where the constraints for `m` contain a self
                 // reference, i.e. "m = Plus(E, m)", in which case the variable
                 // should be instantiated to E rather than the empty set.
-                let solution =
-                    ExtensionSet::from_iter(self.get_constraints(&m).unwrap().iter().filter_map(
-                        |c| match c {
-                            // If `m` has been merged, [`self.variables`] entry
-                            // will have already been updated to the merged
-                            // value by [`self.merge_equal_metas`] so we don't
-                            // need to worry about resolving it.
-                            Constraint::Plus(x, other_m) if m == self.resolve(*other_m) => {
-                                Some(x.clone())
-                            }
+                let plus_constraints =
+                    self.get_constraints(&m)
+                        .unwrap()
+                        .iter()
+                        .cloned()
+                        .flat_map(|c| match c {
+                            Constraint::Plus(r, other_m) => Some((r, self.resolve(other_m))),
                             _ => None,
-                        },
-                    ));
-                self.add_solution(m, solution);
+                        });
+
+                let (rs, other_ms): (Vec<_>, Vec<_>) = plus_constraints.unzip();
+                let solution = ExtensionSet::from_iter(rs.into_iter());
+                let unresolved_metas = other_ms
+                    .into_iter()
+                    .filter(|other_m| m != *other_m)
+                    .collect::<Vec<_>>();
+
+                // If `m` doesn't depend on any other metas then we have all the
+                // information we need to come up with a solution for it.
+                if unresolved_metas.is_empty() {
+                    self.add_solution(m, solution.clone());
+                } else {
+                    unresolved_metas
+                        .iter()
+                        .for_each(|other_m| relations.add_edge(m, *other_m));
+                }
+                solutions.insert(m, solution);
+            }
+        }
+
+        // Strongly connected components are looping constraint dependencies.
+        // This means that each metavariable in the CC has the same solution.
+        for cc in relations.ccs() {
+            let mut combined_solution = ExtensionSet::new();
+            for sol in cc.iter().filter_map(|m| solutions.get(m)) {
+                combined_solution = combined_solution.union(sol);
+            }
+            for m in cc.iter() {
+                self.add_solution(*m, combined_solution.clone());
             }
         }
         self.variables = HashSet::new();
@@ -1509,17 +1581,14 @@ mod test {
     #[test]
     fn test_cfg_loops() -> Result<(), Box<dyn Error>> {
         let just_a = ExtensionSet::singleton(&A);
-        let variants = vec![
-            (
-                ExtensionSet::new(),
-                ExtensionSet::new(),
-                ExtensionSet::new(),
-            ),
-            (just_a.clone(), ExtensionSet::new(), ExtensionSet::new()),
-            (ExtensionSet::new(), just_a.clone(), ExtensionSet::new()),
-            (ExtensionSet::new(), ExtensionSet::new(), just_a.clone()),
-        ];
-
+        let mut variants = Vec::new();
+        for entry in [ExtensionSet::new(), just_a.clone()] {
+            for bb1 in [ExtensionSet::new(), just_a.clone()] {
+                for bb2 in [ExtensionSet::new(), just_a.clone()] {
+                    variants.push((entry.clone(), bb1.clone(), bb2.clone()));
+                }
+            }
+        }
         for (bb0, bb1, bb2) in variants.into_iter() {
             let mut hugr = make_looping_cfg(bb0, bb1, bb2)?;
             hugr.infer_and_validate(&PRELUDE_REGISTRY)?;