From 9310da974a926ea7f90396bc1d132fcc76285f31 Mon Sep 17 00:00:00 2001
From: Craig Roy <craig.roy@cambridgequantum.com>
Date: Fri, 22 Sep 2023 07:42:04 +0100
Subject: [PATCH] fix: Extension validation for looping graphs

---
 src/extension/infer.rs | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/extension/infer.rs b/src/extension/infer.rs
index f3283a6c3..c12ca05bd 100644
--- a/src/extension/infer.rs
+++ b/src/extension/infer.rs
@@ -658,7 +658,17 @@ impl UnificationContext {
     pub fn instantiate_variables(&mut self) {
         for m in self.variables.clone().into_iter() {
             if !self.solved.contains_key(&m) {
-                self.add_solution(m, ExtensionSet::new());
+                // 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 {
+                            Constraint::Plus(x, other_m) if &m == other_m => Some(x.clone()),
+                            _ => None,
+                        },
+                    ));
+                self.add_solution(m, solution);
             }
         }
         self.variables = HashSet::new();