diff --git a/artiq/compiler/embedding.py b/artiq/compiler/embedding.py index f4a9bd4cb..b2e47510d 100644 --- a/artiq/compiler/embedding.py +++ b/artiq/compiler/embedding.py @@ -47,9 +47,20 @@ class EmbeddingMap: 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 @@ class EmbeddingMap: 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 @@ class ASTSynthesizer: 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 @@ class StitchingInferencer(Inferencer): 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 @@ class StitchingInferencer(Inferencer): 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 index 9f1eb25d9..b1d464dca 100644 --- a/artiq/compiler/transforms/llvm_ir_generator.py +++ b/artiq/compiler/transforms/llvm_ir_generator.py @@ -1345,15 +1345,9 @@ class LLVMIRGenerator: 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 @@ class LLVMIRGenerator: 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 index cca1d15a9..ac1e9377c 100644 --- a/artiq/compiler/types.py +++ b/artiq/compiler/types.py @@ -425,6 +425,21 @@ class TInstance(TMono): 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 new file mode 100644 index 000000000..d25475b6d --- /dev/null +++ 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")