From 19e111c2ff063558a4a8cf9da2e47d29fe027202 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 13 Feb 2019 16:29:10 -0800
Subject: [PATCH] feat(library/compiler): allow `main` function to also have
 type `io uint32`

---
 src/library/compiler/compiler.cpp | 29 ++++++++++++++++++-----------
 src/library/compiler/emit_cpp.cpp | 21 +++++++++++++--------
 tests/compiler/expr.lean          |  2 +-
 3 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/src/library/compiler/compiler.cpp b/src/library/compiler/compiler.cpp
index 6ddefebbe5..98b0ecde69 100644
--- a/src/library/compiler/compiler.cpp
+++ b/src/library/compiler/compiler.cpp
@@ -109,18 +109,25 @@ bool is_main_fn(environment const & env, name const & n) {
     return false;
 }
 
-/* Return true iff type is `list string -> io uint32` */
+/* Return true iff type is `list string -> io uint32` or `io uint32` */
 bool is_main_fn_type(expr const & type) {
-    if (!is_arrow(type)) return false;
-    expr d = binding_domain(type);
-    expr r = binding_body(type);
-    return
-        is_app(r) &&
-        is_constant(app_fn(r), get_io_name()) &&
-        is_constant(app_arg(r), get_uint32_name()) &&
-        is_app(d) &&
-        is_constant(app_fn(d), get_list_name()) &&
-        is_constant(app_arg(d), get_string_name());
+    if (is_arrow(type)) {
+        expr d = binding_domain(type);
+        expr r = binding_body(type);
+        return
+            is_app(r) &&
+            is_constant(app_fn(r), get_io_name()) &&
+            is_constant(app_arg(r), get_uint32_name()) &&
+            is_app(d) &&
+            is_constant(app_fn(d), get_list_name()) &&
+            is_constant(app_arg(d), get_string_name());
+    } else if (is_app(type)) {
+        return
+            is_constant(app_fn(type), get_io_name()) &&
+            is_constant(app_arg(type), get_uint32_name());
+    } else {
+        return false;
+    }
 }
 
 environment compile(environment const & env, options const & opts, names const & cs) {
diff --git a/src/library/compiler/emit_cpp.cpp b/src/library/compiler/emit_cpp.cpp
index f26e4e8041..e3b9776b41 100644
--- a/src/library/compiler/emit_cpp.cpp
+++ b/src/library/compiler/emit_cpp.cpp
@@ -916,7 +916,8 @@ static bool uses_lean_namespace(environment const & env) {
 }
 
 static void emit_main_fn(std::ostream & out, environment const & env, module_name const & m, comp_decl const & d) {
-    if (get_num_nested_lambdas(d.snd()) != 2) {
+    unsigned arity = get_num_nested_lambdas(d.snd());
+    if (arity != 2 && arity != 1) {
         throw exception("invalid main function, incorrect arity when generating code");
     }
     bool uses_lean_api = uses_lean_namespace(env);
@@ -928,13 +929,17 @@ static void emit_main_fn(std::ostream & out, environment const & env, module_nam
     else
         out << "lean::initialize_runtime_module();\n";
     out << "initialize_" << mangle(m, false) << "();\n";
-    out << "obj* in = lean::box(0);\n";
-    out << "int i = argc;\n";
-    out << "while (i > 1) {\n i--;\n";
-    out << " obj* n = lean::alloc_cnstr(1,2,0); lean::cnstr_set(n, 0, lean::mk_string(argv[i])); lean::cnstr_set(n, 1, in);\n";
-    out << " in = n;\n";
-    out << "}\n";
-    out << "obj * r = " << g_lean_main << "(in, lean::box(0));\n";
+    if (arity == 2) {
+        out << "obj* in = lean::box(0);\n";
+        out << "int i = argc;\n";
+        out << "while (i > 1) {\n i--;\n";
+        out << " obj* n = lean::alloc_cnstr(1,2,0); lean::cnstr_set(n, 0, lean::mk_string(argv[i])); lean::cnstr_set(n, 1, in);\n";
+        out << " in = n;\n";
+        out << "}\n";
+        out << "obj * r = " << g_lean_main << "(in, lean::box(0));\n";
+    } else {
+        out << "obj * r = " << g_lean_main << "(lean::box(0));\n";
+    }
     out << "int ret = lean::unbox(lean::cnstr_get(r, 0));\n";
     out << "lean::dec(r);\n";
     out << "return ret;\n";
diff --git a/tests/compiler/expr.lean b/tests/compiler/expr.lean
index 8f20c70ab7..7319159a0e 100644
--- a/tests/compiler/expr.lean
+++ b/tests/compiler/expr.lean
@@ -1,7 +1,7 @@
 import init.lean.expr
 open lean
 
-def main (xs : list string) : io uint32 :=
+def main : io uint32 :=
 let e := expr.app (expr.const `f []) (expr.const `a []) in
 io.println' e.dbg_to_string *>
 io.println' ("hash: " ++ to_string e.hash) *>