From acb8810e628602b8cf23b6bf704e836d93f4b0fe Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Thu, 23 Jul 2015 03:07:30 +0300
Subject: [PATCH] Add tests for lambdas and functions.

---
 artiq/compiler/module.py                      |  2 +
 .../compiler/transforms/artiq_ir_generator.py | 41 ++++++++++++++-----
 .../compiler/transforms/llvm_ir_generator.py  |  2 +-
 lit-test/compiler/integration/function.py     | 10 +++++
 lit-test/compiler/integration/lambda.py       |  9 ++++
 5 files changed, 52 insertions(+), 12 deletions(-)
 create mode 100644 lit-test/compiler/integration/function.py
 create mode 100644 lit-test/compiler/integration/lambda.py

diff --git a/artiq/compiler/module.py b/artiq/compiler/module.py
index 38bccda92..80e44acd3 100644
--- a/artiq/compiler/module.py
+++ b/artiq/compiler/module.py
@@ -32,6 +32,8 @@ class Module:
         monomorphism_validator.visit(self.typedtree)
         escape_validator.visit(self.typedtree)
         self.artiq_ir = artiq_ir_generator.visit(self.typedtree)
+        print(self.artiq_ir[0])
+        print(self.artiq_ir[1])
         dead_code_eliminator.process(self.artiq_ir)
         local_access_validator.process(self.artiq_ir)
         self.llvm_ir = llvm_ir_generator.process(self.artiq_ir)
diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py
index 1c49508e2..3b9e5be13 100644
--- a/artiq/compiler/transforms/artiq_ir_generator.py
+++ b/artiq/compiler/transforms/artiq_ir_generator.py
@@ -1250,18 +1250,37 @@ class ARTIQIRGenerator(algorithm.Visitor):
         else:
             typ = node.func.type.find()
             func = self.visit(node.func)
-            args = [self.visit(arg) for arg in node.args]
-            for index, optarg_name in enumerate(typ.optargs):
-                if len(typ.args) + index >= len(args):
-                    optarg_typ = ir.TOption(typ.optargs[optarg_name])
-                    for keyword in node.keywords:
-                        if keyword.arg == optarg_name:
-                            value = self.append(ir.Alloc([self.visit(keyword.value)], optarg_typ))
-                            args.append(value)
+            args = [None] * (len(typ.args) + len(typ.optargs))
+
+            for index, arg_node in enumerate(node.args):
+                arg = self.visit(arg_node)
+                if index < len(typ.args):
+                    args[index] = arg
+                else:
+                    args[index] = self.append(ir.Alloc([arg], ir.TOption(arg.type)))
+
+            for keyword in node.keywords:
+                arg = self.visit(keyword.value)
+                if keyword.arg in typ.args:
+                    for index, arg_name in enumerate(typ.args):
+                        if keyword.arg == arg_name:
+                            assert args[index] is None
+                            args[index] = arg
                             break
-                    else:
-                        value = self.append(ir.Alloc([], optarg_typ))
-                        args.append(value)
+                elif keyword.arg in typ.optargs:
+                    for index, optarg_name in enumerate(typ.optargs):
+                        if keyword.arg == optarg_name:
+                            assert args[len(typ.args) + index] is None
+                            args[len(typ.args) + index] = \
+                                    self.append(ir.Alloc([arg], ir.TOption(arg.type)))
+                            break
+
+            for index, optarg_name in enumerate(typ.optargs):
+                if args[len(typ.args) + index] is None:
+                    args[len(typ.args) + index] = \
+                            self.append(ir.Alloc([], ir.TOption(typ.optargs[optarg_name])))
+
+            assert None not in args
 
             if self.unwind_target is None:
                 return self.append(ir.Call(func, args))
diff --git a/artiq/compiler/transforms/llvm_ir_generator.py b/artiq/compiler/transforms/llvm_ir_generator.py
index 5d4075f8b..8749d3348 100644
--- a/artiq/compiler/transforms/llvm_ir_generator.py
+++ b/artiq/compiler/transforms/llvm_ir_generator.py
@@ -519,7 +519,7 @@ class LLVMIRGenerator:
         if builtins.is_none(insn.value().type):
             return self.llbuilder.ret_void()
         else:
-            return self.llbuilder.ret(self.llmap[insn.value()])
+            return self.llbuilder.ret(self.map(insn.value()))
 
     def process_Unreachable(self, insn):
         return self.llbuilder.unreachable()
diff --git a/lit-test/compiler/integration/function.py b/lit-test/compiler/integration/function.py
new file mode 100644
index 000000000..e1ae25fd9
--- /dev/null
+++ b/lit-test/compiler/integration/function.py
@@ -0,0 +1,10 @@
+# RUN: %python -m artiq.compiler.testbench.jit %s
+
+def fib(x):
+    if x == 1:
+        return x
+    else:
+        return x * fib(x - 1)
+assert fib(5) == 120
+
+# argument combinations handled in lambda.py
diff --git a/lit-test/compiler/integration/lambda.py b/lit-test/compiler/integration/lambda.py
new file mode 100644
index 000000000..7bbe8824f
--- /dev/null
+++ b/lit-test/compiler/integration/lambda.py
@@ -0,0 +1,9 @@
+# RUN: %python -m artiq.compiler.testbench.jit %s
+
+assert (lambda: 1)() == 1
+assert (lambda x: x)(1) == 1
+assert (lambda x, y: x + y)(1, 2) == 3
+assert (lambda x, y=1: x + y)(1) == 2
+assert (lambda x, y=1: x + y)(1, 2) == 3
+assert (lambda x, y=1: x + y)(x=3) == 4
+assert (lambda x, y=1: x + y)(y=2, x=3) == 5