From c065faaf1f10a180d9a020b1da73e5933fbef3fa Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 18 Feb 2017 16:20:21 -0800
Subject: [PATCH] feat(frontends/lean/elaborator): improve ^. notation

---
 src/frontends/lean/elaborator.cpp | 23 ++++++++++++++---------
 src/frontends/lean/elaborator.h   |  1 +
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/frontends/lean/elaborator.cpp b/src/frontends/lean/elaborator.cpp
index 1cf61f1bbc..fa1a2cc56f 100644
--- a/src/frontends/lean/elaborator.cpp
+++ b/src/frontends/lean/elaborator.cpp
@@ -2410,24 +2410,29 @@ name elaborator::field_to_decl(expr const & e, expr const & s, expr const & s_ty
     }
 }
 
-expr elaborator::visit_field(expr const & e, optional<expr> const & expected_type) {
-    lean_assert(is_field_notation(e));
-    expr s      = visit(macro_arg(e, 0), none_expr());
-    expr s_type = head_beta_reduce(instantiate_mvars(infer_type(s)));
-    name full_fname;
+name elaborator::find_field_fn(expr const & e, expr const & s, expr const & s_type) {
     try {
-        full_fname = field_to_decl(e, s, s_type);
+        return field_to_decl(e, s, s_type);
     } catch (elaborator_exception & ex1) {
-        /* Try again using whnf */
-        expr new_s_type = whnf(s_type);
+        expr new_s_type = s_type;
+        if (auto d = unfold_term(env(), new_s_type))
+            new_s_type = *d;
+        new_s_type = m_ctx.whnf_head_pred(new_s_type, [](expr const &) { return false; });
         if (new_s_type == s_type)
             throw;
         try {
-            full_fname = field_to_decl(e, s, new_s_type);
+            return find_field_fn(e, s, new_s_type);
         } catch (elaborator_exception & ex2) {
             throw nested_elaborator_exception(ex2.get_pos(), ex1, ex2.pp());
         }
     }
+}
+
+expr elaborator::visit_field(expr const & e, optional<expr> const & expected_type) {
+    lean_assert(is_field_notation(e));
+    expr s      = visit(macro_arg(e, 0), none_expr());
+    expr s_type = head_beta_reduce(instantiate_mvars(infer_type(s)));
+    name full_fname = find_field_fn(e, s, s_type);
     expr proj  = copy_tag(e, mk_constant(full_fname));
     expr new_e = copy_tag(e, mk_app(proj, copy_tag(e, mk_as_is(s))));
     return visit(new_e, expected_type);
diff --git a/src/frontends/lean/elaborator.h b/src/frontends/lean/elaborator.h
index 6a312d8b89..3701de08b5 100644
--- a/src/frontends/lean/elaborator.h
+++ b/src/frontends/lean/elaborator.h
@@ -230,6 +230,7 @@ private:
     expr visit_equation(expr const & eq);
     expr visit_inaccessible(expr const & e, optional<expr> const & expected_type);
     name field_to_decl(expr const & e, expr const & s, expr const & s_type);
+    name find_field_fn(expr const & e, expr const & s, expr const & s_type);
     expr visit_field(expr const & e, optional<expr> const & expected_type);
     expr visit_structure_instance(expr const & e, optional<expr> const & expected_type);
     expr visit(expr const & e, optional<expr> const & expected_type);