From 82179bc9039e3b39ce83305a9de45c7bc23050a2 Mon Sep 17 00:00:00 2001
From: josh11b <josh11b@users.noreply.github.com>
Date: Thu, 4 Jan 2024 13:51:19 -0800
Subject: [PATCH] Add `Tree::TryAs` to cast to a typed parse node id (#3565)

Co-authored-by: Richard Smith <richard@metafoo.co.uk>
---
 toolchain/parse/tree.h        | 43 +++++++++++++++++++++++++++++++++++
 toolchain/parse/tree_test.cpp | 41 +++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/toolchain/parse/tree.h b/toolchain/parse/tree.h
index a611de1209de..86e519eaff20 100644
--- a/toolchain/parse/tree.h
+++ b/toolchain/parse/tree.h
@@ -125,6 +125,27 @@ class Tree : public Printable<Tree> {
     return !node_has_error(id);
   }
 
+  // Converts `n` to a constrained node id `T` if the `node_kind(n)` matches
+  // the constraint on `T`.
+  template <typename T>
+  auto TryAs(NodeId n) const -> std::optional<T> {
+    CARBON_DCHECK(n.is_valid());
+    if (ConvertTo<T>::AllowedFor(node_kind(n))) {
+      return T(n);
+    } else {
+      return std::nullopt;
+    }
+  }
+
+  // Converts to `n` to a constrained node id `T`. Checks that the
+  // `node_kind(n)` matches the constraint on `T`.
+  template <typename T>
+  auto As(NodeId n) const -> T {
+    CARBON_DCHECK(n.is_valid());
+    CARBON_CHECK(ConvertTo<T>::AllowedFor(node_kind(n)));
+    return T(n);
+  }
+
   auto packaging_directive() const -> const std::optional<PackagingDirective>& {
     return packaging_directive_;
   }
@@ -226,6 +247,9 @@ class Tree : public Printable<Tree> {
  private:
   friend class Context;
 
+  template <typename T>
+  struct ConvertTo;
+
   // The in-memory representation of data used for a particular node in the
   // tree.
   struct NodeImpl {
@@ -452,6 +476,25 @@ auto Tree::Extract(IdT id) const
   return ExtractNodeFromChildren<T>(children(id));
 }
 
+template <const NodeKind& K>
+struct Tree::ConvertTo<NodeIdForKind<K>> {
+  static auto AllowedFor(NodeKind kind) -> bool { return kind == K; }
+};
+
+template <NodeCategory C>
+struct Tree::ConvertTo<NodeIdInCategory<C>> {
+  static auto AllowedFor(NodeKind kind) -> bool {
+    return !!(kind.category() & C);
+  }
+};
+
+template <typename T, typename U>
+struct Tree::ConvertTo<NodeIdOneOf<T, U>> {
+  static auto AllowedFor(NodeKind kind) -> bool {
+    return kind == T::Kind || kind == U::Kind;
+  }
+};
+
 }  // namespace Carbon::Parse
 
 #endif  // CARBON_TOOLCHAIN_PARSE_TREE_H_
diff --git a/toolchain/parse/tree_test.cpp b/toolchain/parse/tree_test.cpp
index dcbe0ad13bf6..075cbd954967 100644
--- a/toolchain/parse/tree_test.cpp
+++ b/toolchain/parse/tree_test.cpp
@@ -56,6 +56,47 @@ TEST_F(TreeTest, IsValid) {
   EXPECT_TRUE((*tree.postorder().begin()).is_valid());
 }
 
+TEST_F(TreeTest, AsAndTryAs) {
+  Lex::TokenizedBuffer& tokens = GetTokenizedBuffer("fn F();");
+  Tree tree = Parse(tokens, consumer_, /*vlog_stream=*/nullptr);
+  ASSERT_FALSE(tree.has_errors());
+  auto it = tree.roots().begin();
+  // A FileEnd node, so won't match.
+  NodeId n = *it;
+
+  // NodeIdForKind
+  std::optional<FunctionDeclId> fn_decl_id = tree.TryAs<FunctionDeclId>(n);
+  EXPECT_FALSE(fn_decl_id.has_value());
+  // NodeIdOneOf
+  std::optional<AnyFunctionDeclId> any_fn_decl_id =
+      tree.TryAs<AnyFunctionDeclId>(n);
+  EXPECT_FALSE(any_fn_decl_id.has_value());
+  // NodeIdInCategory
+  std::optional<AnyDeclId> any_decl_id = tree.TryAs<AnyDeclId>(n);
+  EXPECT_FALSE(any_decl_id.has_value());
+
+  ++it;
+  n = *it;
+  // A FunctionDecl node, so will match.
+  fn_decl_id = tree.TryAs<FunctionDeclId>(n);
+  ASSERT_TRUE(fn_decl_id.has_value());
+  EXPECT_TRUE(*fn_decl_id == n);
+  FunctionDeclId fn_decl_id2 = tree.As<FunctionDeclId>(n);
+  EXPECT_TRUE(*fn_decl_id == fn_decl_id2);
+
+  any_fn_decl_id = tree.TryAs<AnyFunctionDeclId>(n);
+  ASSERT_TRUE(any_fn_decl_id.has_value());
+  EXPECT_TRUE(*any_fn_decl_id == n);
+  AnyFunctionDeclId any_fn_decl_id2 = tree.As<AnyFunctionDeclId>(n);
+  EXPECT_TRUE(*any_fn_decl_id == any_fn_decl_id2);
+
+  any_decl_id = tree.TryAs<AnyDeclId>(n);
+  ASSERT_TRUE(any_decl_id.has_value());
+  EXPECT_TRUE(*any_decl_id == n);
+  AnyDeclId any_decl_id2 = tree.As<AnyDeclId>(n);
+  EXPECT_TRUE(*any_decl_id == any_decl_id2);
+}
+
 TEST_F(TreeTest, PrintPostorderAsYAML) {
   Lex::TokenizedBuffer& tokens = GetTokenizedBuffer("fn F();");
   Tree tree = Parse(tokens, consumer_, /*vlog_stream=*/nullptr);