Add typed forms of `TokenIndex` for tokens with a known kind. (#4016)

Instead of adding a `Token<Kind>` field to typed parse nodes that
contains a `TokenIndex`, make the parse node directly hold a
`TokenIndex` of a more specific kind.

toolchain/lex/BUILD

@@ -199,7 +199,10 @@ cc_library(
 cc_library(
     name = "token_index",
     hdrs = ["token_index.h"],
-    deps = ["//toolchain/base:index_base"],
+    deps = [
+        ":token_kind",
+        "//toolchain/base:index_base",
+    ],
 )
 
 cc_library(

toolchain/lex/token_index.h

@@ -6,6 +6,7 @@
 #define CARBON_TOOLCHAIN_LEX_TOKEN_INDEX_H_
 
 #include "toolchain/base/index_base.h"
+#include "toolchain/lex/token_kind.h"
 
 namespace Carbon::Lex {
 
@@ -32,6 +33,21 @@ struct TokenIndex : public IndexBase {
 constexpr TokenIndex TokenIndex::Invalid(TokenIndex::InvalidIndex);
 constexpr TokenIndex TokenIndex::FirstNonCommentToken(1);
 
+// A lightweight handle to a lexed token in a `TokenizedBuffer` whose kind is
+// known to be `Kind`.
+template <const TokenKind& K>
+struct TokenIndexForKind : public TokenIndex {
+  // NOLINTNEXTLINE(readability-identifier-naming)
+  static const TokenKind& Kind;
+  constexpr explicit TokenIndexForKind(TokenIndex index) : TokenIndex(index) {}
+};
+template <const TokenKind& K>
+const TokenKind& TokenIndexForKind<K>::Kind = K;
+
+#define CARBON_TOKEN(TokenName) \
+  using TokenName##TokenIndex = TokenIndexForKind<TokenKind::TokenName>;
+#include "toolchain/lex/token_kind.def"
+
 }  // namespace Carbon::Lex
 
 #endif  // CARBON_TOOLCHAIN_LEX_TOKEN_INDEX_H_

toolchain/parse/extract.cpp

@@ -330,10 +330,11 @@ auto NodeExtractor::MatchesTokenKind(Lex::TokenKind expected_kind) const
 
 // Extract the token corresponding to a node.
 template <const Lex::TokenKind& Kind>
-struct Extractable<Token<Kind>> {
-  static auto Extract(NodeExtractor& extractor) -> std::optional<Token<Kind>> {
+struct Extractable<Lex::TokenIndexForKind<Kind>> {
+  static auto Extract(NodeExtractor& extractor)
+      -> std::optional<Lex::TokenIndexForKind<Kind>> {
     if (extractor.MatchesTokenKind(Kind)) {
-      return Token<Kind>{.index = extractor.token()};
+      return static_cast<Lex::TokenIndexForKind<Kind>>(extractor.token());
     } else {
       return std::nullopt;
     }
@@ -342,15 +343,16 @@ struct Extractable<Token<Kind>> {
 
 // Extract the token corresponding to a node.
 template <>
-struct Extractable<AnyToken> {
-  static auto Extract(NodeExtractor& extractor) -> std::optional<AnyToken> {
+struct Extractable<Lex::TokenIndex> {
+  static auto Extract(NodeExtractor& extractor)
+      -> std::optional<Lex::TokenIndex> {
     if (!extractor.has_token()) {
       if (auto* trace = extractor.trace()) {
         *trace << "Token expected but processing root node\n";
       }
       return std::nullopt;
     }
-    return AnyToken{.index = extractor.token()};
+    return extractor.token();
   }
 };
 

toolchain/parse/node_ids.h

@@ -111,32 +111,6 @@ struct NodeIdNot : public NodeId {
       : NodeId(NodeId::InvalidIndex) {}
 };
 
-// This class holds the token corresponding to a parse node, and defines the
-// expected token kind. The specified token index will always have kind K if the
-// enclosing node doesn't have errors. Note that we never try to extract a node
-// that has errors, and there are no restrictions on the token kind in that
-// case.
-template <const Lex::TokenKind& K>
-struct Token {
-  static constexpr const Lex::TokenKind& Kind = K;
-  Lex::TokenIndex index;
-};
-
-// This class holds the token corresponding to a parse node in the case where
-// the parse node can correspond to any token. This should only be used when the
-// node kind is either not used in a finished tree, such as `Placeholder`, or is
-// always invalid, such as `InvalidParse`.
-struct AnyToken {
-  Lex::TokenIndex index;
-};
-
-// This class holds the token corresponding to a parse node in the case of a
-// virtual token. The parse node doesn't actually own the token in this case.
-template <const Lex::TokenKind& K>
-struct VirtualToken {
-  Token<K> token;
-};
-
 // Note that the support for extracting these types using the `Tree::Extract*`
 // functions is defined in `extract.cpp`.
 

toolchain/parse/typed_nodes_test.cpp

@@ -315,15 +315,15 @@ TEST_F(TypedNodeTest, Token) {
 
   auto n_var = tree->ExtractAs<VariableDecl>(file.decls[0]);
   ASSERT_TRUE(n_var.has_value());
-  EXPECT_EQ(tokens->GetKind(n_var->token.index), Lex::TokenKind::Semi);
+  EXPECT_EQ(tokens->GetKind(n_var->token), Lex::TokenKind::Semi);
 
   auto n_intro = tree->ExtractAs<VariableIntroducer>(n_var->introducer);
   ASSERT_TRUE(n_intro.has_value());
-  EXPECT_EQ(tokens->GetKind(n_intro->token.index), Lex::TokenKind::Var);
+  EXPECT_EQ(tokens->GetKind(n_intro->token), Lex::TokenKind::Var);
 
   auto n_patt = tree->ExtractAs<BindingPattern>(n_var->pattern);
   ASSERT_TRUE(n_patt.has_value());
-  EXPECT_EQ(tokens->GetKind(n_patt->token.index), Lex::TokenKind::Colon);
+  EXPECT_EQ(tokens->GetKind(n_patt->token), Lex::TokenKind::Colon);
 }
 
 TEST_F(TypedNodeTest, VerifyInvalid) {