Add the basics of generics support behind a "generics" feature flag.

I believe what's left after this is primarily:
- Fixing colon refs through generics. Currently ColonRef resolution is too early for this. Static trait functions also need it moved later.
- Checking if imports work.
- Fixing invocation of impl functions from generified callers in IR converter.

PiperOrigin-RevId: 830639806

Author: richmckeever

xls/data_structures/inline_bitmap.h

@@ -21,6 +21,7 @@
 #include <cstring>
 #include <iterator>
 #include <optional>
+#include <string_view>
 #include <utility>
 
 #include "absl/base/casts.h"
@@ -80,6 +81,19 @@ class InlineBitmap {
     return result;
   }
 
+  // Variant of FromBytes for converting a string.
+  static InlineBitmap FromBytes(std::string_view bytes) {
+    InlineBitmap result(bytes.size() * 8, false);
+    // memcpy() requires valid pointers even when the number of bytes copied is
+    // zero, and an empty absl::Span's data() pointer may not be valid. Guard
+    // the memcpy with a check that the span is not empty.
+    if (!result.data_.empty()) {
+      std::memcpy(result.data_.data(), bytes.data(), bytes.size());
+    }
+    result.MaskLastWord();
+    return result;
+  }
+
   // Constructs a bitmap of width `bits.size()` using the given bits,
   // interpreting index 0 as the *least* significant bit.
   //

xls/dslx/fmt/ast_fmt.cc

@@ -3093,6 +3093,10 @@ absl::StatusOr<DocRef> Formatter::Format(const Module& n) {
           pieces.push_back(arena_.MakeText("#![feature(channel_attributes)]"));
           pieces.push_back(arena_.hard_line());
           break;
+        case ModuleAttribute::kGenerics:
+          pieces.push_back(arena_.MakeText("#![feature(generics)]"));
+          pieces.push_back(arena_.hard_line());
+          break;
       }
     }
     pieces.push_back(arena_.hard_line());

xls/dslx/frontend/ast.cc

@@ -1122,6 +1122,10 @@ TypeVariableTypeAnnotation::TypeVariableTypeAnnotation(
       internal_(internal) {}
 
 std::string TypeVariableTypeAnnotation::ToString() const {
+  if (!internal_) {
+    return type_variable_->ToString();
+  }
+
   return absl::StrCat("TypeVariableTypeAnnotation: ",
                       type_variable_->ToString());
 }

xls/dslx/frontend/module.cc

@@ -94,6 +94,9 @@ std::string Module::ToString() const {
             case ModuleAttribute::kChannelAttributes:
               absl::StrAppend(out, "#![feature(channel_attributes)]");
               break;
+            case ModuleAttribute::kGenerics:
+              absl::StrAppend(out, "#![feature(generics)]");
+              break;
           }
         });
     return absl::StrCat(header, "\n\n", body);

xls/dslx/frontend/module.h

@@ -81,6 +81,8 @@ enum class ModuleAttribute : uint8_t {
 
   // Enable #channel() attributes for this module.
   kChannelAttributes,
+
+  kGenerics,
 };
 
 // Represents a syntactic module in the AST.

xls/dslx/frontend/parser.cc

