def initialize(self, mod, mng):
"""Add stub types to the module."""
mod[env_type] = adt.TypeData(env_type, [a], [empty_env, cons_env])
for node in mng.all_nodes:
if isinstance(node.abstract, AbstractTaggedUnion):
for opt in node.abstract.options:
get_union_ctr(*opt)
mod[union_type] = adt.TypeData(union_type, [], list(tag_map.values()))
def finalize(self, mod):
"""Fill in stub type definitions."""
mod[env_val] = adt.TypeData(
env_val, [], list(m['ctr'] for m in self.env_val_map.values()))
for m in self.env_val_map.values():
for k, v in m.items():
if k == 'ctr':
continue
mod[v[0]] = v[1]
def initialize(self, mod, mng):
"""Add types to the module."""
if mng is not None:
for node in mng.all_nodes:
if isinstance(node.abstract, AbstractTaggedUnion):
for opt in node.abstract.options:
get_union_ctr(*opt)
elif node.is_apply(P.env_setitem):
key = node.inputs[2]
tt = to_relay_type(node.inputs[3].abstract)
assert key.is_constant()
self.env_val_map[key.value] = tt
env_val_keys = sorted(list(self.env_val_map.keys()))
for i, k in enumerate(env_val_keys):
self.env_val_map[k] = (i, self.env_val_map[k])
mod[union_type] = adt.TypeData(union_type, [], list(tag_map.values()))
mod[option_type] = adt.TypeData(option_type, [a], [nil, some])
self.env_ctr = adt.Constructor("v", [self._build_env_type()], env_type)
mod[env_type] = adt.TypeData(env_type, [], [self.env_ctr, dead_env])