compiler: add support for Python modules.

Fixes #408.
m-labs · Jun 21, 2016 · f2ae24d · f2ae24d
1 parent 5c54a6a
commit f2ae24d
Showing 4 changed files with 75 additions and 10 deletions.
diff --git a/artiq/compiler/embedding.py b/artiq/compiler/embedding.py
@@ -47,9 +47,20 @@ def __init__(self):
         self.object_current_key = 0
         self.object_forward_map = {}
         self.object_reverse_map = {}
+        self.module_map = {}
         self.type_map = {}
         self.function_map = {}
 
+    # Modules
+    def store_module(self, module, module_type):
+        self.module_map[module] = module_type
+
+    def retrieve_module(self, module):
+        return self.module_map[module]
+
+    def has_module(self, module):
+        return module in self.module_map
+
     # Types
     def store_type(self, host_type, instance_type, constructor_type):
         self.type_map[host_type] = (instance_type, constructor_type)
@@ -88,7 +99,8 @@ def iter_objects(self):
         for obj_id in self.object_forward_map.keys():
             obj_ref = self.object_forward_map[obj_id]
             if isinstance(obj_ref, (pytypes.FunctionType, pytypes.MethodType,
-                                    pytypes.BuiltinFunctionType, SpecializedFunction)):
+                                    pytypes.BuiltinFunctionType, pytypes.ModuleType,
+                                    SpecializedFunction)):
                 continue
             elif isinstance(obj_ref, type):
                 _, obj_typ = self.type_map[obj_ref]
@@ -178,6 +190,21 @@ def quote(self, value):
                 unquote_loc = self._add('`')
                 loc         = quote_loc.join(unquote_loc)
                 return asttyped.QuoteT(value=value, type=function_type, loc=loc)
+        elif isinstance(value, pytypes.ModuleType):
+            if self.embedding_map.has_module(value):
+                module_type = self.embedding_map.retrieve_module(value)
+            else:
+                module_type = types.TModule(value.__name__, OrderedDict())
+                module_type.attributes['__objectid__'] = builtins.TInt32()
+                self.embedding_map.store_module(value, module_type)
+
+            quote_loc   = self._add('`')
+            repr_loc    = self._add(repr(value))
+            unquote_loc = self._add('`')
+            loc         = quote_loc.join(unquote_loc)
+
+            self.value_map[module_type].append((value, loc))
+            return asttyped.QuoteT(value=value, type=module_type, loc=loc)
         else:
             quote_loc   = self._add('`')
             repr_loc    = self._add(repr(value))
@@ -409,7 +436,7 @@ def __init__(self, engine, value_map, quote):
         self.quote = quote
         self.attr_type_cache = {}
 
-    def _compute_value_type(self, object_value, object_type, object_loc, attr_name, loc):
+    def _compute_attr_type(self, object_value, object_type, object_loc, attr_name, loc):
         if not hasattr(object_value, attr_name):
             if attr_name.startswith('_'):
                 names = set(filter(lambda name: not name.startswith('_'),
@@ -519,7 +546,7 @@ def _unify_attribute(self, result_type, value_node, attr_name, attr_loc, loc):
                 attributes, attr_value_type = self.attr_type_cache[attr_type_key]
             except KeyError:
                 attributes, attr_value_type = \
-                    self._compute_value_type(object_value, object_type, object_loc, attr_name, loc)
+                    self._compute_attr_type(object_value, object_type, object_loc, attr_name, loc)
                 self.attr_type_cache[attr_type_key] = attributes, attr_value_type
 
             if attr_name not in attributes:

diff --git a/artiq/compiler/transforms/llvm_ir_generator.py b/artiq/compiler/transforms/llvm_ir_generator.py
@@ -1345,15 +1345,9 @@ def _quote(self, value, typ, path):
         value_id = id(value)
         if value_id in self.llobject_map:
             return self.llobject_map[value_id]
-
         llty = self.llty_of_type(typ)
-        if types.is_constructor(typ) or types.is_instance(typ):
-            if types.is_instance(typ):
-                # Make sure the class functions are quoted, as this has the side effect of
-                # initializing the global closures.
-                self._quote(type(value), typ.constructor,
-                            lambda: path() + ['__class__'])
 
+        def _quote_attributes():
             llglobal = None
             llfields = []
             for attr in typ.attributes:
@@ -1386,6 +1380,16 @@ def _quote(self, value, typ, path):
             llglobal.initializer = ll.Constant(llty.pointee, llfields)
             llglobal.linkage = "private"
             return llglobal
+
+        if types.is_constructor(typ) or types.is_instance(typ):
+            if types.is_instance(typ):
+                # Make sure the class functions are quoted, as this has the side effect of
+                # initializing the global closures.
+                self._quote(type(value), typ.constructor,
+                            lambda: path() + ['__class__'])
+            return _quote_attributes()
+        elif types.is_module(typ):
+            return _quote_attributes()
         elif builtins.is_none(typ):
             assert value is None
             return ll.Constant.literal_struct([])

diff --git a/artiq/compiler/types.py b/artiq/compiler/types.py
@@ -425,6 +425,21 @@ def __repr__(self):
         return "artiq.compiler.types.TInstance({}, {})".format(
                     repr(self.name), repr(self.attributes))
 
+class TModule(TMono):
+    """
+    A type of a module.
+    """
+
+    def __init__(self, name, attributes):
+        assert isinstance(attributes, OrderedDict)
+        super().__init__(name)
+        self.attributes = attributes
+        self.constant_attributes = set()
+
+    def __repr__(self):
+        return "artiq.compiler.types.TModule({}, {})".format(
+                    repr(self.name), repr(self.attributes))
+
 class TMethod(TMono):
     """
     A type of a method.
@@ -608,6 +623,14 @@ def is_instance(typ, name=None):
     else:
         return isinstance(typ, TInstance)
 
+def is_module(typ, name=None):
+    typ = typ.find()
+    if name is not None:
+        return isinstance(typ, TModule) and \
+            typ.name == name
+    else:
+        return isinstance(typ, TModule)
+
 def is_method(typ):
     return isinstance(typ.find(), TMethod)
 

diff --git a/artiq/test/lit/embedding/module.py b/artiq/test/lit/embedding/module.py
@@ -0,0 +1,11 @@
+# RUN: %python -m artiq.compiler.testbench.embedding %s
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+import time, os
+
+@kernel
+def entrypoint():
+    time.sleep(10)
+    os.mkdir("foo")