feat(frontends/lean/notation_cmd): improve 'notation' cmd

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

src/frontends/lean/notation_cmd.cpp

@@ -135,7 +135,7 @@ environment infixl_cmd_core(parser & p, bool overload) { return mixfix_cmd(p, mi
 environment infixr_cmd_core(parser & p, bool overload) { return mixfix_cmd(p, mixfix_kind::infixr, overload); }
 environment postfix_cmd_core(parser & p, bool overload) { return mixfix_cmd(p, mixfix_kind::postfix, overload); }
 
-static name parse_quoted_symbol(parser & p, buffer<token_entry> & new_tokens) {
+static name parse_quoted_symbol_or_token(parser & p, buffer<token_entry> & new_tokens) {
     if (p.curr_is_quoted_symbol()) {
         environment const & env = p.env();
         auto tk   = p.get_name_val();
@@ -152,8 +152,12 @@ static name parse_quoted_symbol(parser & p, buffer<token_entry> & new_tokens) {
             new_tokens.push_back(token_entry(tkcs, 0));
         }
         return tk;
+    } else if (p.curr_is_keyword()) {
+        auto tk = p.get_token_info().token();
+        p.next();
+        return tk;
     } else {
-        throw parser_error("invalid notation declaration, quoted symbol expected", p.pos());
+        throw parser_error("invalid notation declaration, symbol expected", p.pos());
     }
 }
 
@@ -176,7 +180,7 @@ static void parse_notation_local(parser & p, buffer<expr> & locals) {
     }
 }
 
-static action parse_action(parser & p, buffer<expr> & locals, buffer<token_entry> & new_tokens) {
+static action parse_action(parser & p, name const & prev_token, buffer<expr> & locals, buffer<token_entry> & new_tokens) {
     if (p.curr_is_token(g_colon)) {
         p.next();
         if (p.curr_is_numeral()) {
@@ -195,7 +199,7 @@ static action parse_action(parser & p, buffer<expr> & locals, buffer<token_entry
                 bool is_fold_right = p.curr_is_token_or_id(g_foldr);
                 p.next();
                 auto prec = parse_optional_precedence(p);
-                name sep  = parse_quoted_symbol(p, new_tokens);
+                name sep  = parse_quoted_symbol_or_token(p, new_tokens);
                 expr rec;
                 {
                     parser::local_scope scope(p);
@@ -234,7 +238,11 @@ static action parse_action(parser & p, buffer<expr> & locals, buffer<token_entry
             }
         }
     } else {
-        return mk_expr_action();
+        auto prec = get_precedence(get_token_table(p.env()), prev_token.to_string().c_str());
+        if (prec)
+            return mk_expr_action(*prec);
+        else
+            return mk_expr_action();
     }
 }
 
@@ -248,10 +256,8 @@ notation_entry parse_notation_core(parser & p, bool overload, buffer<token_entry
         is_nud = false;
     }
     while (!p.curr_is_token(g_assign)) {
-        name tk = parse_quoted_symbol(p, new_tokens);
-        if (p.curr_is_quoted_symbol() || p.curr_is_token(g_assign)) {
-            ts.push_back(transition(tk, mk_skip_action()));
-        } else if (p.curr_is_token_or_id(g_binder)) {
+        name tk = parse_quoted_symbol_or_token(p, new_tokens);
+        if (p.curr_is_token_or_id(g_binder)) {
             p.next();
             ts.push_back(transition(tk, mk_binder_action()));
         } else if (p.curr_is_token_or_id(g_binders)) {
@@ -260,11 +266,13 @@ notation_entry parse_notation_core(parser & p, bool overload, buffer<token_entry
         } else if (p.curr_is_identifier()) {
             name n   = p.get_name_val();
             p.next();
-            action a = parse_action(p, locals, new_tokens);
+            action a = parse_action(p, tk, locals, new_tokens);
             expr l = mk_local(n, g_local_type);
             p.add_local_expr(n, l);
             locals.push_back(l);
             ts.push_back(transition(tk, a));
+        } else if (p.curr_is_quoted_symbol() || p.curr_is_keyword() || p.curr_is_token(g_assign)) {
+            ts.push_back(transition(tk, mk_skip_action()));
         } else {
             throw parser_error("invalid notation declaration, quoted-symbol, identifier, 'binder', 'binders' expected", p.pos());
         }

src/frontends/lean/parser.cpp

@@ -7,6 +7,7 @@ Author: Leonardo de Moura
 #include <utility>
 #include <string>
 #include <limits>
+#include <vector>
 #include "util/interrupt.h"
 #include "util/script_exception.h"
 #include "util/sstream.h"

src/frontends/lean/token_table.cpp

@@ -18,13 +18,13 @@ token_table add_command_token(token_table const & s, char const * token) {
     return insert(s, token, token_info(token));
 }
 token_table add_command_token(token_table const & s, char const * token, char const * val) {
-    return insert(s, token, token_info(val));
+    return insert(s, token, token_info(token, val));
 }
 token_table add_token(token_table const & s, char const * token, unsigned prec) {
     return insert(s, token, token_info(token, prec));
 }
 token_table add_token(token_table const & s, char const * token, char const * val, unsigned prec) {
-    return insert(s, token, token_info(val, prec));
+    return insert(s, token, token_info(token, val, prec));
 }
 token_table const * find(token_table const & s, char c) {
     return s.find(c);

src/frontends/lean/token_table.h

@@ -15,13 +15,21 @@ namespace lean {
 unsigned get_arrow_prec();
 class token_info {
     bool        m_command;
+    name        m_token;
     name        m_value;
     unsigned    m_precedence;
 public:
     token_info():m_command(true) {}
-    token_info(char const * val):m_command(true), m_value(val), m_precedence(0) {}
-    token_info(char const * val, unsigned prec):m_command(false), m_value(val), m_precedence(prec) {}
+    token_info(char const * val):
+        m_command(true), m_token(val), m_value(val), m_precedence(0) {}
+    token_info(char const * token, char const * val):
+        m_command(true), m_token(token), m_value(val), m_precedence(0) {}
+    token_info(char const * val, unsigned prec):
+        m_command(false), m_token(val), m_value(val), m_precedence(prec) {}
+    token_info(char const * token, char const * val, unsigned prec):
+        m_command(false), m_token(token), m_value(val), m_precedence(prec) {}
     bool is_command() const { return m_command; }
+    name const & token() const { return m_token; }
     name const & value() const { return m_value; }
     unsigned precedence() const { return m_precedence; }
 };

tests/lean/run/n3.lean

@@ -0,0 +1,19 @@
+definition [inline] Bool : Type.{1} := Type.{0}
+variable N : Type.{1}
+variable and : Bool → Bool → Bool
+infixr `∧` 35 := and
+variable le : N → N → Bool
+variable lt : N → N → Bool
+precedence `≤`:50
+precedence `<`:50
+infixl ≤ := le
+infixl < := lt
+notation A ≤ B ≤ C := A ≤ B ∧ B ≤ C
+notation A ≤ B < C := A ≤ B ∧ B < C
+notation A < B ≤ C := A < B ∧ B ≤ C
+variables a b c d e : N
+check a ≤ b ≤ c
+check a ≤ d
+check a < b ≤ c
+check a ≤ b < c
+check a < b