From fc0230730d162dcff5458c8d0992f3b2f8aa2645 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Fri, 9 Sep 2016 12:11:35 -0700
Subject: [PATCH] feat(frontends/lean/elaborator): make sure all equations have
 the same number of patterns

---
 src/frontends/lean/elaborator.cpp        | 39 ++++++++++++++++++++++--
 tests/lean/wrong_arity.lean              |  6 ++++
 tests/lean/wrong_arity.lean.expected.out |  1 +
 3 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 tests/lean/wrong_arity.lean
 create mode 100644 tests/lean/wrong_arity.lean.expected.out

diff --git a/src/frontends/lean/elaborator.cpp b/src/frontends/lean/elaborator.cpp
index d1fe6f97c6..4f0404ff3b 100644
--- a/src/frontends/lean/elaborator.cpp
+++ b/src/frontends/lean/elaborator.cpp
@@ -1332,6 +1332,40 @@ static expr copy_domain(unsigned num, expr const & source, expr const & target)
     }
 }
 
+static void check_equations_arity(buffer<expr> const & eqns) {
+    buffer<optional<unsigned>> fidx2arity;
+    for (expr eqn : eqns) {
+        unsigned nbinders = 0;
+        expr curr = eqn;
+        while (is_lambda(eqn)) {
+            nbinders++;
+            eqn = binding_body(eqn);
+        }
+        if (is_equation(eqn)) {
+            expr const & lhs = equation_lhs(eqn);
+            expr const & fn  = get_app_fn(lhs);
+            unsigned arity   = get_app_num_args(lhs);
+            if (!is_var(fn) || var_idx(fn) >= nbinders)
+                throw elaborator_exception(eqn, "ill-formed match/equations expression");
+            unsigned fidx    = nbinders - var_idx(fn) - 1;
+            if (fidx >= fidx2arity.size())
+                fidx2arity.resize(fidx+1, optional<unsigned>());
+            if (auto r = fidx2arity[fidx]) {
+                if (*r != arity) {
+                    throw elaborator_exception(eqn, "invalid match/equations expression, "
+                                               "each case must have the same number of patterns");
+                }
+            } else {
+                fidx2arity[fidx] = arity;
+            }
+        } else if (is_no_equation(eqn)) {
+            /* do nothing */
+        } else {
+            throw elaborator_exception(curr, "ill-formed match/equations expression");
+        }
+    }
+}
+
 expr elaborator::visit_equations(expr const & e) {
     expr const & ref = e;
     buffer<expr> eqs;
@@ -1361,13 +1395,14 @@ expr elaborator::visit_equations(expr const & e) {
             // Replace first num_fns domains of eq with the ones in first_eq.
             // This is a trick/hack to ensure the fns in each equation have
             // the same elaborated type.
-            new_eq   = visit(copy_domain(num_fns, *first_eq, eq), none_expr());
+            new_eq   = copy_tag(eq, visit(copy_domain(num_fns, *first_eq, eq), none_expr()));
         } else {
-            new_eq   = visit(eq, none_expr());
+            new_eq   = copy_tag(eq, visit(eq, none_expr()));
             first_eq = new_eq;
         }
         new_eqs.push_back(new_eq);
     }
+    check_equations_arity(new_eqs);
     process_checkpoint(C);
     expr new_e;
     if (new_R) {
diff --git a/tests/lean/wrong_arity.lean b/tests/lean/wrong_arity.lean
new file mode 100644
index 0000000000..deb4aa8e4f
--- /dev/null
+++ b/tests/lean/wrong_arity.lean
@@ -0,0 +1,6 @@
+set_option new_elaborator true
+open nat
+
+theorem succ_ne_self : ∀ (n : ℕ), succ n ≠ n
+| 0 h := absurd h (nat.succ_ne_zero 0)
+| n   := sorry
diff --git a/tests/lean/wrong_arity.lean.expected.out b/tests/lean/wrong_arity.lean.expected.out
new file mode 100644
index 0000000000..93715733ad
--- /dev/null
+++ b/tests/lean/wrong_arity.lean.expected.out
@@ -0,0 +1 @@
+wrong_arity.lean:6:6: error: invalid match/equations expression, each case must have the same number of patterns