[DSLX] Improve type handling in parametrics

xls/dslx/frontend/parser.cc

@@ -2779,10 +2779,81 @@ absl::StatusOr<std::vector<ExprOrType>> Parser::ParseParametrics(
     }
 
     if (peek->kind() == TokenKind::kIdentifier) {
+      // We have to distinguish between:
+      // type identifiers and
+      // value identifiers
       XLS_ASSIGN_OR_RETURN(auto nocr, ParseNameOrColonRef(bindings));
-      return ToExprNode(nocr);
+      XLS_ASSIGN_OR_RETURN(bool complex_,
+          PeekTokenIn({TokenKind::kOAngle, TokenKind::kOBrack}));
+      if (!complex_) {
+        return ToExprNode(nocr);
+      }
+
+      // Complex parametrics should only be types:
+      // - imports can't be parametric.
+      // - existing code didn't allow for parametric enums
+      TypeRef* type_ref;
+      if (std::holds_alternative<ColonRef*>(nocr)) {
+        auto cr = std::get<ColonRef*>(nocr);
+        if (!cr->ResolveImportSubject().has_value()) {
+          XLS_ASSIGN_OR_RETURN(
+              BoundNode bn,
+              bindings.ResolveNodeOrError(
+                  ToAstNode(TypeDefinitionGetNameDef(cr))->ToString(),
+                  cr->span()));
+          return ParseErrorStatus(
+              cr->span(),
+              absl::StrFormat("Expected module for module-reference; got %s",
+                              ToAstNode(bn)->ToString()));
+        }
+        type_ref = module_->Make<TypeRef>(cr->span(), cr);
+      } else {
+        auto nr = std::get<NameRef*>(nocr);
+        XLS_ASSIGN_OR_RETURN(
+            BoundNode type_def,
+            bindings.ResolveNodeOrError(nr->ToString(), nr->span()));
+        if (!IsOneOf<TypeAlias, EnumDef, StructDef>(ToAstNode(type_def))) {
+          return ParseErrorStatus(
+              nr->span(),
+              absl::StrFormat(
+                  "Expected a type, but identifier '%s' doesn't resolve to "
+                  "a type, it resolved to a %s",
+                  nr->ToString(), BoundNodeGetTypeString(type_def)));
+        }
+
+        XLS_ASSIGN_OR_RETURN(TypeDefinition type_definition,
+                             BoundNodeToTypeDefinition(type_def));
+        type_ref = module_->Make<TypeRef>(nr->span(), type_definition);
+      }
+
+      std::vector<ExprOrType> parametrics;
+      XLS_ASSIGN_OR_RETURN(bool peek_is_oangle, PeekTokenIs(TokenKind::kOAngle));
+      if (peek_is_oangle) {
+        XLS_ASSIGN_OR_RETURN(parametrics, ParseParametrics(bindings));
+      }
+
+      std::vector<Expr*> dims;
+      XLS_ASSIGN_OR_RETURN(bool peek_is_obrack, PeekTokenIs(TokenKind::kOBrack));
+      if (peek_is_obrack) {  // Array type annotation.
+        XLS_ASSIGN_OR_RETURN(dims, ParseDims(bindings));
+      }
+
+      Span span(ToAstNode(nocr)->GetSpan().value().start(), GetPos());
+      XLS_ASSIGN_OR_RETURN(TypeAnnotation * type_annotation,
+                           MakeTypeRefTypeAnnotation(
+                             span, type_ref, dims,
+                             std::move(parametrics)));
+
+      {
+        XLS_ASSIGN_OR_RETURN(const Token* peek, PeekToken());
+        if (peek->kind() == TokenKind::kColon) {
+          return ParseCast(bindings, type_annotation);
+        }
+      }
+      return type_annotation;
     }
 
+
     XLS_ASSIGN_OR_RETURN(TypeAnnotation * type_annotation,
                          ParseTypeAnnotation(bindings));
 

xls/dslx/frontend/parser_test.cc

@@ -836,10 +836,154 @@ TEST_F(ParserTest, CmpChainParensOnLhsAndRhs) {
                 "((x) == (y)) == ((y) == (z))");
 }
 
