impl fold_const for RegionFudger

Signed-off-by: Gabriel Smith <ga29smith@gmail.com>

2 files changed, 58 insertions, 1 deletions

diff --git a/src/librustc/infer/const_variable.rs b/src/librustc/infer/const_variable.rs
index 5691bf08eae..ac758add872 100644
--- a/src/librustc/infer/const_variable.rs
+++ b/src/librustc/infer/const_variable.rs
@@ -5,6 +5,7 @@ use crate::ty::{self, InferConst};
 
 use std::cmp;
 use std::marker::PhantomData;
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::snapshot_vec as sv;
 use rustc_data_structures::unify as ut;
 
@@ -23,6 +24,8 @@ pub enum ConstVariableOrigin {
     SubstitutionPlaceholder(Span),
 }
 
+pub type ConstVariableMap<'tcx> = FxHashMap<ty::ConstVid<'tcx>, ConstVariableOrigin>;
+
 struct ConstVariableData {
     origin: ConstVariableOrigin,
 }
@@ -184,6 +187,32 @@ impl<'tcx> ConstVariableTable<'tcx> {
         self.values.commit(snapshot);
         self.relations.commit(relation_snapshot);
     }
+
+    /// Returns a map `{V1 -> V2}`, where the keys `{V1}` are
+    /// const-variables created during the snapshot, and the values
+    /// `{V2}` are the root variables that they were unified with,
+    /// along with their origin.
+    pub fn consts_created_since_snapshot(
+        &mut self,
+        s: &Snapshot<'tcx>
+    ) -> ConstVariableMap<'tcx> {
+        let actions_since_snapshot = self.values.actions_since_snapshot(&s.snapshot);
+
+        actions_since_snapshot
+            .iter()
+            .filter_map(|action| match action {
+                &sv::UndoLog::NewElem(index) => Some(ty::ConstVid {
+                    index: index as u32,
+                    phantom: PhantomData,
+                }),
+                _ => None,
+            })
+            .map(|vid| {
+                let origin = self.values.get(vid.index as usize).origin.clone();
+                (vid, origin)
+            })
+            .collect()
+    }
 }
 
 impl<'tcx> ut::UnifyKey for ty::ConstVid<'tcx> {
diff --git a/src/librustc/infer/fudge.rs b/src/librustc/infer/fudge.rs
index 4bab0bedbd6..5bb007dbb00 100644
--- a/src/librustc/infer/fudge.rs
+++ b/src/librustc/infer/fudge.rs
@@ -1,5 +1,6 @@
 use crate::ty::{self, Ty, TyCtxt, TyVid, IntVid, FloatVid, RegionVid};
 use crate::ty::fold::{TypeFoldable, TypeFolder};
+use crate::mir::interpret::ConstValue;
 
 use super::InferCtxt;
 use super::RegionVariableOrigin;
@@ -176,6 +177,33 @@ impl<'a, 'gcx, 'tcx> TypeFolder<'gcx, 'tcx> for InferenceFudger<'a, 'gcx, 'tcx>
     }
 
     fn fold_const(&mut self, ct: &'tcx ty::LazyConst<'tcx>) -> &'tcx ty::LazyConst<'tcx> {
-        ct // FIXME(const_generics)
+        if let ty::LazyConst::Evaluated(ty::Const {
+            val: ConstValue::Infer(ty::InferConst::Var(vid)),
+            ty,
+        }) = *ct {
+            match self.const_variables.get(&vid) {
+                None => {
+                    // This variable was created before the
+                    // "fudging".  Since we refresh all
+                    // variables to their binding anyhow, we know
+                    // that it is unbound, so we can just return
+                    // it.
+                    debug_assert!(
+                        self.infcx.const_unification_table.borrow_mut()
+                        .probe(vid)
+                        .is_unknown()
+                    );
+                    ct
+                }
+                Some(&origin) => {
+                    // This variable was created during the
+                    // fudging. Recreate it with a fresh variable
+                    // here.
+                    self.infcx.next_const_var(ty, origin)
+                }
+            }
+        } else {
+            ct.super_fold_with(self)
+        }
     }
 }