@@ -310,6 +310,8 @@ absl::Status Parser::ParseModuleAttribute() {
     } else if (feature == "channel_attributes") {
       module_->AddAttribute(ModuleAttribute::kChannelAttributes,
                             attribute_span);
+    } else if (feature == "generics") {
+      module_->AddAttribute(ModuleAttribute::kGenerics, attribute_span);
     } else {
       return ParseErrorStatus(
           attribute_span,
@@ -4114,10 +4116,14 @@ absl::StatusOr<std::vector<ParametricBinding*>> Parser::ParseParametricBindings(
     XLS_RETURN_IF_ERROR(
         DropTokenOrError(TokenKind::kColon, /*start=*/nullptr,
                          "Expect type annotation on parametric"));
-    XLS_ASSIGN_OR_RETURN(TypeAnnotation * type,
-                         ParseTypeAnnotation(bindings,
-                                             /*first=*/std::nullopt,
-                                             /*allow_generic_type=*/true));
+    XLS_ASSIGN_OR_RETURN(
+        TypeAnnotation * type,
+        ParseTypeAnnotation(
+            bindings,
+            /*first=*/std::nullopt,
+            /*allow_generic_type=*/
+            parse_fn_stubs_ ||
+                module_->attributes().contains(ModuleAttribute::kGenerics)));
     if (GenericTypeAnnotation* gta =
             dynamic_cast<GenericTypeAnnotation*>(type)) {
       name_def->set_definer(gta);
@@ -4222,6 +4228,17 @@ absl::StatusOr<ExprOrType> Parser::ParseParametricArg(Bindings& bindings) {
         return ParseCast(bindings, type_annotation);
       }
     }
+    if (auto name_ref = std::get_if<NameRef*>(&nocr); name_ref) {
+      // In a case like foo<T>(), the `nocr` for `T` is  a `NameRef` to the type
+      // variable, and what we want to yield is a TVTA for that.
+      AnyNameDef any_name_def = (*name_ref)->name_def();
+      if (auto name_def = std::get_if<const NameDef*>(&any_name_def);
+          name_def && (*name_def)->definer() &&
+          (*name_def)->definer()->kind() == AstNodeKind::kTypeAnnotation) {
+        return module_->Make<TypeVariableTypeAnnotation>(*name_ref,
+                                                         /*internal=*/false);
+      }
+    }
     // Otherwise, it's a value or an unadorned imported type.
     return ToExprNode(nocr);
   }

xls/dslx/frontend/parser_test.cc

@@ -2584,15 +2584,20 @@ TEST_F(ParserTest, ModuleWithParametric) {
 }
 
 TEST_F(ParserTest, ModuleWithGenericType) {
-  RoundTrip(R"(fn parametric<T: type>() -> u32 {
+  RoundTrip(R"(#![feature(generics)]
+
+fn parametric<T: type>() -> u32 {
     zero!<T>()
 })");
 }
 
 TEST_F(ParserTest, ModuleWithInvalidGenericType) {
-  constexpr std::string_view text = R"(fn non_parametric(T: type) -> u32 {
+  constexpr std::string_view text = R"(#![feature(generics)]
+
+fn non_parametric(T: type) -> u32 {
     zero!<T>()
-})";
+}
+)";
   Scanner s{file_table_, Fileno(0), std::string{text}};
   Parser parser{"test", &s};
   auto module_status = parser.ParseModule();

xls/dslx/interp_value.cc

@@ -66,6 +66,8 @@ std::string TagToString(InterpValueTag tag) {
       return "token";
     case InterpValueTag::kChannelReference:
       return "channel_reference";
+    case InterpValueTag::kTypeReference:
+      return "type_reference";
   }
   return absl::StrFormat("<invalid InterpValueTag(%d)>",
                          static_cast<int64_t>(tag));
@@ -187,7 +189,8 @@ std::string InterpValue::ToString(bool humanize,
       tag_ == InterpValueTag::kSBits || tag_ == InterpValueTag::kTuple ||
       tag_ == InterpValueTag::kArray || tag_ == InterpValueTag::kEnum ||
       tag_ == InterpValueTag::kFunction || tag_ == InterpValueTag::kToken ||
-      tag_ == InterpValueTag::kChannelReference) {
+      tag_ == InterpValueTag::kChannelReference ||
+      tag_ == InterpValueTag::kTypeReference) {
     return result;
   }
   LOG(FATAL) << "Unhandled tag: " << tag_;
@@ -240,6 +243,8 @@ std::string InterpValue::ToStringInternal(bool humanize,
           GetChannelReferenceOrDie().GetChannelId().has_value()
               ? absl::StrCat(*GetChannelReferenceOrDie().GetChannelId())
               : "none");
+    case InterpValueTag::kTypeReference:
+      return std::get<TypeReference>(payload_).string;
   }
   return "<INVALID>";
 }
