From 7f77632f1a15294c6ad4522b5939d71f242cea71 Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Sat, 6 Jun 2015 16:58:23 +0300
Subject: [PATCH] Add lit-based tests for type inferencer.

---
 .gitignore                       |  1 +
 artiq/py2llvm/typing.py          | 29 ++++++++++++++++++++++-------
 lit-test/harness.py              | 23 +++++++++++++++++++++++
 lit-test/lit.cfg                 | 12 ++++++++++++
 lit-test/py2llvm/typing/unify.py | 26 ++++++++++++++++++++++++++
 setup.py                         |  2 +-
 6 files changed, 85 insertions(+), 8 deletions(-)
 create mode 100644 lit-test/harness.py
 create mode 100644 lit-test/lit.cfg
 create mode 100644 lit-test/py2llvm/typing/unify.py

diff --git a/.gitignore b/.gitignore
index a8b35b740..22d473a3f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,3 +13,4 @@ doc/manual/_build
 /*.egg-info
 /.coverage
 examples/master/results
+Output/
diff --git a/artiq/py2llvm/typing.py b/artiq/py2llvm/typing.py
index e93af3225..310819530 100644
--- a/artiq/py2llvm/typing.py
+++ b/artiq/py2llvm/typing.py
@@ -129,10 +129,10 @@ class LocalExtractor(algorithm.Visitor):
 class Inferencer(algorithm.Transformer):
     def __init__(self, engine):
         self.engine = engine
-        self.env_stack = [{}]
+        self.env_stack = []
         self.function  = None # currently visited function
 
-    def _unify(self, typea, typeb, loca, locb, kind):
+    def _unify(self, typea, typeb, loca, locb, kind='generic'):
         try:
             typea.unify(typeb)
         except types.UnificationError as e:
@@ -186,6 +186,13 @@ class Inferencer(algorithm.Transformer):
             "name '{name}' is not bound to anything", {"name":name}, loc)
         self.engine.process(diag)
 
+    def visit_root(self, node):
+        extractor = LocalExtractor(env_stack=self.env_stack, engine=self.engine)
+        extractor.visit(node)
+        self.env_stack.append(extractor.typing_env)
+
+        return self.visit(node)
+
     # Visitors that replace node with a typed node
     #
     def visit_arg(self, node):
@@ -349,17 +356,25 @@ class Printer(algorithm.Visitor):
         super().generic_visit(node)
 
 def main():
-    import sys, fileinput
+    import sys, fileinput, os
+
+    inference_mode = True
+
     engine = diagnostic.Engine(all_errors_are_fatal=True)
     try:
-        buf = source.Buffer("".join(fileinput.input()), fileinput.filename())
-        parsed = parse_buffer(buf, engine=engine)
-        typed = Inferencer(engine=engine).visit(parsed)
+        buf = source.Buffer("".join(fileinput.input()), os.path.basename(fileinput.filename()))
+        parsed, comments = parse_buffer(buf, engine=engine)
+        typed = Inferencer(engine=engine).visit_root(parsed)
         printer = Printer(buf)
         printer.visit(typed)
+        for comment in comments:
+            if comment.text.find("CHECK") >= 0:
+                printer.rewriter.remove(comment.loc)
         print(printer.rewrite().source)
     except diagnostic.Error as e:
-        print("\n".join(e.diagnostic.render()), file=sys.stderr)
+        if inference_mode:
+            print("\n".join(e.diagnostic.render()), file=sys.stderr)
+            exit(1)
 
 if __name__ == "__main__":
     main()
diff --git a/lit-test/harness.py b/lit-test/harness.py
new file mode 100644
index 000000000..1779e3266
--- /dev/null
+++ b/lit-test/harness.py
@@ -0,0 +1,23 @@
+"""
+The purpose of this harness is to emulate the behavior of
+the python executable, but add the ARTIQ root to sys.path
+beforehand.
+
+This is necessary because eggs override the PYTHONPATH environment
+variable, but not current directory; therefore `python -m artiq...`
+ran from the ARTIQ root would work, but there is no simple way to
+emulate the same behavior when invoked under lit.
+"""
+
+import sys, os, argparse, importlib
+
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument('-m', metavar='mod', type=str, help='run library module as a script')
+parser.add_argument('args', type=str, nargs='+', help='arguments passed to program in sys.argv[1:]')
+args = parser.parse_args(sys.argv[1:])
+
+artiq_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+sys.path.insert(1, artiq_path)
+
+sys.argv[1:] = args.args
+importlib.import_module(args.m).main()
diff --git a/lit-test/lit.cfg b/lit-test/lit.cfg
new file mode 100644
index 000000000..398164291
--- /dev/null
+++ b/lit-test/lit.cfg
@@ -0,0 +1,12 @@
+import lit.util
+import lit.formats
+
+config.name = 'ARTIQ'
+config.test_format = lit.formats.ShTest()
+config.suffixes = ['.py']
+config.excludes = ['harness.py']
+config.test_source_root = os.path.dirname(__file__)
+
+python_executable = 'python3'
+harness = '{} {}'.format(python_executable, os.path.join(config.test_source_root, 'harness.py'))
+config.substitutions.append( ('%python', harness) )
diff --git a/lit-test/py2llvm/typing/unify.py b/lit-test/py2llvm/typing/unify.py
new file mode 100644
index 000000000..879c55d0f
--- /dev/null
+++ b/lit-test/py2llvm/typing/unify.py
@@ -0,0 +1,26 @@
+# RUN: %python -m artiq.py2llvm.typing %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+a = 1
+# CHECK-L: a:int(width='a)
+
+b = a
+# CHECK-L: b:int(width='a)
+
+c = True
+# CHECK-L: c:bool
+
+d = False
+# CHECK-L: d:bool
+
+e = None
+# CHECK-L: e:NoneType
+
+f = 1.0
+# CHECK-L: f:float
+
+g = []
+# CHECK-L: g:list(elt='b)
+
+h = [1]
+# CHECK-L: h:list(elt=int(width='c))
diff --git a/setup.py b/setup.py
index 44087dc4c..427b9e9aa 100755
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,7 @@ from setuptools import setup, find_packages
 requirements = [
     "sphinx", "sphinx-argparse", "pyserial", "numpy", "scipy",
     "python-dateutil", "prettytable", "h5py", "pydaqmx", "pyelftools",
-    "quamash", "pyqtgraph", "pythonparser"
+    "quamash", "pyqtgraph", "pythonparser", "lit", "OutputCheck"
 ]
 
 scripts = [