diff --git a/src/undirected_graph.rs b/src/undirected_graph.rs
index 127a48c31a6f171965c7dddb2d7483712f47bdac..5194186a642c618542ba29184f7056bb9c77609b 100644
--- a/src/undirected_graph.rs
+++ b/src/undirected_graph.rs
@@ -285,8 +285,12 @@ impl UndirectedGraph {
         return diff;
     }
 
-    pub fn contractions(&self, edge_orbit_reps: &Vec<usize>) -> HashSet<UndirectedGraph> {
-        let mut contractions: HashSet<UndirectedGraph> = HashSet::new();
+    pub fn contractions(
+        &self,
+        edge_orbit_reps: &Vec<usize>,
+    ) -> HashMap<UndirectedGraph, HashMap<usize, usize>> {
+        // TODO: Check that passing around hashmaps is not too expensive.
+        let mut contractions: HashMap<UndirectedGraph, HashMap<usize, usize>> = HashMap::new();
         for &edge_orbit_rep in edge_orbit_reps {
             let edge: (usize, usize) = self.edges[edge_orbit_rep];
             // TODO: Optimize the computation of new_edges.
@@ -319,9 +323,8 @@ impl UndirectedGraph {
                 .all(|new_edge| unique_edges.insert(*new_edge));
             if has_no_double_edges {
                 let cg: UndirectedGraph = UndirectedGraph::new(self.num_vertices - 1, new_edges);
-                let (_normal_cg_sign, normal_cg, _normal_cg_orbits) = cg.normal_form();
-                // TODO: Also return vertex orbit data.
-                contractions.insert(normal_cg);
+                let (_normal_cg_sign, normal_cg, normal_cg_orbits) = cg.normal_form();
+                contractions.insert(normal_cg, normal_cg_orbits);
             }
         }
         return contractions;