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)