def test_visit_index(self):
fake_span = self.fake_span
# Make a t[i] inde xnode.
t = ast.NameRef(fake_span, 't', ast.NameDef(fake_span, 't'))
i = ast.NameRef(fake_span, 'i', ast.NameDef(fake_span, 'i'))
index = ast.Index(fake_span, t, i)
c = _Collector()
index.accept(c)
self.assertEqual(c.collected, [t, i, index])
def test_stringify_type(self):
fake_span = self.fake_span
m = self.m
number_5 = ast.Number(m, fake_span, '5')
number_2 = ast.Number(m, fake_span, '2')
number_3 = ast.Number(m, fake_span, '3')
bits_token = Token(TokenKind.KEYWORD, value=Keyword.BITS, span=fake_span)
type_ = ast_helpers.make_builtin_type_annotation(m, fake_span, bits_token,
(number_5,))
self.assertEqual('bits[5]', str(type_))
type_ = ast_helpers.make_builtin_type_annotation(
m, fake_span,
Token(TokenKind.KEYWORD, value=Keyword.U32,
span=fake_span), (number_5,))
self.assertEqual('u32[5]', str(type_))
# "no-volume" bits array.
# TODO(leary): 2020-08-24 delete bits in favor of uN
type_ = ast_helpers.make_builtin_type_annotation(m, fake_span, bits_token,
())
self.assertEqual('bits', str(type_))
# TypeRef with dims.
my_type_tok = Token(TokenKind.IDENTIFIER, value='MyType', span=fake_span)
name_def = ast.NameDef(m, fake_span, 'MyType')
type_def = ast.TypeDef(m, False, name_def, type_)
type_ref = ast.TypeRef(m, fake_span, my_type_tok.value, type_def=type_def)
type_ = ast_helpers.make_type_ref_type_annotation(m, fake_span, type_ref,
(number_2, number_3))
self.assertEqual('MyType[2][3]', str(type_))
def test_type_annotation_properties(self):
m = self.m
fake_span = self.fake_span
number_5 = ast.Number(m, fake_span, '5')
number_2 = ast.Number(m, fake_span, '2')
number_3 = ast.Number(m, fake_span, '3')
bits_token = Token(TokenKind.KEYWORD, value=Keyword.BITS, span=fake_span)
un_token = Token(TokenKind.KEYWORD, value=Keyword.UN, span=fake_span)
u32_token = Token(TokenKind.KEYWORD, value=Keyword.U32, span=fake_span)
type_ = ast_helpers.make_builtin_type_annotation(m, fake_span, bits_token,
(number_5,))
self.assertEqual('bits[5]', str(type_))
type_ = ast_helpers.make_builtin_type_annotation(m, fake_span, bits_token,
(number_5, number_2))
self.assertIsInstance(type_, ast.ArrayTypeAnnotation)
self.assertEqual('bits[5][2]', str(type_))
type_ = ast_helpers.make_builtin_type_annotation(m, fake_span, u32_token,
())
self.assertEqual('u32', str(type_))
type_ = ast_helpers.make_builtin_type_annotation(m, fake_span, u32_token,
(number_3,))
self.assertIsInstance(type_, ast.ArrayTypeAnnotation)
self.assertEqual('u32[3]', str(type_))
type_ = ast_helpers.make_builtin_type_annotation(m, fake_span, un_token,
(number_2,))
self.assertEqual('uN[2]', str(type_))
type_ = ast_helpers.make_builtin_type_annotation(m, fake_span, un_token,
(number_2, number_3))
self.assertIsInstance(type_, ast.ArrayTypeAnnotation)
self.assertEqual('uN[2][3]', str(type_))
# TODO(leary): 2020-08-24 delete bits in favor of uN
# "no-volume" bits array.
type_ = ast_helpers.make_builtin_type_annotation(m, fake_span, bits_token,
())
self.assertEqual('bits', str(type_))
# TypeRef with dims.
name_def = ast.NameDef(m, fake_span, 'MyType')
type_def = ast.TypeDef(m, False, name_def, type_)
type_ref = ast.TypeRef(m, fake_span, 'MyType', type_def=type_def)
type_ = ast_helpers.make_type_ref_type_annotation(m, fake_span, type_ref,
(number_2, number_3))
self.assertIsInstance(type_, ast.ArrayTypeAnnotation)
self.assertEqual('MyType[2][3]', str(type_))
type_ = ast.TupleTypeAnnotation(
m, fake_span,
(ast_helpers.make_builtin_type_annotation(m, fake_span, bits_token,
(number_5,)),
ast_helpers.make_builtin_type_annotation(m, fake_span, bits_token,
(number_2,))))
self.assertIsInstance(type_, ast.TupleTypeAnnotation)
self.assertEqual('(bits[5], bits[2])', str(type_))
def test_visit_unop(self):
fake_span = self.fake_span
i_def = ast.NameDef(fake_span, 'i')
i_ref = ast.NameRef(fake_span, 'i', i_def)
negated = ast.Unop(Token(TokenKind.MINUS, fake_span), i_ref)
c = _Collector()
negated.accept(c)
self.assertEqual(c.collected, [i_ref, negated])
def _generate_expr(self, env: Env,
level: int) -> Tuple[ast.Expr, ast.TypeAnnotation]:
"""Generates an expression AST node and returns it."""
if self.should_nest(level):
identifier = self.gensym()
name_def = ast.NameDef(self.m, self.fake_span, identifier)
choices = collections.OrderedDict()
if self.options.emit_loops:
choices[self._generate_counted_for] = 1.0
choices[self._generate_tuple_or_index] = 1.0
choices[self._generate_concat] = 1.0
choices[self._generate_binop] = 1.0
choices[self._generate_unop] = 1.0
choices[self._generate_unop_builtin] = 1.0
choices[self._generate_one_hot_select_builtin] = 1.0
choices[self._generate_number] = 1.0
choices[self._generate_bit_slice] = 1.0
if not self._codegen_ops_only:
choices[self._generate_bitwise_reduction] = 1.0
choices[self._generate_cast_bits_to_array] = 1.0
# If maps recurse with equal probability, then the output will grow
# exponentially with level, so we need to scale inversely.
choices[lambda env: self._generate_map(level, env)] = 1.0 / (10**level)
rng_choices = getattr(self.rng, 'choices')
while True:
expr_generator = rng_choices(
population=list(choices.keys()),
weights=list(choices.values()),
k=1)[0]
try:
rhs, rhs_type = expr_generator(env)
except EmptyEnvError:
# Try a different generator that may be more accepting of env values.
del choices[expr_generator]
if not choices: # If we ran out of things to try, bail.
return self._generate_retval(env)
continue
else:
break
new_env = collections.OrderedDict(env)
new_env[identifier] = (
lambda: self._make_name_ref(name_def)), rhs_type, False
body, body_type = self._generate_expr(new_env, level + 1)
name_def_tree = ast.NameDefTree(self.m, self.fake_span, name_def)
let = ast.Let(
self.m,
name_def_tree,
rhs_type,
rhs,
body,
self.fake_span,
const=None)
return let, body_type
else: # Should not nest any more -- select return values.
return self._generate_retval(env)
def test_ndt_preorder(self):
fake_pos = self.fake_pos
fake_span = Span(fake_pos, fake_pos)
t = ast.NameDef(fake_span, 't')
u = ast.NameDef(fake_span, 'u')
wrapped_t = ast.NameDefTree(fake_span, t)
wrapped_u = ast.NameDefTree(fake_span, u)
interior = ast.NameDefTree(fake_span, (wrapped_t, wrapped_u))
outer = ast.NameDefTree(fake_span, (interior, ))
walk_data = []
def walk(item: ast.NameDefTree, level: int, i: int):
walk_data.append((item, level, i))
outer.do_preorder(walk)
self.assertLen(walk_data, 3)
self.assertEqual(walk_data[0], (interior, 1, 0))
self.assertEqual(walk_data[1], (wrapped_t, 2, 0))
self.assertEqual(walk_data[2], (wrapped_u, 2, 1))
def generate_function(self,
name: Text,
level: int = 0,
param_count: int = None) -> ast.Function:
if param_count is None:
param_count = int(math.ceil(self.rng.weibullvariate(1, 1.15) * 4))
params = self._generate_params(param_count)
body, body_type = self._generate_body(level, params)
return ast.Function(self.fake_span,
ast.NameDef(self.fake_span, name),
parametric_bindings=(),
params=params,
return_type=body_type,
body=body,
public=False)
def _generate_param(self) -> ast.Param:
identifier = self.gensym()
type_ = self._generate_bits_type()
name_def = ast.NameDef(self.fake_span, identifier)
return ast.Param(name_def, type_)
def _make_name_def(self, identifier: Text) -> ast.NameDef:
return ast.NameDef(self.fake_span, identifier)
