From 973a5c4812dea2116495660176076746aa072769 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 9 Nov 2015 13:36:49 -0800
Subject: [PATCH] feat(library/blast): improve blast type_context push/pop
 operations

---
 src/library/blast/blast.cpp |  9 ++++-----
 src/library/blast/state.cpp | 10 ++++++++++
 src/library/blast/state.h   | 22 +++++++++-------------
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/src/library/blast/blast.cpp b/src/library/blast/blast.cpp
index d2bffd664..72d35ea5e 100644
--- a/src/library/blast/blast.cpp
+++ b/src/library/blast/blast.cpp
@@ -87,8 +87,8 @@ class blastenv {
     unsigned                   m_inc_depth;
 
     class tctx : public type_context {
-        blastenv &         m_benv;
-        std::vector<state> m_stack;
+        blastenv &                              m_benv;
+        std::vector<state::assignment_snapshot> m_stack;
     public:
         tctx(blastenv & benv):
             type_context(benv.m_env, benv.m_ios, benv.m_tmp_local_generator),
@@ -199,12 +199,11 @@ class blastenv {
         }
 
         virtual void push() {
-            // TODO(Leo): we only need to save the assignment and metavar_decls.
-            m_stack.push_back(m_benv.m_curr_state);
+            m_stack.push_back(m_benv.m_curr_state.save_assignment());
         }
 
         virtual void pop() {
-            m_benv.m_curr_state = m_stack.back();
+            m_benv.m_curr_state.restore_assignment(m_stack.back());
             m_stack.pop_back();
         }
 
diff --git a/src/library/blast/state.cpp b/src/library/blast/state.cpp
index 783fe98d8..7aa3935d1 100644
--- a/src/library/blast/state.cpp
+++ b/src/library/blast/state.cpp
@@ -514,6 +514,16 @@ expr state::expand_hrefs(expr const & e, list<expr> const & hrefs) const {
     return expand_hrefs_fn(*this, hrefs)(e);
 }
 
+auto state::save_assignment() -> assignment_snapshot {
+    return assignment_snapshot(m_uassignment, m_metavar_decls, m_eassignment);
+}
+
+void state::restore_assignment(assignment_snapshot const & s) {
+    m_uassignment   = s.m_uassignment;
+    m_metavar_decls = s.m_metavar_decls;
+    m_eassignment   = s.m_eassignment;
+}
+
 void initialize_state() {
     g_prefix     = new name(name::mk_internal_unique_name());
 }
diff --git a/src/library/blast/state.h b/src/library/blast/state.h
index 756b07ee0..129b011b6 100644
--- a/src/library/blast/state.h
+++ b/src/library/blast/state.h
@@ -246,23 +246,19 @@ public:
 
     void display(environment const & env, io_state const & ios) const;
 
-    /** Auxiliary object for creating snapshots of the metavariable assignments.
-        \remark The snapshots are created (and restored) in constant time */
     class assignment_snapshot {
-        state &     m_state;
-        uassignment m_old_uassignment;
-        eassignment m_old_eassignment;
+        friend class state;
+        uassignment        m_uassignment;
+        metavar_decls      m_metavar_decls;
+        eassignment        m_eassignment;
+        assignment_snapshot(uassignment const & u, metavar_decls const & decls, eassignment const & e):
+            m_uassignment(u), m_metavar_decls(decls), m_eassignment(e) {}
     public:
-        assignment_snapshot(state & s):
-            m_state(s),
-            m_old_uassignment(s.m_uassignment),
-            m_old_eassignment(s.m_eassignment) {}
-        void restore() {
-            m_state.m_uassignment = m_old_uassignment;
-            m_state.m_eassignment = m_old_eassignment;
-        }
     };
 
+    assignment_snapshot save_assignment();
+    void restore_assignment(assignment_snapshot const & s);
+
     void push_proof_step(proof_step const & ps) {
         m_depth++;
         m_proof_steps = cons(ps, m_proof_steps);