perf(library/type_context): add caching for type_context::infer

src/library/blast/blast.cpp

@@ -264,11 +264,11 @@ class blastenv {
             return m_benv.m_curr_state.mk_metavar(type);
         }
 
-        virtual void push() override {
+        virtual void push_core() override {
             m_stack.push_back(m_benv.m_curr_state.save_assignment());
         }
 
-        virtual void pop() override {
+        virtual void pop_core() override {
             m_benv.m_curr_state.restore_assignment(m_stack.back());
             m_stack.pop_back();
         }

src/library/tmp_type_context.cpp

@@ -35,6 +35,7 @@ void tmp_type_context::clear() {
     m_eassignment.clear();
     m_trail.clear();
     m_scopes.clear();
+    clear_infer_cache();
 }
 
 void tmp_type_context::set_next_uvar_idx(unsigned next_idx) {
@@ -111,7 +112,7 @@ expr tmp_type_context::mk_mvar(expr const & type) {
     return mk_idx_metavar(idx, type);
 }
 
-void tmp_type_context::push() {
+void tmp_type_context::push_core() {
     m_scopes.push_back(scope());
     scope & s = m_scopes.back();
     s.m_uassignment_sz     = m_uassignment.size();
@@ -119,7 +120,7 @@ void tmp_type_context::push() {
     s.m_trail_sz           = m_trail.size();
 }
 
-void tmp_type_context::pop() {
+void tmp_type_context::pop_core() {
     lean_assert(!m_scopes.empty());
     scope const & s  = m_scopes.back();
     unsigned old_sz  = s.m_trail_sz;

src/library/tmp_type_context.h

@@ -77,8 +77,8 @@ public:
     virtual level mk_uvar();
     virtual expr mk_mvar(expr const &);
 
-    virtual void push();
+    virtual void push_core();
-    virtual void pop();
+    virtual void pop_core();
     virtual unsigned get_num_check_points() const;
     virtual void commit();
 

src/library/type_context.cpp

@@ -124,6 +124,16 @@ type_context::~type_context() {
         delete m_local_gen;
 }
 
+void type_context::push() {
+    m_infer_cache.push();
+    push_core();
+}
+
+void type_context::pop() {
+    pop_core();
+    m_infer_cache.pop();
+}
+
 expr type_context::mk_internal_local(name const & n, expr const & type, binder_info const & bi) {
     return mk_local(m_ngen.next(), n, type, bi);
 }
@@ -1187,7 +1197,9 @@ expr type_context::infer(expr const & e) {
     lean_assert(!is_var(e));
     lean_assert(closed(e));
     check_system("infer_type");
-
+    auto it = m_infer_cache.find(e);
+    if (it != m_infer_cache.end())
+        return it->second;
     expr r;
     switch (e.kind()) {
     case expr_kind::Local:
@@ -1217,13 +1229,18 @@ expr type_context::infer(expr const & e) {
         r = infer_app(e);
         break;
     }
-    // TODO(Leo): cache results if we have performance problems
+    m_infer_cache.insert(mk_pair(e, r));
     return r;
 }
 
 void type_context::clear_cache() {
     m_ci_cache.clear();
     m_ss_cache.clear();
+    clear_infer_cache();
+}
+
+void type_context::clear_infer_cache() {
+    m_infer_cache.clear();
 }
 
 /** \brief If the constant \c e is a class, return its name */

src/library/type_context.h

@@ -5,8 +5,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Author: Leonardo de Moura
 */
 #pragma once
+#include <functional>
 #include <memory>
 #include <vector>
+#include "util/scoped_map.h"
 #include "kernel/environment.h"
 #include "library/io_state.h"
 #include "library/io_state_stream.h"
@@ -184,6 +186,9 @@ class type_context {
         We use it when inferring types and we want to be sure a type is Pi-type. */
     bool                            m_relax_is_opaque;
 
+    typedef scoped_map<expr, expr, expr_hash, std::equal_to<expr>> infer_cache;
+    infer_cache m_infer_cache;
+
     bool is_opaque(declaration const & d) const;
     optional<expr> reduce_projection(expr const & e);
     optional<expr> norm_ext(expr const & e);
@@ -430,9 +435,9 @@ public:
     virtual expr mk_mvar(expr const &) = 0;
 
     /** \brief Save the current assignment and metavariable declarations */
-    virtual void push() = 0;
+    virtual void push_core() = 0;
     /** \brief Retore assignment (inverse for push) */
-    virtual void pop() = 0;
+    virtual void pop_core() = 0;
     /** \brief Return the number of checkpoints created using \c push and not popped yet. */
     virtual unsigned get_num_check_points() const = 0;
     /** \brief Keep the changes since last push */
@@ -453,6 +458,9 @@ public:
     bool has_assigned_uvar(levels const & ls) const;
     bool has_assigned_uvar_mvar(expr const & e) const;
 
+    void push();
+    void pop();
+
     /** \brief Expand macro using extension context */
     optional<expr> expand_macro(expr const & m);
 
@@ -524,8 +532,10 @@ public:
     /** \brief Similar to \c force_assign but sets m_relax_is_opaque */
     bool relaxed_force_assign(expr const & ma, expr const & v);
 
-    /** \brief Clear internal caches used to speedup computation */
+    /** \brief Clear all internal caches used to speedup computation */
     void clear_cache();
+    /** \brief Clear internal type inference cache used to speedup computation */
+    void clear_infer_cache();
 
     /** \brief Update configuration options.
         Return true iff the new options do not change the behavior of the object.
@@ -604,8 +614,8 @@ public:
     virtual expr mk_mvar(expr const &);
     virtual expr infer_local(expr const & e) const { return mlocal_type(e); }
     virtual expr infer_metavar(expr const & e) const { return mlocal_type(e); }
-    virtual void push() { m_trail.push_back(m_assignment); }
+    virtual void push_core() { m_trail.push_back(m_assignment); }
-    virtual void pop() { lean_assert(!m_trail.empty()); m_assignment = m_trail.back(); m_trail.pop_back(); }
+    virtual void pop_core() { lean_assert(!m_trail.empty()); m_assignment = m_trail.back(); m_trail.pop_back(); }
     virtual unsigned get_num_check_points() const { return m_trail.size(); }
     virtual void commit() { lean_assert(!m_trail.empty()); m_trail.pop_back(); }
     virtual optional<expr> mk_subsingleton_instance(expr const & type);