From 6a82f33c722fcd1a6dddef7bbe4f837d1e5238dc Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sun, 14 Aug 2016 15:45:02 -0700
Subject: [PATCH] feat(library): add replace_visitor_with_tc

---
 src/library/CMakeLists.txt              |  2 +-
 src/library/replace_visitor_with_tc.cpp | 51 +++++++++++++++++++++++++
 src/library/replace_visitor_with_tc.h   | 29 ++++++++++++++
 3 files changed, 81 insertions(+), 1 deletion(-)
 create mode 100644 src/library/replace_visitor_with_tc.cpp
 create mode 100644 src/library/replace_visitor_with_tc.h

diff --git a/src/library/CMakeLists.txt b/src/library/CMakeLists.txt
index e8c60387df..1301eda391 100644
--- a/src/library/CMakeLists.txt
+++ b/src/library/CMakeLists.txt
@@ -14,7 +14,7 @@ add_library(library OBJECT deep_copy.cpp expr_lt.cpp io_state.cpp
   attribute_manager.cpp error_handling.cpp unification_hint.cpp
   local_context.cpp metavar_context.cpp type_context.cpp export_decl.cpp delayed_abstraction.cpp
   fun_info.cpp congr_lemma.cpp defeq_canonizer.cpp scope_pos_info_provider.cpp
-  inductive.cpp mpq_macro.cpp arith_instance_manager.cpp
+  inductive.cpp mpq_macro.cpp arith_instance_manager.cpp replace_visitor_with_tc.cpp
 
   # Legacy -- The following files will be eventually deleted
   normalize.cpp justification.cpp constraint.cpp metavar.cpp choice.cpp locals.cpp
diff --git a/src/library/replace_visitor_with_tc.cpp b/src/library/replace_visitor_with_tc.cpp
new file mode 100644
index 0000000000..ff20f81c9f
--- /dev/null
+++ b/src/library/replace_visitor_with_tc.cpp
@@ -0,0 +1,51 @@
+/*
+Copyright (c) 2016 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#include "kernel/instantiate.h"
+#include "library/replace_visitor_with_tc.h"
+namespace lean {
+expr replace_visitor_with_tc::visit_lambda_pi_let(bool is_lam, expr const & e) {
+    type_context::tmp_locals locals(m_ctx);
+    expr t = e;
+    while (true) {
+        if ((is_lam && is_lambda(t)) || (!is_lam && is_pi(t))) {
+            expr d = instantiate_rev(binding_domain(t), locals.size(), locals.data());
+            d = visit(d);
+            locals.push_local(binding_name(t), d, binding_info(t));
+            t = binding_body(t);
+        } else if (is_let(t)) {
+            expr d = instantiate_rev(let_type(t), locals.size(), locals.data());
+            expr v = instantiate_rev(let_value(t), locals.size(), locals.data());
+            d = visit(d);
+            v = visit(v);
+            locals.push_let(let_name(t), d, v);
+            t = let_body(t);
+        } else {
+            break;
+        }
+    }
+    t = instantiate_rev(t, locals.size(), locals.data());
+    t = visit(t);
+    return is_lam ? locals.mk_lambda(t) : locals.mk_pi(t);
+}
+
+expr replace_visitor_with_tc::visit_app(expr const & e) {
+    buffer<expr> args;
+    expr const & fn = get_app_args(e, args);
+    expr new_fn   = visit(fn);
+    bool modified = !is_eqp(fn, new_fn);
+    for (expr & arg : args) {
+        expr new_arg = visit(arg);
+        if (!is_eqp(new_arg, arg))
+            modified = true;
+        arg = new_arg;
+    }
+    if (!modified)
+        return e;
+    else
+        return mk_app(new_fn, args);
+}
+}
diff --git a/src/library/replace_visitor_with_tc.h b/src/library/replace_visitor_with_tc.h
new file mode 100644
index 0000000000..0d0449d898
--- /dev/null
+++ b/src/library/replace_visitor_with_tc.h
@@ -0,0 +1,29 @@
+/*
+Copyright (c) 2016 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Leonardo de Moura
+*/
+#pragma once
+#include "library/replace_visitor.h"
+#include "library/type_context.h"
+namespace lean {
+/** \brief Version of replace_visitor with a nested type_context object. */
+class replace_visitor_with_tc : public replace_visitor {
+protected:
+    type_context & m_ctx;
+    expr visit_lambda_pi_let(bool is_lam, expr const & e);
+    virtual expr visit_lambda(expr const & e) override {
+        return visit_lambda_pi_let(true, e);
+    }
+    virtual expr visit_let(expr const & e) override {
+        return visit_lambda_pi_let(true, e);
+    }
+    virtual expr visit_pi(expr const & e) override {
+        return visit_lambda_pi_let(false, e);
+    }
+    virtual expr visit_app(expr const & e) override;
+public:
+    replace_visitor_with_tc(type_context & ctx):m_ctx(ctx) {}
+};
+}