def collect_literals(fn: FuncIR, literals: Literals) -> None:
"""Store all Python literal object refs in fn.
Collecting literals must happen only after we have the final IR.
This way we won't include literals that have been optimized away.
"""
for block in fn.blocks:
for op in block.ops:
if isinstance(op, LoadLiteral):
literals.record_literal(op.value)
def __init__(
self,
names: NameGenerator,
group_name: Optional[str] = None,
group_map: Optional[Dict[str, Optional[str]]] = None,
) -> None:
"""Setup shared emitter state.
Args:
names: The name generator to use
group_map: Map from module names to group name
group_name: Current group name
"""
self.temp_counter = 0
self.names = names
self.group_name = group_name
self.group_map = group_map or {}
# Groups that this group depends on
self.group_deps = set() # type: Set[str]
# The map below is used for generating declarations and
# definitions at the top of the C file. The main idea is that they can
# be generated at any time during the emit phase.
# A map of a C identifier to whatever the C identifier declares. Currently this is
# used for declaring structs and the key corresponds to the name of the struct.
# The declaration contains the body of the struct.
self.declarations = OrderedDict() # type: Dict[str, HeaderDeclaration]
self.literals = Literals()
def test_tuple_literal(self) -> None:
lit = Literals()
lit.record_literal((1, 'y', None, (b'a', 'b')))
lit.record_literal((b'a', 'b'))
lit.record_literal(())
assert lit.literal_index((b'a', 'b')) == 7
assert lit.literal_index((1, 'y', None, (b'a', 'b'))) == 8
assert lit.literal_index(()) == 9
print(lit.encoded_tuple_values())
assert lit.encoded_tuple_values() == [
'3', # Number of tuples
'2',
'5',
'4', # First tuple (length=2)
'4',
'6',
'3',
'0',
'7', # Second tuple (length=4)
'0', # Third tuple (length=0)
]
def test_simple_literal_index(self) -> None:
lit = Literals()
lit.record_literal(1)
lit.record_literal('y')
lit.record_literal(True)
lit.record_literal(None)
lit.record_literal(False)
assert lit.literal_index(None) == 0
assert lit.literal_index(False) == 1
assert lit.literal_index(True) == 2
assert lit.literal_index('y') == 3
assert lit.literal_index(1) == 4