From fe3396e1ae08b71cd9056203a90dd712f9d9d7f0 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 12 Dec 2016 09:05:07 -0800
Subject: [PATCH] perf(library/vm/vm): cache the result of 0-ary vm_decls

---
 src/library/vm/vm.cpp | 15 +++++++++++++++
 src/library/vm/vm.h   |  4 +++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/library/vm/vm.cpp b/src/library/vm/vm.cpp
index 3fe622cd38..014186ccc8 100644
--- a/src/library/vm/vm.cpp
+++ b/src/library/vm/vm.cpp
@@ -1842,6 +1842,13 @@ unsigned vm_state::pop_frame_core() {
     lean_assert(sz - 1 < m_stack.size());
     swap(m_stack[sz - fr.m_num - 1], m_stack[sz - 1]);
     m_stack.resize(sz - fr.m_num);
+    unsigned curr_fidx = fr.m_curr_fn_idx;
+    if (curr_fidx != g_null_fn_idx && get_decl(curr_fidx).get_arity() == 0) {
+        /* cache result */
+        if (curr_fidx >= m_cache_vector.size())
+            m_cache_vector.resize(curr_fidx+1);
+        m_cache_vector[curr_fidx] = m_stack.back();
+    }
     if (m_debugging) shrink_stack_info();
     m_code   = fr.m_code;
     m_fn_idx = fr.m_fn_idx;
@@ -2305,6 +2312,14 @@ void vm_state::run() {
                where n is fn.arity
             */
             vm_decl decl = get_decl(instr.get_fn_idx());
+            /* If d is 0-ary, then check if value is cached */
+            if (decl.get_arity() == 0 && decl.get_idx() < m_cache_vector.size()) {
+                if (auto r = m_cache_vector[decl.get_idx()]) {
+                    m_stack.push_back(*r);
+                    m_pc++;
+                    goto main_loop;
+                }
+            }
             invoke_global(decl);
             goto main_loop;
         }
diff --git a/src/library/vm/vm.h b/src/library/vm/vm.h
index 99a54f3078..d6b879ce88 100644
--- a/src/library/vm/vm.h
+++ b/src/library/vm/vm.h
@@ -519,6 +519,7 @@ public:
 /** \brief Virtual machine for executing VM bytecode. */
 class vm_state {
     typedef std::vector<vm_decl> decl_vector;
+    typedef std::vector<optional<vm_obj>> cache_vector;
     typedef unsigned_map<vm_decl> decl_map;
     typedef std::vector<vm_cases_function> builtin_cases_vector;
     typedef unsigned_map<vm_cases_function> builtin_cases_map;
@@ -526,6 +527,7 @@ class vm_state {
     options                     m_options;
     decl_map                    m_decl_map;
     decl_vector                 m_decl_vector;
+    cache_vector                m_cache_vector; /* for 0-ary declarations */
     builtin_cases_map           m_builtin_cases_map;
     builtin_cases_vector        m_builtin_cases_vector;
     unsigned_map<name>          m_builtin_cases_names;
@@ -543,8 +545,8 @@ class vm_state {
         unsigned                m_num;
         unsigned                m_pc;
         unsigned                m_bp;
-        /* The following two fields are only used for profiling the code */
         unsigned                m_curr_fn_idx;
+        /* The following two fields are only used for profiling the code */
         unsigned                m_frame_idx;
         frame(vm_instr const * code, unsigned fn_idx, unsigned num, unsigned pc, unsigned bp,
               unsigned curr_fn_idx, unsigned frame_idx):