feat(frontends/lean/notation_cmd): add 'notation' command

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

src/frontends/lean/builtin_cmds.cpp

@@ -328,6 +328,7 @@ cmd_table init_cmd_table() {
     add_cmd(r, cmd_info("infix",        "declare a new infix (left) notation", infixl_cmd));
     add_cmd(r, cmd_info("infixr",       "declare a new infix (right) notation", infixr_cmd));
     add_cmd(r, cmd_info("postfix",      "declare a new postfix notation", postfix_cmd));
+    add_cmd(r, cmd_info("notation",     "declare a new notation", notation_cmd));
     add_cmd(r, cmd_info("#setline",     "modify the current line number, it only affects error/report messages", set_line_cmd));
     return r;
 }

src/frontends/lean/builtin_exprs.cpp

@@ -40,7 +40,7 @@ parse_table init_nud_table() {
 
 parse_table init_led_table() {
     parse_table r(false);
-    r = r.add({transition("->", mk_expr_action(get_arrow_prec()-1))}, mk_arrow(Var(0), Var(2)));
+    r = r.add({transition("->", mk_expr_action(get_arrow_prec()-1))}, mk_arrow(Var(1), Var(1)));
     return r;
 }
 }

src/frontends/lean/notation_cmd.cpp

@@ -6,12 +6,21 @@ Author: Leonardo de Moura
 */
 #include <limits>
 #include <string>