@@ -426,6 +431,7 @@ bool InterpValue::Eq(const InterpValue& other) const {
     // bit value can be used in any place an enum type is annotated.
     case InterpValueTag::kSBits:
     case InterpValueTag::kUBits:
+    case InterpValueTag::kTypeReference:
     case InterpValueTag::kEnum: {
       return other.HasBits() && GetBitsOrDie() == other.GetBitsOrDie();
     }
@@ -472,6 +478,7 @@ bool InterpValue::operator==(const InterpValue& rhs) const { return Eq(rhs); }
   }
   switch (lhs.tag_) {
     case InterpValueTag::kUBits:
+    case InterpValueTag::kTypeReference:
       return MakeBool(ucmp(lhs.GetBitsOrDie(), rhs.GetBitsOrDie()));
     case InterpValueTag::kSBits:
       return MakeBool(scmp(lhs.GetBitsOrDie(), rhs.GetBitsOrDie()));
@@ -759,10 +766,20 @@ const Bits& InterpValue::GetBitsOrDie() const {
   if (std::holds_alternative<Bits>(payload_)) {
     return std::get<Bits>(payload_);
   }
+  if (auto* ref = std::get_if<TypeReference>(&payload_)) {
+    return ref->bits;
+  }
 
   return std::get<EnumData>(payload_).value;
 }
 
+absl::StatusOr<const TypeAnnotation*> InterpValue::GetTypeReference() const {
+  if (std::holds_alternative<TypeReference>(payload_)) {
+    return std::get<TypeReference>(payload_).annotation;
+  }
+  return absl::InvalidArgumentError("Value does not contain a type reference.");
+}
+
 absl::StatusOr<InterpValue::ChannelReference> InterpValue::GetChannelReference()
     const {
   if (std::holds_alternative<ChannelReference>(payload_)) {
@@ -1031,6 +1048,10 @@ absl::StatusOr<xls::Value> InterpValue::ConvertToIr() const {
       return absl::InvalidArgumentError(absl::StrFormat(
           "Cannot convert channel-reference-typed values to IR."));
     }
+    case InterpValueTag::kTypeReference: {
+      return absl::InvalidArgumentError(
+          absl::StrFormat("Cannot convert type-reference-typed values to IR."));
+    }
   }
   LOG(FATAL) << "Unhandled tag: " << tag_;
 }

xls/dslx/interp_value.h

@@ -54,6 +54,7 @@ enum class InterpValueTag : uint8_t {
   kFunction,
   kToken,
   kChannelReference,
+  kTypeReference,
 };
 
 std::string TagToString(InterpValueTag tag);
@@ -169,6 +170,14 @@ class InterpValue {
     return InterpValue(InterpValueTag::kToken, std::make_shared<TokenData>());
   }
 
+  static InterpValue MakeTypeReference(const TypeAnnotation* annotation) {
+    std::string string = annotation->ToString();
+    return InterpValue(InterpValueTag::kTypeReference,
+                       TypeReference{.annotation = annotation,
+                                     .string = string,
+                                     .bits = Bits::FromBytes(string)});
+  }
+
   // Queries
 
   bool IsTuple() const { return tag_ == InterpValueTag::kTuple; }
@@ -187,6 +196,9 @@ class InterpValue {
     return IsBuiltinFunction() &&
            std::get<Builtin>(GetFunctionOrDie()) == Builtin::kTrace;
   }
+  bool IsTypeReference() const {
+    return tag_ == InterpValueTag::kTypeReference;
+  }
 
   bool IsFalse() const { return IsBool() && GetBitsOrDie().IsZero(); }
   bool IsTrue() const { return IsBool() && GetBitsOrDie().IsAllOnes(); }
@@ -344,6 +356,7 @@ class InterpValue {
   const FnData& GetFunctionOrDie() const { return *GetFunction().value(); }
   absl::StatusOr<Bits> GetBits() const;
   const Bits& GetBitsOrDie() const;
+  absl::StatusOr<const TypeAnnotation*> GetTypeReference() const;
   absl::StatusOr<ChannelReference> GetChannelReference() const;
   ChannelReference GetChannelReferenceOrDie() const {
     return std::get<ChannelReference>(payload_);
@@ -370,7 +383,8 @@ class InterpValue {
   // apply to enum values as well.
   bool HasBits() const {
     return std::holds_alternative<Bits>(payload_) ||
-           std::holds_alternative<EnumData>(payload_);
+           std::holds_alternative<EnumData>(payload_) ||
+           std::holds_alternative<TypeReference>(payload_);
   }
 
   bool HasValues() const {
@@ -427,8 +441,15 @@ class InterpValue {
   //
   // TODO(leary): 2020-02-10 When all Python bindings are eliminated we can more
   // easily make an interpreter scoped lifetime that InterpValues can live in.
-  using Payload = std::variant<Bits, EnumData, std::vector<InterpValue>, FnData,
-                               std::shared_ptr<TokenData>, ChannelReference>;
+  struct TypeReference {
+    const TypeAnnotation* annotation;
+    std::string string;
+    Bits bits;
+  };
+
+  using Payload =
+      std::variant<Bits, EnumData, std::vector<InterpValue>, FnData,
+                   std::shared_ptr<TokenData>, ChannelReference, TypeReference>;
 
   InterpValue(InterpValueTag tag, Payload payload)
       : tag_(tag), payload_(std::move(payload)) {}

xls/dslx/ir_convert/function_converter.cc

@@ -3252,21 +3252,25 @@ absl::Status FunctionConverter::HandleFunction(
     VLOG(5) << "Resolving parametric binding: "
             << parametric_binding->ToString();
 
-    std::optional<InterpValue> parametric_value =
-        GetParametricBinding(parametric_binding->identifier());
-    XLS_RET_CHECK(parametric_value.has_value());
-    XLS_ASSIGN_OR_RETURN(std::unique_ptr<Type> parametric_type,
-                         ResolveType(parametric_binding->name_def()));
-    XLS_RET_CHECK(!parametric_type->IsMeta());
-    XLS_ASSIGN_OR_RETURN(TypeDim parametric_width_ctd,
-                         parametric_type->GetTotalBitCount());
-    XLS_ASSIGN_OR_RETURN(Value param_value,
-                         InterpValueToValue(*parametric_value));
-    const CValue evaluated = DefConst(parametric_binding, param_value);
-    const_prefill.SetNamedValue(parametric_binding->name_def()->identifier(),
-                                evaluated.ir_value);
-    XLS_RETURN_IF_ERROR(
-        DefAlias(parametric_binding, /*to=*/parametric_binding->name_def()));
+    if (!parametric_binding->type_annotation()
+             ->IsAnnotation<GenericTypeAnnotation>()) {
+      std::optional<InterpValue> parametric_value =
+          GetParametricBinding(parametric_binding->identifier());
+      XLS_RET_CHECK(parametric_value.has_value());
+      XLS_ASSIGN_OR_RETURN(std::unique_ptr<Type> parametric_type,
+                           ResolveType(parametric_binding->name_def()));
+      XLS_RET_CHECK(!parametric_type->IsMeta());
+      XLS_ASSIGN_OR_RETURN(TypeDim parametric_width_ctd,
+                           parametric_type->GetTotalBitCount());
+
+      XLS_ASSIGN_OR_RETURN(Value param_value,
+                           InterpValueToValue(*parametric_value));
+      const CValue evaluated = DefConst(parametric_binding, param_value);
+      const_prefill.SetNamedValue(parametric_binding->name_def()->identifier(),
+                                  evaluated.ir_value);
+      XLS_RETURN_IF_ERROR(
+          DefAlias(parametric_binding, /*to=*/parametric_binding->name_def()));
+    }
   }
 
   // If there is foreign function data, all constant values are replaced now.