From 65e64d100669c3f0fde64d5b0cd154a0317a2c50 Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Mon, 1 Apr 2013 09:54:32 -0700
Subject: [PATCH] loop counting

Signed-off-by: Nikolaj Bjorner <nbjorner@microsoft.com>
---
 src/muz_qe/dl_mk_loop_counter.cpp | 24 +++++++++++++++++-------
 src/muz_qe/dl_mk_loop_counter.h   |  6 ++++++
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/muz_qe/dl_mk_loop_counter.cpp b/src/muz_qe/dl_mk_loop_counter.cpp
index 1e7d0c3fa..172198746 100644
--- a/src/muz_qe/dl_mk_loop_counter.cpp
+++ b/src/muz_qe/dl_mk_loop_counter.cpp
@@ -25,24 +25,34 @@ namespace datalog {
     mk_loop_counter::mk_loop_counter(context & ctx, unsigned priority):
         plugin(priority),
         m(ctx.get_manager()),
-        a(m) {        
+        a(m),
+        m_refs(m) {        
     }
 
     mk_loop_counter::~mk_loop_counter() { }
 
     app_ref mk_loop_counter::add_arg(app* fn, unsigned idx) {
-        ptr_vector<sort> domain;
         expr_ref_vector args(m);
-        domain.append(fn->get_num_args(), fn->get_decl()->get_domain());
-        domain.push_back(a.mk_int());
+        func_decl* new_fn, *old_fn = fn->get_decl();
         args.append(fn->get_num_args(), fn->get_args());
         args.push_back(m.mk_var(idx, a.mk_int()));
-        func_decl_ref new_fn(m);
-        new_fn = m.mk_func_decl(fn->get_decl()->get_name(), domain.size(), domain.c_ptr(), m.mk_bool_sort());
+        
+        if (!m_old2new.find(old_fn, new_fn)) {
+            ptr_vector<sort> domain;
+            domain.append(fn->get_num_args(), old_fn->get_domain());
+            domain.push_back(a.mk_int());
+            new_fn = m.mk_func_decl(old_fn->get_name(), domain.size(), domain.c_ptr(), old_fn->get_range());
+            m_old2new.insert(old_fn, new_fn);
+            m_new2old.insert(new_fn, old_fn);
+            m_refs.push_back(new_fn);
+        }
         return app_ref(m.mk_app(new_fn, args.size(), args.c_ptr()), m);
     }
         
     rule_set * mk_loop_counter::operator()(rule_set const & source) {
+        m_refs.reset();
+        m_old2new.reset();
+        m_new2old.reset();
         context& ctx = source.get_context();
         rule_manager& rm = source.get_rule_manager();
         rule_set * result = alloc(rule_set, ctx);
@@ -95,7 +105,7 @@ namespace datalog {
         }
 
         // model converter: remove references to extra argument.
-        // proof converter: etc.
+        // proof converter: remove references to extra argument as well.
 
         return result;
     }
diff --git a/src/muz_qe/dl_mk_loop_counter.h b/src/muz_qe/dl_mk_loop_counter.h
index c16159b1c..6601f837f 100644
--- a/src/muz_qe/dl_mk_loop_counter.h
+++ b/src/muz_qe/dl_mk_loop_counter.h
@@ -27,12 +27,18 @@ namespace datalog {
     class mk_loop_counter : public rule_transformer::plugin {
         ast_manager& m;
         arith_util   a;
+        func_decl_ref_vector m_refs;
+        obj_map<func_decl, func_decl*> m_new2old;
+        obj_map<func_decl, func_decl*> m_old2new;
+
         app_ref add_arg(app* fn, unsigned idx);        
     public:
         mk_loop_counter(context & ctx, unsigned priority = 33000);
         ~mk_loop_counter();
         
         rule_set * operator()(rule_set const & source);
+
+        func_decl* get_old(func_decl* f) const { return m_new2old.find(f); }
     };
 
 };