+#include "kernel/abstract.h"
 #include "frontends/lean/parser.h"
 
 namespace lean {
 static name g_max("max");
 static name g_colon(":");
+static name g_comma(",");
 static name g_assign(":=");
+static name g_lparen("(");
+static name g_rparen(")");
+static name g_scoped("scoped");
+static name g_foldr("foldr");
+static name g_foldl("foldl");
+static name g_binder("binder");
+static name g_binders("binders");
 
 static std::string parse_symbol(parser & p, char const * msg) {
     name n;
@@ -54,8 +63,13 @@ environment precedence_cmd(parser & p) {
 enum class mixfix_kind { infixl, infixr, postfix };
 
 using notation::mk_expr_action;
+using notation::mk_binder_action;
+using notation::mk_binders_action;
+using notation::mk_exprs_action;
+using notation::mk_scoped_expr_action;
 using notation::mk_skip_action;
 using notation::transition;
+using notation::action;
 
 environment mixfix_cmd(parser & p, mixfix_kind k) {
     std::string tk = parse_symbol(p, "invalid notation declaration, quoted symbol or identifier expected");
@@ -77,9 +91,9 @@ environment mixfix_cmd(parser & p, mixfix_kind k) {
     char const * tks = tk.c_str();
     switch (k) {
     case mixfix_kind::infixl:
-        return add_led_notation(env, {transition(tks, mk_expr_action(*prec))}, mk_app(f, Var(0), Var(1)));
+        return add_led_notation(env, {transition(tks, mk_expr_action(*prec))}, mk_app(f, Var(1), Var(0)));
     case mixfix_kind::infixr:
-        return add_led_notation(env, {transition(tks, mk_expr_action(*prec-1))}, mk_app(f, Var(0), Var(1)));
+        return add_led_notation(env, {transition(tks, mk_expr_action(*prec-1))}, mk_app(f, Var(1), Var(0)));
     case mixfix_kind::postfix:
         return add_led_notation(env, {transition(tks, mk_skip_action())}, mk_app(f, Var(0)));
     }
@@ -90,8 +104,138 @@ environment infixl_cmd(parser & p) { return mixfix_cmd(p, mixfix_kind::infixl);
 environment infixr_cmd(parser & p) { return mixfix_cmd(p, mixfix_kind::infixr); }
 environment postfix_cmd(parser & p) { return mixfix_cmd(p, mixfix_kind::postfix); }
 
+static name parse_quoted_symbol(parser & p, environment & env) {
+    if (p.curr_is_quoted_symbol()) {
+        auto tk   = p.get_name_val();
+        auto tks  = tk.to_string();
+        auto tkcs = tks.c_str();
+        p.next();
+        if (p.curr_is_token(g_colon)) {
+            p.next();
+            unsigned prec = parse_precedence(p, "invalid notation declaration, precedence (small numeral) expected");
+            auto old_prec = get_precedence(get_token_table(env), tkcs);
+            if (!old_prec || prec != *old_prec)
+                env = add_token(env, tkcs, prec);
+        } else if (!get_precedence(get_token_table(env), tkcs)) {
+            env = add_token(env, tkcs, 0);
+        }
+        return tk;
+    } else {
+        throw parser_error("invalid notation declaration, quoted symbol expected", p.pos());
+    }
+}
+
+static expr parse_notation_expr(parser & p, buffer<expr> const & locals) {
+    expr r = p.parse_expr();
+    return abstract(r, locals.size(), locals.data());
+}
+
+static expr g_local_type = Bool; // type used in notation local declarations, it is irrelevant
+
+static void parse_notation_local(parser & p, buffer<expr> & locals) {
+    if (p.curr_is_identifier()) {
+        name n = p.get_name_val();
+        p.next();
+        expr l = mk_local(n, n, g_local_type); // remark: the type doesn't matter
+        p.add_local_expr(n, l);
+        locals.push_back(l);
+    } else {
+        throw parser_error("invalid notation declaration, identifier expected", p.pos());
+    }
+}
+
+action parse_action(parser & p, environment & env, buffer<expr> & locals) {
+    if (p.curr_is_token(g_colon)) {
+        p.next();
+        if (p.curr_is_numeral()) {
+            unsigned prec = parse_precedence(p, "invalid notation declaration, small numeral expected");
+            return mk_expr_action(prec);
+        } else if (p.curr_is_token_or_id(g_scoped)) {
+            p.next();
+            return mk_scoped_expr_action(mk_var(0));
+        } else {
+            p.check_token_next(g_lparen, "invalid notation declaration, '(', numeral or 'scoped' expected");
+            if (p.curr_is_token_or_id(g_foldl) || p.curr_is_token_or_id(g_foldr)) {
+                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, env);
+                expr rec;
+                {
+                    parser::local_scope scope(p);
+                    p.check_token_next(g_lparen, "invalid fold notation argument, '(' expected");
+                    parse_notation_local(p, locals);
+                    parse_notation_local(p, locals);
+                    p.check_token_next(g_comma,  "invalid fold notation argument, ',' expected");
+                    rec  = parse_notation_expr(p, locals);
+                    p.check_token_next(g_rparen, "invalid fold notation argument, ')' expected");
+                    locals.pop_back();
+                    locals.pop_back();
+                }
+                expr ini  = parse_notation_expr(p, locals);
+                p.check_token_next(g_rparen, "invalid fold notation argument, ')' expected");
+                return mk_exprs_action(sep, rec, ini, is_fold_right, prec ? *prec : 0);
+            } else if (p.curr_is_token_or_id(g_scoped)) {
+                p.next();
+                auto prec = parse_optional_precedence(p);
+                expr rec;
+                {
+                    parser::local_scope scope(p);
+                    parse_notation_local(p, locals);
+                    p.check_token_next(g_comma,  "invalid scoped notation argument, ',' expected");
+                    rec  = parse_notation_expr(p, locals);
+                    locals.pop_back();
+                }
+                p.check_token_next(g_rparen, "invalid scoped notation argument, ')' expected");
+                return mk_scoped_expr_action(rec, prec ? *prec : 0);
+            } else {
+                throw parser_error("invalid notation declaration, 'foldl', 'foldr' or 'scoped' expected", p.pos());
+            }
+        }
+    } else {
+        return mk_expr_action();
+    }
+}
+
 environment notation_cmd(parser & p) {
-    // TODO(Leo)
+    environment env = p.env();
-    return p.env();
+    buffer<expr>       locals;
+    buffer<transition> ts;
+    parser::local_scope scope(p);
+    bool is_nud = true;
+    if (p.curr_is_identifier()) {
+        parse_notation_local(p, locals);
+        is_nud = false;
+    }
+    while (!p.curr_is_token(g_assign)) {
+        name tk = parse_quoted_symbol(p, env);
+        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)) {
+            p.next();
+            ts.push_back(transition(tk, mk_binder_action()));
+        } else if (p.curr_is_token_or_id(g_binders)) {
+            p.next();
+            ts.push_back(transition(tk, mk_binders_action()));
+        } else if (p.curr_is_identifier()) {
+            name n   = p.get_name_val();
+            p.next();
+            action a = parse_action(p, env, locals);
+            expr l = mk_local(n, n, g_local_type);
+            p.add_local_expr(n, l);
+            locals.push_back(l);
+            ts.push_back(transition(tk, a));
+        } else {
+            throw parser_error("invalid notation declaration, quoted-symbol, identifier, 'binder', 'binders' expected", p.pos());
+        }
+    }
+    p.next();
+    if (ts.empty())
+        throw parser_error("invalid notation declaration, empty notation is not allowed", p.pos());
+    expr n = parse_notation_expr(p, locals);
+    if (is_nud)
+        return add_nud_notation(env, ts.size(), ts.data(), n);
+    else
+        return add_led_notation(env, ts.size(), ts.data(), n);
 }
 }

src/frontends/lean/parser.cpp

@@ -164,7 +164,32 @@ expr parser::rec_save_pos(expr const & e, pos_info p) {
 
 /** \brief Create a copy of \c e, and the position of new expression with p */
 expr parser::copy_with_new_pos(expr const & e, pos_info p) {
-    return rec_save_pos(deep_copy(e), p);
+    switch (e.kind()) {
+    case expr_kind::Sort: case expr_kind::Constant: case expr_kind::Meta:
+    case expr_kind::Local: case expr_kind::Var:
+        return save_pos(copy(e), p);
+    case expr_kind::App:
+        return save_pos(::lean::mk_app(copy_with_new_pos(app_fn(e), p),
+                                       copy_with_new_pos(app_arg(e), p)),
+                        p);
+    case expr_kind::Lambda: case expr_kind::Pi:
+        return save_pos(update_binding(e,
+                                       copy_with_new_pos(binding_domain(e), p),
+                                       copy_with_new_pos(binding_body(e), p)),
+                        p);
+    case expr_kind::Let:
+        return save_pos(update_let(e,
+                                   copy_with_new_pos(let_type(e), p),
+                                   copy_with_new_pos(let_value(e), p),
+                                   copy_with_new_pos(let_body(e), p)),
+                        p);
+    case expr_kind::Macro: {
+        buffer<expr> args;
+        for (unsigned i = 0; i < macro_num_args(e); i++)
+            args.push_back(copy_with_new_pos(macro_arg(e, i), p));
+        return save_pos(update_macro(e, args.size(), args.data()), p);
+    }}
+    lean_unreachable(); // LCOV_EXCL_LINE
 }
 
 /** \brief Add \c ls to constants occurring in \c e. */
@@ -548,23 +573,24 @@ expr parser::parse_notation(parse_table t, expr * left) {
             buffer<expr> r_args;
             r_args.push_back(parse_expr(a.rbp()));
             while (curr_is_token(a.get_sep())) {
+                next();
                 r_args.push_back(parse_expr(a.rbp()));
             }
-            expr r = instantiate(a.get_initial(), args.size(), args.data());
+            expr r = instantiate_rev(a.get_initial(), args.size(), args.data());
             if (a.is_fold_right()) {
                 unsigned i = r_args.size();
-                while (!i > 0) {
+                while (i > 0) {
                     --i;
                     args.push_back(r_args[i]);
                     args.push_back(r);
-                    r = instantiate(a.get_rec(), args.size(), args.data());
+                    r = instantiate_rev(a.get_rec(), args.size(), args.data());
                     args.pop_back(); args.pop_back();
                 }
             } else {
                 for (unsigned i = 0; i < r_args.size(); i++) {
                     args.push_back(r_args[i]);
                     args.push_back(r);
-                    r = instantiate(a.get_rec(), args.size(), args.data());
+                    r = instantiate_rev(a.get_rec(), args.size(), args.data());
                     args.pop_back(); args.pop_back();
                 }
             }
@@ -591,7 +617,7 @@ expr parser::parse_notation(parse_table t, expr * left) {
                     else
                         r = Pi(l, r, ps[i].m_bi);
                     args.push_back(r);
-                    r = instantiate(a.get_rec(), args.size(), args.data());
+                    r = instantiate_rev(a.get_rec(), args.size(), args.data());
                     args.pop_back();
                 }
             }
@@ -613,7 +639,8 @@ expr parser::parse_notation(parse_table t, expr * left) {
     }
     buffer<expr> cs;
     for (expr const & a : as) {
-        cs.push_back(instantiate(a, args.size(), args.data()));
+        expr r = instantiate_rev(copy_with_new_pos(a, p), args.size(), args.data());
+        cs.push_back(r);
     }
     return mk_choice(cs.size(), cs.data());
 }

src/frontends/lean/token_table.cpp

@@ -65,7 +65,7 @@ token_table init_token_table() {
                                "variables", "{variables}", "[variables]", "[private]", "[inline]", "abbreviation",
                                "evaluate", "check", "print", "end", "namespace", "section", "import",
                                "abbreviation", "inductive", "record", "structure", "module", "universe",
-                               "precedence", "infixl", "infixr", "infix", "postfix", "#setline", nullptr};
+                               "precedence", "infixl", "infixr", "infix", "postfix", "notation", "#setline", nullptr};
 
     std::pair<char const *, char const *> aliases[] =
         {{g_lambda_unicode, "fun"}, {"forall", "Pi"}, {g_forall_unicode, "Pi"}, {g_pi_unicode, "Pi"},

src/kernel/instantiate.cpp

@@ -40,6 +40,27 @@ expr instantiate(expr const & e, std::initializer_list<expr> const & l) {  retur
 expr instantiate(expr const & e, unsigned i, expr const & s) { return instantiate(e, i, 1, &s); }
 expr instantiate(expr const & e, expr const & s) { return instantiate(e, 0, s); }
 
+expr instantiate_rev(expr const & a, unsigned n, expr const * subst) {
+    if (closed(a))
+        return a;
+    return replace(a, [=](expr const & m, unsigned offset) -> optional<expr> {
+            if (offset >= get_free_var_range(m))
+                return some_expr(m); // expression m does not contain free variables with idx >= offset
+            if (is_var(m)) {
+                unsigned vidx = var_idx(m);
+                if (vidx >= offset) {
+                    unsigned h = offset + n;
+                    if (h < offset /* overflow, h is bigger than any vidx */ || vidx < h) {
+                        return some_expr(lift_free_vars(subst[n - (vidx - offset) - 1], offset));
+                    } else {
+                        return some_expr(mk_var(vidx - n));
+                    }
+                }
+            }
+            return none_expr();
+        });
+}
+
 bool is_head_beta(expr const & t) {
     expr const * it = &t;
     while (is_app(*it)) {

src/kernel/instantiate.h

@@ -18,6 +18,9 @@ expr instantiate(expr const & e, unsigned i, expr const & s);
 /** \brief Replace free variable \c 0 with \c s in \c e. */
 expr instantiate(expr const & e, expr const & s);
 
+/** \brief Replace the free variables with indices 0, ..., n-1 with s[n-1], ..., s[0] in e. */
+expr instantiate_rev(expr const & e, unsigned n, expr const * s);
+
 expr apply_beta(expr f, unsigned num_args, expr const * args);
 bool is_head_beta(expr const & t);
 expr head_beta_reduce(expr const & t);

tests/lean/t10.lean

@@ -0,0 +1,23 @@
+variable N : Type.{1}
+definition B : Type.{1} := Type.{0}
+variable ite : B → N → N → N
+variable and : B → B → B
+variable f : N → N
+variable p : B
+variable q : B
+variable x : N
+variable y : N
+variable z : N
+infixr `∧` 25 := and
+notation `if` c `then` t `else` e := ite c t e
+check if p ∧ q then f x else y
+check if p ∧ q then q else y
+variable list : Type.{1}
+variable nil : list
+variable cons : N → list → list
+-- Non empty lists
+notation `[` l:(foldr `,` (h t, cons h t) nil) `]` := l
+check [x, y, z, x, y, y]
+check [x]
+notation `[` `]` := nil
+check []

tests/lean/t10.lean.expected.out

@@ -0,0 +1,10 @@
+ite (and p q) (f x) y : N
+t10.lean:14:22: error: type mismatch at application
+    ite (and p q) q
+expected type:
+    N
+given type:
+    B
+cons x (cons y (cons z (cons x (cons y (cons y nil))))) : list
+cons x nil : list
+nil : list

tests/lean/t11.lean

@@ -0,0 +1,11 @@
+variable A    : Type.{1}
+definition bool : Type.{1} := Type.{0}
+variable Exists (P : A → bool) : bool
+notation `exists` binders `,` b:(scoped b, Exists b) := b
+notation `∃` binders `,` b:(scoped b, Exists b) := b
+variable p : A → bool
+variable q : A → A → bool
+check exists x : A, p x
+check ∃ x y : A, q x y
+notation `{` binder `|` b:scoped `}` := b
+check {x : A | x}

tests/lean/t11.lean.expected.out

@@ -0,0 +1,3 @@
+Exists (fun (x : A), (p x)) : bool
+Exists (fun (x : A), (Exists (fun (y : A), (q x y)))) : bool
+fun (x : A), x : A -> A