feat(frontends/lean/parse_table): add get_head_index auxiliary function for indexing notation declarations

src/frontends/lean/parse_table.cpp

@@ -9,6 +9,7 @@ Author: Leonardo de Moura
 #include "util/rb_map.h"
 #include "util/sstream.h"
 #include "kernel/free_vars.h"
+#include "kernel/replace_fn.h"
 #include "library/kernel_bindings.h"
 #include "frontends/lean/parse_table.h"
 #include "frontends/lean/no_info.h"
@@ -201,6 +202,9 @@ void action::display(std::ostream & out) const {
         break;
     }
 }
+bool action::is_simple() const {
+    return kind() != action_kind::Ext && kind() != action_kind::LuaExt;
+}
 
 void action_cell::dealloc() {
     switch (m_kind) {
@@ -351,6 +355,72 @@ parse_table parse_table::add_core(unsigned num, transition const * ts, expr cons
     return r;
 }
 
+static bool is_simple(unsigned num, transition const * ts) {
+    return std::all_of(ts, ts+num, [](transition const & t) { return t.is_simple(); });
+}
+
+bool is_safe_ascii(unsigned num, transition const * ts) {
+    return std::all_of(ts, ts+num, [](transition const & t) { return t.is_safe_ascii(); });
+}
+
+/** \brief Given \c a, an expression that is the denotation of an expression, if \c a is a variable,
+    then use the actions in the transitions \c ts to expand \c a. The idea is to produce a head symbol
+    we can use to decide whether the notation should be considered during pretty printing.
+
+    \see get_head_index
+*/
+static expr expand_pp_pattern(unsigned num, transition const * ts, expr const & a) {
+    lean_assert(is_simple(num, ts));
+    if (!is_var(a))
+        return a;
+    return replace(a, [&](expr const & e) {
+            if (is_var(e)) {
+                unsigned vidx = var_idx(e);
+                unsigned i = num;
+                unsigned offset = 0;
+                while (i > 0) {
+                    --i;
+                    action const & act = ts[i].get_action();
+                    switch (act.kind()) {
+                    case action_kind::Binder: case action_kind::Binders: case action_kind::Skip:
+                        break;
+                    case action_kind::Ext: case action_kind::LuaExt:
+                        lean_unreachable();
+                    case action_kind::Expr:
+                        if (vidx == 0) return none_expr();
+                        offset++;
+                        vidx--;
+                        break;
+                    case action_kind::Exprs:
+                        if (vidx == 0)
+                            return some_expr(lift_free_vars(act.get_rec(), offset));
+                        offset++;
+                        vidx--;
+                        break;
+                    case action_kind::ScopedExpr:
+                        if (vidx == 0)
+                            return some_expr(lift_free_vars(act.get_rec(), offset));
+                        offset++;
+                        vidx--;
+                        break;
+                    }
+                }
+                return none_expr();
+            } else {
+                return none_expr();
+            }
+        });
+}
+
+optional<head_index> get_head_index(unsigned num, transition const * ts, expr const & a) {
+    if (is_simple(num, ts)) {
+        expr n = expand_pp_pattern(num, ts, a);
+        if (!is_var(n))
+            return some(head_index(n));
+    }
+    return optional<head_index>();
+}
+
 parse_table parse_table::add(unsigned num, transition const * ts, expr const & a, bool overload) const {
     expr new_a = annotate_macro_subterms(a);
     validate_transitions(is_nud(), num, ts, new_a);

src/frontends/lean/parse_table.h

@@ -10,6 +10,7 @@ Author: Leonardo de Moura
 #include "util/buffer.h"
 #include "util/lua.h"
 #include "kernel/expr.h"
+#include "library/head_map.h"
 #include "frontends/lean/token_table.h"
 #include "frontends/lean/parser_pos_provider.h"
 
@@ -78,6 +79,9 @@ public:
 
     bool is_equal(action const & a) const;
     void display(std::ostream & out) const;
+
+    /** \brief Return true iff the action is not Ext or LuaExt */
+    bool is_simple() const;
 };
 
 action mk_skip_action();
@@ -101,8 +105,12 @@ public:
         m_token(t), m_action(a) {}
     name const & get_token() const { return m_token; }
     action const & get_action() const { return m_action; }
+    bool is_simple() const { return m_action.is_simple(); }
+    bool is_safe_ascii() const { return m_token.is_safe_ascii(); }
 };
 
+bool is_safe_ascii(unsigned num, transition const * ts);
+
 /** \brief Apply \c f to expressions embedded in the given transition */
 transition replace(transition const & t, std::function<expr(expr const &)> const & f);
 
@@ -136,6 +144,15 @@ public:
 
     void display(std::ostream & out) const;
 };
+
+/** \brief Given a notation definition, return the "head symbol" of
+    every instance of the given notation.
+
+    \remark The result is none if the notation uses actions implemented in C++ or Lua.
+    The result is none if the denotation is a variable.
+*/
+optional<head_index> get_head_index(unsigned num, transition const * ts, expr const & a);
+
 void initialize_parse_table();
 void finalize_parse_table();
 }