feat(kernel): add tag to kernel expressions

Frontends can used tags to associate expressions with line number information.
The update_* procedures automatically propagate tags.

Signed-off-by: Leonardo de Moura <leonardo@microsoft.com>

src/kernel/expr.cpp

@@ -31,12 +31,13 @@ unsigned hash_levels(levels const & ls) {
 }
 
 expr_cell::expr_cell(expr_kind k, unsigned h, bool has_mv, bool has_local, bool has_param_univ):
-    m_kind(static_cast<unsigned>(k)),
     m_flags(0),
+    m_kind(static_cast<unsigned>(k)),
     m_has_mv(has_mv),
     m_has_local(has_local),
     m_has_param_univ(has_param_univ),
     m_hash(h),
+    m_tag(nulltag),
     m_rc(0) {
     // m_hash_alloc does not need to be a unique identifier.
     // We want diverse hash codes because given expr_cell * c1 and expr_cell * c2,
@@ -77,6 +78,10 @@ void expr_cell::set_is_arrow(bool flag) {
     lean_assert(is_arrow() && *is_arrow() == flag);
 }
 
+void expr_cell::set_tag(tag t) {
+    m_tag = t;
+}
+
 bool is_meta(expr const & e) {
     expr const * it = &e;
     while (is_app(*it)) {
@@ -376,9 +381,16 @@ unsigned get_free_var_range(expr const & e) {
 
 bool operator==(expr const & a, expr const & b) { return expr_eq_fn()(a, b); }
 
+static expr copy_tag(expr const & e, expr && new_e) {
+    tag t = e.get_tag();
+    if (t != nulltag)
+        new_e.set_tag(t);
+    return new_e;
+}
+
 expr update_app(expr const & e, expr const & new_fn, expr const & new_arg) {
     if (!is_eqp(app_fn(e), new_fn) || !is_eqp(app_arg(e), new_arg))
-        return mk_app(new_fn, new_arg);
+        return copy_tag(e, mk_app(new_fn, new_arg));
     else
         return e;
 }
@@ -387,45 +399,45 @@ expr update_rev_app(expr const & e, unsigned num, expr const * new_args) {
     expr const * it = &e;
     for (unsigned i = 0; i < num - 1; i++) {
         if (!is_app(*it) || !is_eqp(app_arg(*it), new_args[i]))
-            return mk_rev_app(num, new_args);
+            return copy_tag(e, mk_rev_app(num, new_args));
         it = &app_fn(*it);
     }
     if (!is_eqp(*it, new_args[num - 1]))
-        return mk_rev_app(num, new_args);
+        return copy_tag(e, mk_rev_app(num, new_args));
     return e;
 }
 
 expr update_binder(expr const & e, expr const & new_domain, expr const & new_body) {
     if (!is_eqp(binder_domain(e), new_domain) || !is_eqp(binder_body(e), new_body))
-        return mk_binder(e.kind(), binder_name(e), new_domain, new_body, binder_info(e));
+        return copy_tag(e, mk_binder(e.kind(), binder_name(e), new_domain, new_body, binder_info(e)));
     else
         return e;
 }
 
 expr update_let(expr const & e, expr const & new_type, expr const & new_val, expr const & new_body) {
     if (!is_eqp(let_type(e), new_type) || !is_eqp(let_value(e), new_val) || !is_eqp(let_body(e), new_body))
-        return mk_let(let_name(e), new_type, new_val, new_body);
+        return copy_tag(e, mk_let(let_name(e), new_type, new_val, new_body));
     else
         return e;
 }
 
 expr update_mlocal(expr const & e, expr const & new_type) {
     if (!is_eqp(mlocal_type(e), new_type))
-        return mk_mlocal(is_metavar(e), mlocal_name(e), new_type);
+        return copy_tag(e, mk_mlocal(is_metavar(e), mlocal_name(e), new_type));
     else
         return e;
 }
 
 expr update_sort(expr const & e, level const & new_level) {
     if (!is_eqp(sort_level(e), new_level))
-        return mk_sort(new_level);
+        return copy_tag(e, mk_sort(new_level));
     else
         return e;
 }
 
 expr update_constant(expr const & e, levels const & new_levels) {
     if (!is_eqp(const_level_params(e), new_levels))
-        return mk_constant(const_name(e), new_levels);
+        return copy_tag(e, mk_constant(const_name(e), new_levels));
     else
         return e;
 }
@@ -440,7 +452,7 @@ expr update_macro(expr const & e, unsigned num, expr const * args) {
         if (i == num)
             return e;
     }
-    return mk_macro(to_macro(e)->m_definition, num, args);
+    return copy_tag(e, mk_macro(to_macro(e)->m_definition, num, args));
 }
 
 bool is_atomic(expr const & e) {

src/kernel/expr.h

@@ -25,6 +25,10 @@ Author: Leonardo de Moura
 #include "kernel/extension_context.h"
 
 namespace lean {
+// Tags are used by frontends to mark expressions. They are automatically propagated by
+// procedures such as update_app, update_binder, etc.
+typedef unsigned tag;
+constexpr tag nulltag = std::numeric_limits<unsigned>::max();
 class expr;
 /* =======================================
    Expressions
@@ -43,16 +47,17 @@ class expr;
 enum class expr_kind { Var, Sort, Constant, Meta, Local, App, Lambda, Pi, Let, Macro };
 class expr_cell {
 protected:
-    unsigned short     m_kind;
     // The bits of the following field mean:
     //    0-1  - term is an arrow (0 - not initialized, 1 - is arrow, 2 - is not arrow)
     // Remark: we use atomic_uchar because these flags are computed lazily (i.e., after the expression is created)
     atomic_uchar       m_flags;
+    unsigned           m_kind:8;
     unsigned           m_has_mv:1;         // term contains metavariables
     unsigned           m_has_local:1;      // term contains local constants
     unsigned           m_has_param_univ:1; // term constains parametric universe levels
-    unsigned           m_hash;        // hash based on the structure of the expression (this is a good hash for structural equality)
-    unsigned           m_hash_alloc;  // hash based on 'time' of allocation (this is a good hash for pointer-based equality)
+    unsigned           m_hash;             // hash based on the structure of the expression (this is a good hash for structural equality)
+    unsigned           m_hash_alloc;       // hash based on 'time' of allocation (this is a good hash for pointer-based equality)
+    atomic_uint        m_tag;
     MK_LEAN_RC(); // Declare m_rc counter
     void dealloc();
 
@@ -69,6 +74,8 @@ public:
     bool has_metavar() const { return m_has_mv; }
     bool has_local() const { return m_has_local; }
     bool has_param_univ() const { return m_has_param_univ; }
+    void set_tag(tag t);
+    tag get_tag() const { return m_tag; }
 };
 
 class macro_definition;
@@ -110,6 +117,10 @@ public:
     bool has_local() const { return m_ptr->has_local(); }
     bool has_param_univ() const { return m_ptr->has_param_univ(); }
 
+    void set_tag(tag t) { m_ptr->set_tag(t); }
+    tag get_tag() const { return m_ptr->get_tag(); }
+
+
     expr_cell * raw() const { return m_ptr; }
 
     friend expr mk_var(unsigned idx);

src/tests/kernel/expr.cpp

@@ -388,8 +388,10 @@ int main() {
     tst17();
     tst18();
     std::cout << "sizeof(expr):            " << sizeof(expr) << "\n";
-    std::cout << "sizeof(expr_app):        " << sizeof(expr_app) << "\n";
     std::cout << "sizeof(expr_cell):       " << sizeof(expr_cell) << "\n";
+    std::cout << "sizeof(expr_app):        " << sizeof(expr_app) << "\n";
+    std::cout << "sizeof(expr_var):        " << sizeof(expr_var) << "\n";
+    std::cout << "sizeof(expr_const):      " << sizeof(expr_const) << "\n";
     std::cout << "sizeof(optional<expr>):  " << sizeof(optional<expr>) << "\n";
     std::cout << "sizeof(optional<sexpr>): " << sizeof(optional<sexpr>) << "\n";
     std::cout << "sizeof(std::function):   " << sizeof(std::function<expr(expr const &, optional<expr> const &)>) << "\n";