-TEST_F(ParserTest, ZeroMacro) {
+TEST_F(ParserTest, ZeroMacroSimple) {
   RoundTripExpr("zero!<u32>()", {}, /*populate_dslx_builtins=*/true);
 }
 
+TEST_F(ParserTest, ZeroMacroSimpleStruct) {
+  RoundTripExpr("zero!<MyType>()", {"MyType"}, /*populate_dslx_builtins=*/true);
+}
+
+TEST_F(ParserTest, ZeroMacroSimpleArray) {
+  RoundTripExpr("zero!<u32[10]>()", {}, /*populate_dslx_builtins=*/true);
+}
+
+TEST_F(ParserTest, ZeroMacroSimpleBitsArray) {
+  RoundTripExpr("zero!<bits[32][10]>()", {}, /*populate_dslx_builtins=*/true);
+}
+
+TEST_F(ParserTest, ZeroMacroSimpleStructArray) {
+  const char* text = R"(zero!<MyType[10]>())";
+  Scanner s{"test.x", std::string{text}};
+  Parser p{"test", &s};
+
+  Bindings b;
+
+  Module& mod = p.module();
+  for (auto const& it : GetParametricBuiltins()) {
+    std::string name(it.first);
+    b.Add(name, mod.GetOrCreateBuiltinNameDef(name));
+  }
+
+  NameDef* name_def = mod.Make<NameDef>(
+    Span::Fake(), std::string("MyType"), nullptr);
+
+  using StructMember = std::pair<NameDef*, TypeAnnotation*>;
+  auto* struct_def = mod.Make<StructDef>(
+    Span::Fake(), name_def, std::vector<ParametricBinding*>(),
+    std::vector<StructMember>(), false);
+  b.Add(name_def->identifier(), struct_def);
+
+  auto expr_or = p.ParseExpression(/*bindings=*/b);
+  if (!expr_or.ok()) {
+    TryPrintError(expr_or.status(),
+                  [&](std::string_view path) -> absl::StatusOr<std::string> {
+                    return std::string{text};
+                  });
+  }
+  ASSERT_TRUE(expr_or.ok());
+}
+
+TEST_F(ParserTest, ZeroMacroParametricStruct) {
+  const char* text = R"(zero!<MyType<MyParm0, MyParm1>>())";
+  Scanner s{"test.x", std::string{text}};
+  Parser p{"test", &s};
+
+  Bindings b;
+
+  Module& mod = p.module();
+  for (auto const& it : GetParametricBuiltins()) {
+    std::string name(it.first);
+    b.Add(name, mod.GetOrCreateBuiltinNameDef(name));
+  }
+
+  std::vector<ParametricBinding*> params;
+
+  NameDef* name_def = mod.Make<NameDef>(
+    Span::Fake(), std::string("MyParm0"), nullptr);
+  b.Add(name_def->identifier(), name_def);
+  BuiltinType builtin_type = BuiltinTypeFromString("u32").value();
+  TypeAnnotation* elem_type = mod.Make<BuiltinTypeAnnotation>(
+    Span::Fake(), builtin_type, mod.GetOrCreateBuiltinNameDef("u32"));
+  params.push_back(mod.Make<ParametricBinding>(name_def, elem_type, nullptr));
+
+  name_def = mod.Make<NameDef>(
+    Span::Fake(), std::string("MyParm1"), nullptr);
+  b.Add(name_def->identifier(), name_def);
+  builtin_type = BuiltinTypeFromString("u32").value();
+  elem_type = mod.Make<BuiltinTypeAnnotation>(
+    Span::Fake(), builtin_type, mod.GetOrCreateBuiltinNameDef("u32"));
+  params.push_back(mod.Make<ParametricBinding>(name_def, elem_type, nullptr));
+
+  name_def = mod.Make<NameDef>(
+    Span::Fake(), std::string("MyType"), nullptr);
+
+  using StructMember = std::pair<NameDef*, TypeAnnotation*>;
+  auto* struct_def = mod.Make<StructDef>(
+    Span::Fake(), name_def, params,
+    std::vector<StructMember>(), false);
+  b.Add(name_def->identifier(), struct_def);
+
+  auto expr_or = p.ParseExpression(/*bindings=*/b);
+  if (!expr_or.ok()) {
+    TryPrintError(expr_or.status(),
+                  [&](std::string_view path) -> absl::StatusOr<std::string> {
+                    return std::string{text};
+                  });
+  }
+  ASSERT_TRUE(expr_or.ok());
+}
+
+TEST_F(ParserTest, ZeroMacroParametricStructArray) {
+  const char* text = R"(zero!<MyType<MyParm0, MyParm1>[10]>())";
+  Scanner s{"test.x", std::string{text}};
+  Parser p{"test", &s};
+
+  Bindings b;
+
+  Module& mod = p.module();
+  for (auto const& it : GetParametricBuiltins()) {
+    std::string name(it.first);
+    b.Add(name, mod.GetOrCreateBuiltinNameDef(name));
+  }
+
+  std::vector<ParametricBinding*> params;
+
+  NameDef* name_def = mod.Make<NameDef>(
+    Span::Fake(), std::string("MyParm0"), nullptr);
+  b.Add(name_def->identifier(), name_def);
+  BuiltinType builtin_type = BuiltinTypeFromString("u32").value();
+  TypeAnnotation* elem_type = mod.Make<BuiltinTypeAnnotation>(
+    Span::Fake(), builtin_type, mod.GetOrCreateBuiltinNameDef("u32"));
+  params.push_back(mod.Make<ParametricBinding>(name_def, elem_type, nullptr));
+
+  name_def = mod.Make<NameDef>(
+    Span::Fake(), std::string("MyParm1"), nullptr);
+  b.Add(name_def->identifier(), name_def);
+  builtin_type = BuiltinTypeFromString("u32").value();
+  elem_type = mod.Make<BuiltinTypeAnnotation>(
+    Span::Fake(), builtin_type, mod.GetOrCreateBuiltinNameDef("u32"));
+  params.push_back(mod.Make<ParametricBinding>(name_def, elem_type, nullptr));
+
+  name_def = mod.Make<NameDef>(
+    Span::Fake(), std::string("MyType"), nullptr);
+
+  using StructMember = std::pair<NameDef*, TypeAnnotation*>;
+  auto* struct_def = mod.Make<StructDef>(
+    Span::Fake(), name_def, params,
+    std::vector<StructMember>(), false);
+  b.Add(name_def->identifier(), struct_def);
+
+  auto expr_or = p.ParseExpression(/*bindings=*/b);
+  if (!expr_or.ok()) {
+    TryPrintError(expr_or.status(),
+                  [&](std::string_view path) -> absl::StatusOr<std::string> {
+                    return std::string{text};
+                  });
+  }
+  ASSERT_TRUE(expr_or.ok());
+}
+
 TEST_F(ParserTest, ParseBlockWithTwoStatements) {
   RoundTripExpr(R"({
   type MyU32 = u32;