feat(kernel/expr): cache is_arrow result

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

src/kernel/expr.cpp

@@ -64,6 +64,25 @@ void expr_cell::dec_ref(optional<expr> & c, buffer<expr_cell*> & todelete) {
         dec_ref(*c, todelete);
 }
 
+optional<bool> expr_cell::is_arrow() const {
+    // it is stored in bits 3-4
+    unsigned r = (m_flags & (8+16)) >> 3;
+    if (r == 0) {
+        return optional<bool>();
+    } else if (r == 1) {
+        return optional<bool>(true);
+    } else {
+        lean_assert(r == 2);
+        return optional<bool>(false);
+    }
+}
+
+void expr_cell::set_is_arrow(bool flag) {
+    unsigned mask = flag ? 8 : 16;
+    m_flags |= mask;
+    lean_assert(is_arrow() && *is_arrow() == flag);
+}
+
 expr_var::expr_var(unsigned idx):
     expr_cell(expr_kind::Var, idx, false),
     m_vidx(idx) {}
@@ -227,7 +246,14 @@ bool operator==(expr const & a, expr const & b) {
 }
 
 bool is_arrow(expr const & t) {
-    return is_pi(t) && !has_free_var(abst_body(t), 0);
+    optional<bool> r = t.raw()->is_arrow();
+    if (r) {
+        return *r;
+    } else {
+        bool res = is_pi(t) && !has_free_var(abst_body(t), 0);
+        t.raw()->set_is_arrow(res);
+        return res;
+    }
 }
 
 bool is_eq(expr const & e, expr & lhs, expr & rhs) {

src/kernel/expr.h

@@ -68,6 +68,11 @@ typedef list<local_entry> local_context;
 class expr_cell {
 protected:
     unsigned short     m_kind;
+    // The bits of the following field mean:
+    //    0    - term is maximally shared
+    //    1    - term is closed
+    //    2    - term contains metavariables
+    //    3-4  - term is an arrow (0 - not initialized, 1 - is arrow, 2 - is not arrow)
     atomic_ushort      m_flags;
     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)
@@ -80,6 +85,11 @@ protected:
 
     bool is_closed() const { return (m_flags & 2) != 0; }
     void set_closed() { m_flags |= 2; }
+
+    optional<bool> is_arrow() const;
+    void set_is_arrow(bool flag);
+    friend bool is_arrow(expr const & e);
+
     friend class has_free_var_fn;
     static void dec_ref(expr & c, buffer<expr_cell*> & todelete);
     static void dec_ref(optional<expr> & c, buffer<expr_cell*> & todelete);