feat(frontends/lean): add notation for 'sep'

library/init/set.lean

@@ -63,7 +63,7 @@ definition set_has_neg : has_neg (set A) :=
 ⟨compl⟩
 
 definition diff (s t : set A) : set A :=
-λ a, a ∈ s ∧ a ∉ t
+{a ∈ s | a ∉ t}
 infix `\`:70 := diff
 
 definition powerset (s : set A) : set (set A) :=

src/frontends/lean/brackets.cpp

@@ -48,6 +48,24 @@ static expr parse_insertable(parser & p, pos_info const & pos, expr const & e) {
     return r;
 }
 
+/* Parse rest of the sep expression '{' id '∈' ... */
+static expr parse_sep(parser & p, pos_info const & pos, name const & id) {
+    expr s = p.parse_expr();
+    p.check_token_next(get_bar_tk(), "invalid sep expression, '|' expected");
+    parser::local_scope scope(p);
+    expr local = p.save_pos(mk_local(id, p.save_pos(mk_expr_placeholder(), pos)), pos);
+    p.add_local(local);
+    expr pred  = p.parse_expr();
+    p.check_token_next(get_rcurly_tk(), "invalid sep expression, '}' expected");
+    pred   = Fun(local, pred);
+    return p.rec_save_pos(mk_app(mk_constant(get_sep_name()), pred, s), pos);
+}
+
+/* Parse rest of the monadic comprehension expression '{' expr '|' ... */
+static expr parse_monadic_comprehension(parser & p, pos_info const & pos, expr const & e) {
+    throw parser_error("monadic comprehension was not implemented yet", pos);
+}
+
 /* Parse rest of the qualified structure instance prefix '{' S '.' ... */
 static expr parse_qualified_structure_instance(parser & p, pos_info const & pos, name const & S) {
     throw parser_error("qualified strucuture instance was not implemented yet", pos);
@@ -63,16 +81,6 @@ static expr parse_structure_instance_update(parser & p, pos_info const & pos, ex
     throw parser_error("strucuture instance update was not implemented yet", pos);
 }
 
-/* Parse rest of the monadic comprehension expression '{' expr '|' ... */
-static expr parse_monadic_comprehension(parser & p, pos_info const & pos, expr const & e) {
-    throw parser_error("monadic comprehension was not implemented yet", pos);
-}
-
-/* Parse rest of the sep expression '{' id '∈' ... */
-static expr parse_sep(parser & p, pos_info const & pos, name const & id) {
-    throw parser_error("sep was not implemented yet", pos);
-}
-
 expr parse_curly_bracket(parser & p, unsigned, expr const *, pos_info const & pos) {
     expr e;
     if (p.curr_is_token(get_rcurly_tk())) {
@@ -98,7 +106,7 @@ expr parse_curly_bracket(parser & p, unsigned, expr const *, pos_info const & po
             return parse_qualified_structure_instance(p, pos, id);
         } else if (p.curr_is_token(get_assign_tk()) || p.curr_is_token(get_fieldarrow_tk())) {
             return parse_structure_instance(p, pos, id);
-        } else if (p.curr_is_token(get_membership_tk())) {
+        } else if (p.curr_is_token(get_membership_tk()) || p.curr_is_token(get_in_tk())) {
             p.next();
             return parse_sep(p, pos, id);
         } else {

src/frontends/lean/pp.cpp

@@ -1510,6 +1510,29 @@ auto pretty_fn::pp_explicit_collection(buffer<expr> const & elems) -> result {
     return result(r);
 }
 
+static bool is_sep(expr const & e) {
+    return
+        is_constant(get_app_fn(e), get_sep_name()) &&
+        get_app_num_args(e) == 5 &&
+        is_lambda(app_arg(app_fn(e)));
+}
+
+auto pretty_fn::pp_sep(expr const & e) -> result {
+    lean_assert(is_sep(e));
+    expr s    = app_arg(e);
+    expr pred = app_arg(app_fn(e));
+    lean_assert(is_lambda(pred));
+    auto p     = binding_body_fresh(pred, true);
+    expr body  = p.first;
+    expr local = p.second;
+    format in  = format(m_unicode ? "∈" : "in");
+    format r   = bracket("{",
+                         format(local_pp_name(local)) + space() + in + space() +
+                         pp_child(s, 0).fmt() + space() + format("|") + space() +
+                         pp_child(body, 0).fmt(), "}");
+    return result(r);
+}
+
 auto pretty_fn::pp(expr const & e, bool ignore_hide) -> result {
     check_system("pretty printer");
     if ((m_depth >= m_max_depth ||
@@ -1538,6 +1561,8 @@ auto pretty_fn::pp(expr const & e, bool ignore_hide) -> result {
     if (m_notation) {
         if (is_subtype(e))
             return pp_subtype(e);
+        if (is_sep(e))
+            return pp_sep(e);
         buffer<expr> elems;
         if (is_explicit_collection(e, elems))
             return pp_explicit_collection(elems);

src/frontends/lean/pp.h

@@ -115,6 +115,7 @@ private:
     result pp_child_core(expr const & e, unsigned bp, bool ignore_hide = false);
     result pp_child(expr const & e, unsigned bp, bool ignore_hide = false);
     result pp_subtype(expr const & e);
+    result pp_sep(expr const & e);
     result pp_explicit_collection(buffer<expr> const & elems);
     result pp_var(expr const & e);
     result pp_sort(expr const & e);

src/library/constants.cpp

@@ -298,6 +298,7 @@ name const * g_real_to_int = nullptr;
 name const * g_rfl = nullptr;
 name const * g_right_distrib = nullptr;
 name const * g_ring = nullptr;
+name const * g_sep = nullptr;
 name const * g_select = nullptr;
 name const * g_semiring = nullptr;
 name const * g_sigma = nullptr;
@@ -655,6 +656,7 @@ void initialize_constants() {
     g_rfl = new name{"rfl"};
     g_right_distrib = new name{"right_distrib"};
     g_ring = new name{"ring"};
+    g_sep = new name{"sep"};
     g_select = new name{"select"};
     g_semiring = new name{"semiring"};
     g_sigma = new name{"sigma"};
@@ -1013,6 +1015,7 @@ void finalize_constants() {
     delete g_rfl;
     delete g_right_distrib;
     delete g_ring;
+    delete g_sep;
     delete g_select;
     delete g_semiring;
     delete g_sigma;
@@ -1370,6 +1373,7 @@ name const & get_real_to_int_name() { return *g_real_to_int; }
 name const & get_rfl_name() { return *g_rfl; }
 name const & get_right_distrib_name() { return *g_right_distrib; }
 name const & get_ring_name() { return *g_ring; }
+name const & get_sep_name() { return *g_sep; }
 name const & get_select_name() { return *g_select; }
 name const & get_semiring_name() { return *g_semiring; }
 name const & get_sigma_name() { return *g_sigma; }

src/library/constants.h

@@ -300,6 +300,7 @@ name const & get_real_to_int_name();
 name const & get_rfl_name();
 name const & get_right_distrib_name();
 name const & get_ring_name();
+name const & get_sep_name();
 name const & get_select_name();
 name const & get_semiring_name();
 name const & get_sigma_name();

src/library/constants.txt

@@ -293,6 +293,7 @@ real.to_int
 rfl
 right_distrib
 ring
+sep
 select
 semiring
 sigma

tests/lean/curly_notation.lean

@@ -10,3 +10,8 @@ print s2
 
 definition s3 : set string := {"hello", "world"}
 print s3
+
+check { a ∈ s1 | a > 1 }
+check { a in s1 | a > 1 }
+set_option pp.unicode false
+check { a ∈ s1 | a > 2 }

tests/lean/curly_notation.lean.expected.out

@@ -0,0 +1,12 @@
+{1, 2, 3} : set ℕ
+{1} : set ℕ
+∅ : set ℕ
+definition s1 : set ℕ :=
+{1, 2 + 3, 3, 4}
+definition s2 : set char :=
+{'a', 'b', 'c'}
+definition s3 : set string :=
+{"hello", "world"}
+{a ∈ s1 | a > 1} : set ℕ
+{a ∈ s1 | a > 1} : set ℕ
+{a in s1 | a > 2} : set nat