fix(frontends/lean/elaborator): bug when elaborating expressions with multiple annotations

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

src/frontends/lean/elaborator.cpp

@@ -964,7 +964,7 @@ public:
 
     bool is_choice_app(expr const & e) {
         expr const & f = get_app_fn(e);
-        return is_choice(f) || (is_explicit(f) && is_choice(get_explicit_arg(f)));
+        return is_choice(f) || (is_annotation(f) && is_choice(get_nested_annotation_arg(f)));
     }
 
     /** \brief Process ((choice f_1 ... f_n) a_1 ... a_k) as
@@ -972,17 +972,14 @@ public:
     */
     expr visit_choice_app(expr const & e, constraint_seq & cs) {
         buffer<expr> args;
-        expr f = get_app_rev_args(e, args);
-        bool expl = is_explicit(f);
-        if (expl)
-            f = get_explicit_arg(f);
+        expr r = get_app_rev_args(e, args);
+        expr f = get_nested_annotation_arg(r);
         lean_assert(is_choice(f));
         buffer<expr> new_choices;
         unsigned num = get_num_choices(f);
         for (unsigned i = 0; i < num; i++) {
             expr f_i = get_choice(f, i);
-            if (expl)
-                f_i = copy_tag(f_i, mk_explicit(f_i));
+            f_i      = copy_annotations(r, f_i);
             new_choices.push_back(mk_rev_app(f_i, args));
         }
         return visit_choice(copy_tag(e, mk_choice(new_choices.size(), new_choices.data())), none_expr(), cs);
@@ -992,7 +989,7 @@ public:
         if (is_choice_app(e))
             return visit_choice_app(e, cs);
         constraint_seq f_cs;
-        bool expl   = is_explicit(get_app_fn(e));
+        bool expl   = is_nested_explicit(get_app_fn(e));
         expr f      = visit(app_fn(e), f_cs);
         auto f_t    = ensure_fun(f, f_cs);
         f           = f_t.first;
@@ -1198,13 +1195,6 @@ public:
         }
     }
 
-    expr visit_extra_info(expr const & e, constraint_seq & cs) {
-        auto ecs = visit(get_annotation_arg(e));
-        cs += ecs.second;
-        save_extra_type_data(e, ecs.first);
-        return ecs.first;
-    }
-
     expr visit_core(expr const & e, constraint_seq & cs) {
         if (is_placeholder(e)) {
             return visit_placeholder(e, cs);
@@ -1219,8 +1209,6 @@ public:
             return visit(get_annotation_arg(e), cs);
         } else if (is_typed_expr(e)) {
             return visit_typed_expr(e, cs);
-        } else if (is_extra_info(e)) {
-            return visit_extra_info(e, cs);
         } else {
             switch (e.kind()) {
             case expr_kind::Local:      return e;
@@ -1238,6 +1226,11 @@ public:
     }
 
     pair<expr, constraint_seq> visit(expr const & e) {
+        if (is_extra_info(e)) {
+            auto ecs = visit(get_annotation_arg(e));
+            save_extra_type_data(e, ecs.first);
+            return ecs;
+        }
         expr r;
         expr b = e;
         constraint_seq cs;

src/library/annotation.cpp

@@ -91,6 +91,39 @@ expr const & get_annotation_arg(expr const & e) {
     return macro_arg(e, 0);
 }
 
+bool is_nested_annotation(expr const & e, name const & kind) {
+    expr const * it = &e;
+    while (is_annotation(*it)) {
+        if (get_annotation_kind(*it) == kind)
+            return true;
+        it = &get_annotation_arg(*it);
+    }
+    return false;
+}
+
+expr const & get_nested_annotation_arg(expr const & e) {
+    expr const * it = &e;
+    while (is_annotation(*it))
+        it = &get_annotation_arg(*it);
+    return *it;
+}
+
+expr copy_annotations(expr const & from, expr const & to) {
+    buffer<expr> trace;
+    expr const * it = &from;
+    while (is_annotation(*it)) {
+        trace.push_back(*it);
+        it = &get_annotation_arg(*it);
+    }
+    expr r     = to;
+    unsigned i = trace.size();
+    while (i > 0) {
+        --i;
+        r = copy_tag(trace[i], mk_annotation(get_annotation_kind(trace[i]), r));
+    }
+    return r;
+}
+
 name const & get_have_name() {
     static name g_have("have");
     static register_annotation_fn g_have_annotation(g_have);

src/library/annotation.h

@@ -29,6 +29,13 @@ expr mk_annotation(name const & kind, expr const & e);
 bool is_annotation(expr const & e);
 /** \brief Return true iff \c e was created using #mk_annotation, and has tag \c kind. */
 bool is_annotation(expr const & e, name const & kind);
+/** \brief Return true iff \c e is of the form (a_1 ... (a_k e') ...)
+    where all a_i's are annotations and one of the is \c kind.
+
+    \remark is_nested_annotation(e, kind) implies is_annotation(e, kind)
+*/
+bool is_nested_annotation(expr const & e, name const & kind);
+
 /** \brief Return the annotated expression, \c e must have been created using #mk_annotation.
 
     \post get_annotation_arg(mk_annotation(k, e)) == e
@@ -40,6 +47,15 @@ expr const & get_annotation_arg(expr const & e);
 */
 name const & get_annotation_kind(expr const & e);
 
+/** \brief Return the nested annotated expression, \c e must have been created using #mk_annotation.
+    This function is the "transitive" version of #get_annotation_arg.
+    It guarantees that the result does not satisfy the predicate is_annotation.
+*/
+expr const & get_nested_annotation_arg(expr const & e);
+
+/** \brief Copy annotation from \c from to \c to. */
+expr copy_annotations(expr const & from, expr const & to);
+
 /** \brief Tag \c e as a 'have'-expression. 'have' is a pre-registered annotation. */
 expr mk_have_annotation(expr const & e);
 /** \brief Tag \c e as a 'show'-expression. 'show' is a pre-registered annotation. */

src/library/explicit.cpp

@@ -33,6 +33,7 @@ static name g_as_is_name    = get_as_is_name();    // force 'as_is' annotation t
 
 expr mk_explicit(expr const & e) { return mk_annotation(get_explicit_name(), e); }
 bool is_explicit(expr const & e) { return is_annotation(e, get_explicit_name()); }
+bool is_nested_explicit(expr const & e) { return is_nested_annotation(e, get_explicit_name()); }
 expr mk_as_is(expr const & e) { return mk_annotation(get_as_is_name(), e); }
 bool is_as_is(expr const & e) { return is_annotation(e, get_as_is_name()); }
 expr mk_implicit(expr const & e) { return mk_annotation(get_implicit_name(), e); }

src/library/explicit.h

@@ -15,6 +15,8 @@ namespace lean {
 expr mk_explicit(expr const & e);
 /** \brief Return true iff \c e is an explicit expression. */
 bool is_explicit(expr const & e);
+/** \brief See #is_nested_annotation */
+bool is_nested_explicit(expr const & e);
 /** \brief Create an explicit expression that is accepted as is
     by the elaborator.
 */