feat(library/compiler): add reduce_arity compilation step

src/library/compiler/CMakeLists.txt

@@ -3,4 +3,4 @@ add_library(compiler OBJECT
   ## New compiler
   util.cpp lcnf.cpp csimp.cpp elim_dead_let.cpp cse.cpp erase_irrelevant.cpp
   specialize.cpp compiler.cpp lambda_lifting.cpp extract_closed.cpp
-  simp_app_args.cpp llnf.cpp ll_infer_type.cpp ir.cpp)
+  simp_app_args.cpp llnf.cpp ll_infer_type.cpp reduce_arity.cpp ir.cpp)

src/library/compiler/compiler.cpp

@@ -19,6 +19,7 @@ Author: Leonardo de Moura
 #include "library/compiler/specialize.h"
 #include "library/compiler/lambda_lifting.h"
 #include "library/compiler/extract_closed.h"
+#include "library/compiler/reduce_arity.h"
 #include "library/compiler/ll_infer_type.h"
 #include "library/compiler/simp_app_args.h"
 #include "library/compiler/llnf.h"
@@ -144,6 +145,8 @@ environment compile(environment const & env, options const & opts, names const &
     trace_compiler(name({"compiler", "elim_dead_let"}), ds);
     ds = apply(ecse, new_env, ds);
     trace_compiler(name({"compiler", "cse"}), ds);
+    ds = reduce_arity(ds);
+    trace_compiler(name({"compiler", "reduce_arity"}), ds);
     ds = lambda_lifting(new_env, ds);
     trace_compiler(name({"compiler", "lambda_lifting"}), ds);
     ds = apply(esimp, new_env, ds);
@@ -190,6 +193,7 @@ void initialize_compiler() {
     register_trace_class({"compiler", "erase_irrelevant"});
     register_trace_class({"compiler", "lambda_lifting"});
     register_trace_class({"compiler", "extract_closed"});
+    register_trace_class({"compiler", "reduce_arity"});
     register_trace_class({"compiler", "simp_app_args"});
     register_trace_class({"compiler", "llnf"});
     register_trace_class({"compiler", "boxed"});

src/library/compiler/reduce_arity.cpp

@@ -0,0 +1,50 @@
+/*
+Copyright (c) 2018 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author: Leonardo de Moura
+*/
+#include "kernel/instantiate.h"
+#include "library/compiler/util.h"
+
+namespace lean {
+comp_decls reduce_arity(comp_decl const & cdecl) {
+    expr code    = cdecl.snd();
+    buffer<expr> fvars;
+    name_generator ngen;
+    local_ctx lctx;
+    while (is_lambda(code)) {
+        lean_assert(!has_loose_bvars(code));
+        expr fvar = lctx.mk_local_decl(ngen, binding_name(code), binding_domain(code));
+        fvars.push_back(fvar);
+        code = binding_body(code);
+    }
+    code = instantiate_rev(code, fvars.size(), fvars.data());
+    buffer<expr> new_fvars;
+    for (expr & fvar : fvars) {
+        if (has_fvar(code, fvar)) {
+            new_fvars.push_back(fvar);
+        }
+    }
+    if (fvars.size() == new_fvars.size()) {
+        /* Do nothing, all arguments are used. */
+        return comp_decls(cdecl);
+    }
+    name red_fn   = name(cdecl.fst(), "_rarg");
+    expr red_code = lctx.mk_lambda(new_fvars, code);
+    comp_decl red_decl(red_fn, red_code);
+    /* Replace `cdecl` code with a call to `red_fn`.
+       We rely on inlining to reduce calls to `cdecl` into calls to `red_decl`. */
+    expr new_code = mk_app(mk_constant(red_fn), new_fvars);
+    new_code      = lctx.mk_lambda(fvars, new_code);
+    comp_decl new_decl(cdecl.fst(), new_code);
+    return comp_decls(red_decl, comp_decls(new_decl));
+}
+
+comp_decls reduce_arity(comp_decls const & ds) {
+    comp_decls r;
+    for (comp_decl const & d : ds) {
+        r = append(r, reduce_arity(d));
+    }
+    return r;
+}
+}

src/library/compiler/reduce_arity.h

@@ -0,0 +1,10 @@
+/*
+Copyright (c) 2018 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author: Leonardo de Moura
+*/
+#pragma once
+#include "library/compiler/util.h"
+namespace lean {
+comp_decls reduce_arity(comp_decls const & cdecls);
+}