From 917fd94f941324f375105a90454b67805b3950a9 Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident@gmail.com>
Date: Wed, 10 Jun 2020 19:17:29 -0700
Subject: [PATCH] Add limited support for function arguments.

---
 pytest/Compiler/structure.py       |  7 +++++++
 python/npcomp/compiler/frontend.py | 14 +++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/pytest/Compiler/structure.py b/pytest/Compiler/structure.py
index 14b4e8307..f09cb2ea3 100644
--- a/pytest/Compiler/structure.py
+++ b/pytest/Compiler/structure.py
@@ -10,6 +10,13 @@ def import_global(f):
   print(fe.ir_module.to_asm())
   return f
 
+# CHECK-LABEL: func @positional_args
+# CHECK-SAME: (%arg0: !basicpy.UnknownType, %arg1: !basicpy.UnknownType) -> !basicpy.UnknownType
+@import_global
+def positional_args(a, b):
+  # CHECK: basicpy.binary_expr %arg0 "Add" %arg1
+  return a + b
+
 # CHECK-LABEL: func @pass_no_return
 @import_global
 def pass_no_return():
diff --git a/python/npcomp/compiler/frontend.py b/python/npcomp/compiler/frontend.py
index 3e1d18dd9..0447baf59 100644
--- a/python/npcomp/compiler/frontend.py
+++ b/python/npcomp/compiler/frontend.py
@@ -79,17 +79,25 @@ class ImportFrontend:
     logging.debug(":::::::")
     logging.debug("::: Importing global function {}:\n{}", ast_fd.name,
                   ast.dump(ast_fd, include_attributes=True))
+
+    # TODO: VERY BAD: Assumes all positional params.
+    unknown_type = h.basicpy_UnknownType
+    f_params = inspect.signature(f).parameters
+    arg_count = len(f_params)
+    ir_f_type = h.function_type([unknown_type for _ in range(arg_count)],
+                                [unknown_type])
+
     h.builder.set_file_line_col(filename_ident, ast_fd.lineno,
                                 ast_fd.col_offset)
     h.builder.insert_before_terminator(ir_m.first_block)
-    ir_f_type = h.function_type([], [h.basicpy_UnknownType])
     ir_f = h.func_op(ast_fd.name, ir_f_type, create_entry_block=True)
-
     fctx = FunctionContext(ir_c=ir_c,
                            ir_f=ir_f,
                            ir_h=h,
                            filename_ident=filename_ident)
+    for f_arg, ir_arg in zip(f_params, ir_f.first_block.args):
+      fctx.map_local_name(f_arg, ir_arg)
+
     fdimport = FunctionDefImporter(fctx, ast_fd)
     fdimport.import_body()
     return ir_f
-