Avoid head and tail when manipulating contexts

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

src/kernel/context.cpp

@@ -58,11 +58,22 @@ std::ostream & operator<<(std::ostream & out, context const & c) {
     return out;
 }
 
-context const & lookup(context const & c, unsigned i) {
+std::pair<context_entry const &, context const &> lookup_ext(context const & c, unsigned i) {
     context const * it1 = &c;
     while (*it1) {
         if (i == 0)
-            return *it1;
+            return std::pair<context_entry const &, context const &>(head(*it1), tail(*it1));
+        --i;
+        it1 = &tail(*it1);
+    }
+    throw exception("unknown free variable");
+}
+
+context_entry const & lookup(context const & c, unsigned i) {
+    context const * it1 = &c;
+    while (*it1) {
+        if (i == 0)
+            return head(*it1);
         --i;
         it1 = &tail(*it1);
     }

src/kernel/context.h

@@ -25,7 +25,17 @@ public:
     expr const & get_domain() const { return m_domain; }
     expr const & get_body() const   { return m_body; }
 };
-context const & lookup(context const & c, unsigned i);
+/**
+   \brief Return the context entry for the free variable with de
+   Bruijn index \c i, and the context for this entry.
+*/
+std::pair<context_entry const &, context const &> lookup_ext(context const & c, unsigned i);
+/**
+   \brief Return the context entry for the free variable with de
+   Bruijn index \c i.
+*/
+context_entry const & lookup(context const & c, unsigned i);
+
 inline context extend(context const & c, name const & n, expr const & d, expr const & b) {
     return context(context_entry(n, d, b), c);
 }

src/kernel/expr_formatter.cpp

@@ -82,8 +82,7 @@ class simple_expr_formatter : public expr_formatter {
         switch (a.kind()) {
         case expr_kind::Var:
             try {
-                context const & c1 = lookup(c, var_idx(a));
-                out() << head(c1).get_name();
+                out() << lookup(c, var_idx(a)).get_name();
             } catch (exception & ex) {
                 out() << "#" << var_idx(a);
             }

src/kernel/normalize.cpp

@@ -64,15 +64,13 @@ class normalize_fn {
             --j;
             it1 = &tail(*it1);
         }
-        context const & c = ::lean::lookup(m_ctx, j);
-        if (c) {
-            context_entry const & entry = head(c);
+        auto p = lookup_ext(m_ctx, j);
+        context_entry const & entry = p.first;
+        context const & entry_c     = p.second;
         if (entry.get_body())
-                return svalue(::lean::normalize(entry.get_body(), m_env, tail(c)));
+            return svalue(::lean::normalize(entry.get_body(), m_env, entry_c));
         else
-                return svalue(length(c) - 1);
-        }
-        throw exception("unknown free variable");
+            return svalue(length(entry_c));
     }
 
     /** \brief Convert the closure \c a into an expression using the given stack in a context that contains \c k binders. */

src/kernel/type_check.cpp

@@ -50,10 +50,11 @@ struct infer_type_fn {
     cache               m_cache;
 
     expr lookup(context const & c, unsigned i) {
-        context const & def_c = ::lean::lookup(c, i);
-        lean_assert(length(c) >= length(def_c));
-        lean_assert(length(def_c) > 0);
-        return lift_free_vars(head(def_c).get_domain(), length(c) - (length(def_c) - 1));
+        auto p = lookup_ext(c, i);
+        context_entry const & def = p.first;
+        context const & def_c     = p.second;
+        lean_assert(length(c) > length(def_c));
+        return lift_free_vars(def.get_domain(), length(c) - length(def_c));
     }
 
     expr_formatter & fmt() { return m_env.get_formatter(); }

src/tests/kernel/occurs.cpp

@@ -44,7 +44,13 @@ static void tst2() {
     expr b = Const("b");
     context c;
     c = extend(c, "a", Type());
+    lean_assert(length(c) == 1);
+    lean_assert(lookup(c, 0).get_name() == "a");
+    auto p = lookup_ext(c, 0);
+    lean_assert(p.first.get_name() == "a");
+    lean_assert(length(p.second) == 0);
     std::cout << sanitize_names(c, f(a)) << "\n";
+    lean_assert(lookup(sanitize_names(c, f(a)), 0).get_name() != name("a"));
     std::cout << sanitize_names(c, f(b)) << "\n";
 }
 

src/tests/util/list.cpp

@@ -40,6 +40,9 @@ static void tst3() {
     lean_assert(head(tail(l)) == 20);
     lean_assert(head(tail(tail(l))) == 30);
     lean_assert(length(l) == 3);
+    lean_assert(length(list<int>()) == 0);
+    lean_assert(length(list<int>(10, list<int>())) == 1);
+    lean_assert(length(tail(list<int>(10, list<int>()))) == 0);
 }
 
 int main() {