Add CARBON_KIND_SWITCH to better handle typed inst switches. (#3820)

Converts a switch in formatter.cpp as an example of resulting syntax.

toolchain/base/BUILD

@@ -15,6 +15,14 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "kind_switch",
+    hdrs = ["kind_switch.h"],
+    deps = [
+        "@llvm-project//llvm:Support",
+    ],
+)
+
 cc_library(
     name = "pretty_stack_trace_function",
     hdrs = ["pretty_stack_trace_function.h"],

toolchain/base/kind_switch.h

@@ -0,0 +1,93 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef CARBON_TOOLCHAIN_BASE_KIND_SWITCH_H_
+#define CARBON_TOOLCHAIN_BASE_KIND_SWITCH_H_
+
+#include "llvm/ADT/STLExtras.h"
+
+// This library provides switch-like behaviors for Carbon's kind-based types.
+//
+// An expected use case is to mix regular switch `case` statements and
+// `CARBON_KIND`. However, the `switch` must be defined using
+// `CARBON_KIND_SWITCH`. For example:
+//
+//   CARBON_KIND_SWITCH(untyped_inst) {
+//     case CARBON_KIND(SomeInstType inst): {
+//       return inst.typed_field;
+//     }
+//     case OtherType1::Kind:
+//     case OtherType2::Kind:
+//       return value;
+//     default:
+//       return default_value;
+//   }
+//
+// For compatibility, this requires:
+//
+// - The type passed to `CARBON_KIND_SWITCH` has `.kind()` to switch on, and
+//   `.As<CaseT>` for `CARBON_KIND` to cast to.
+// - Each type passed to `CARBON_KIND` (`CaseT` above) provides
+//   `CaseT::Kind`, which is passed to the `case` keyword.
+//   `CaseT::Kind::RawEnumType` is the type returned by `.kind()`.
+//
+// Note, this is currently used primarily for Inst in toolchain. When more
+// use-cases are added, it would be worth considering whether the API
+// requirements should change.
+namespace Carbon::Internal::Kind {
+
+// Given `CARBON_KIND_SWITCH(value)` this handles calling `value.kind()`.
+template <typename T>
+auto SwitchOn(T&& switch_value) -> auto {
+  return switch_value.kind();
+}
+
+// Given `CARBON_KIND(CaseT name)` this generates `CaseT::Kind`. It explicitly
+// returns `KindT` because that may differ from `CaseT::Kind`, and may not be
+// copyable.
+template <typename FnT>
+consteval auto ForCase() -> auto {
+  using ArgT = llvm::function_traits<FnT>::template arg_t<0>;
+  return static_cast<decltype(ArgT::Kind)::RawEnumType>(ArgT::Kind);
+}
+
+// Given `CARBON_KIND_SWITCH(value)` and `CARBON_KIND(CaseT name)` this
+// generates `value.As<CaseT>()`.
+template <typename FnT, typename ValueT>
+auto Cast(ValueT&& kind_switch_value) -> auto {
+  using CaseT = llvm::function_traits<FnT>::template arg_t<0>;
+  return kind_switch_value.template As<CaseT>();
+}
+
+#define CARBON_INTERNAL_KIND_MERGE(Prefix, Line) Prefix##Line
+#define CARBON_INTERNAL_KIND_LABEL(Line) \
+  CARBON_INTERNAL_KIND_MERGE(carbon_internal_kind_case_, Line)
+
+}  // namespace Carbon::Internal::Kind
+
+// Produces a switch statement on value.kind().
+#define CARBON_KIND_SWITCH(value)                            \
+  switch (                                                   \
+      const auto& carbon_internal_kind_switch_value = value; \
+      ::Carbon::Internal::Kind::SwitchOn(carbon_internal_kind_switch_value))
+
+// Produces a case-compatible block of code that also instantiates a local typed
+// variable. typed_variable_declaration looks like `int i`, with a space.
+//
+// This uses `if` to scope the variable, and provides a dangling `else` in order
+// to prevent accidental `else` use. The label allows `:` to follow the macro
+// name, making it look more like a typical `case`.
+//
+// NOLINTBEGIN(bugprone-macro-parentheses)
+#define CARBON_KIND(typed_variable_declaration)                                \
+  ::Carbon::Internal::Kind::ForCase<                                           \
+      decltype([]([[maybe_unused]] typed_variable_declaration) {})>()          \
+      : if (typed_variable_declaration = ::Carbon::Internal::Kind::Cast<       \
+                decltype([]([[maybe_unused]] typed_variable_declaration) {})>( \
+                carbon_internal_kind_switch_value);                            \
+            false) {}                                                          \
+  else [[maybe_unused]] CARBON_INTERNAL_KIND_LABEL(__LINE__)
+// NOLINTEND(bugprone-macro-parentheses)
+
+#endif  // CARBON_TOOLCHAIN_BASE_KIND_SWITCH_H_

toolchain/sem_ir/BUILD

@@ -127,6 +127,7 @@ cc_library(
         ":ids",
         ":inst_kind",
         "//common:ostream",
+        "//toolchain/base:kind_switch",
         "//toolchain/base:value_store",
         "//toolchain/lex:tokenized_buffer",
         "//toolchain/parse:tree",

toolchain/sem_ir/formatter.cpp

@@ -9,6 +9,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/SaveAndRestore.h"
+#include "toolchain/base/kind_switch.h"
 #include "toolchain/base/value_store.h"
 #include "toolchain/lex/tokenized_buffer.h"
 #include "toolchain/parse/tree.h"
@@ -446,7 +447,7 @@ class InstNamer {
         continue;
       }
 
-      auto inst = sem_ir_.insts().Get(inst_id);
+      auto untyped_inst = sem_ir_.insts().Get(inst_id);
       auto add_inst_name = [&](std::string name) {
         insts[inst_id.index] = {
             scope_id, scope.insts.AllocateName(
@@ -457,59 +458,55 @@ class InstNamer {
             (sem_ir_.names().GetIRBaseName(name_id).str() + suffix).str());
       };
 
-      if (auto branch = inst.TryAs<AnyBranch>()) {
+      if (auto branch = untyped_inst.TryAs<AnyBranch>()) {
         AddBlockLabel(scope_id, sem_ir_.insts().GetNodeId(inst_id), *branch);
       }
 
-      switch (inst.kind()) {
-        case AddrPattern::Kind: {
+      CARBON_KIND_SWITCH(untyped_inst) {
+        case CARBON_KIND(AddrPattern inst): {
           // TODO: We need to assign names to parameters that appear in
           // function declarations, which may be nested within a pattern. For
           // now, just look through `addr`, but we should find a better way to
           // visit parameters.
-          CollectNamesInBlock(scope_id, inst.As<AddrPattern>().inner_id);
+          CollectNamesInBlock(scope_id, inst.inner_id);
           break;
         }
-        case AssociatedConstantDecl::Kind: {
-          add_inst_name_id(inst.As<AssociatedConstantDecl>().name_id);
+        case CARBON_KIND(AssociatedConstantDecl inst): {
+          add_inst_name_id(inst.name_id);
           continue;
         }
         case BindAlias::Kind:
         case BindName::Kind:
         case BindSymbolicName::Kind: {
-          add_inst_name_id(sem_ir_.bind_names()
-                               .Get(inst.As<AnyBindName>().bind_name_id)
-                               .name_id);
+          auto inst = untyped_inst.As<AnyBindName>();
+          add_inst_name_id(sem_ir_.bind_names().Get(inst.bind_name_id).name_id);
           continue;
         }
-        case ClassDecl::Kind: {
-          add_inst_name_id(
-              sem_ir_.classes().Get(inst.As<ClassDecl>().class_id).name_id,
-              ".decl");
-          CollectNamesInBlock(scope_id, inst.As<ClassDecl>().decl_block_id);
+        case CARBON_KIND(ClassDecl inst): {
+          add_inst_name_id(sem_ir_.classes().Get(inst.class_id).name_id,
+                           ".decl");
+          CollectNamesInBlock(scope_id, inst.decl_block_id);
           continue;
         }
-        case ClassType::Kind: {
-          add_inst_name_id(
-              sem_ir_.classes().Get(inst.As<ClassType>().class_id).name_id);
+        case CARBON_KIND(ClassType inst): {
+          add_inst_name_id(sem_ir_.classes().Get(inst.class_id).name_id);
           continue;
         }
-        case FunctionDecl::Kind: {
-          add_inst_name_id(sem_ir_.functions()
-                               .Get(inst.As<FunctionDecl>().function_id)
-                               .name_id);
-          CollectNamesInBlock(scope_id, inst.As<FunctionDecl>().decl_block_id);
+        case CARBON_KIND(FunctionDecl inst): {
+          add_inst_name_id(sem_ir_.functions().Get(inst.function_id).name_id);
+          CollectNamesInBlock(scope_id, inst.decl_block_id);
           continue;
         }
-        case ImplDecl::Kind: {
-          CollectNamesInBlock(scope_id, inst.As<ImplDecl>().decl_block_id);
+        case CARBON_KIND(ImplDecl inst): {
+          CollectNamesInBlock(scope_id, inst.decl_block_id);
           break;
         }
         case ImportRefUnused::Kind:
         case ImportRefUsed::Kind: {
           add_inst_name("import_ref");
-          // When building import refs, we frequently add instructions without a
-          // block. Constants that refer to them need to be separately named.
+          // When building import refs, we frequently add instructions without
+          // a block. Constants that refer to them need to be separately
+          // named.
           auto const_id = sem_ir_.constant_values().Get(inst_id);
           if (const_id.is_valid() && const_id.is_template() &&
               !insts[const_id.inst_id().index].second) {
@@ -517,35 +514,32 @@ class InstNamer {
           }
           continue;
         }
-        case InterfaceDecl::Kind: {
-          add_inst_name_id(sem_ir_.interfaces()
-                               .Get(inst.As<InterfaceDecl>().interface_id)
-                               .name_id,
+        case CARBON_KIND(InterfaceDecl inst): {
+          add_inst_name_id(sem_ir_.interfaces().Get(inst.interface_id).name_id,
                            ".decl");
-          CollectNamesInBlock(scope_id, inst.As<InterfaceDecl>().decl_block_id);
+          CollectNamesInBlock(scope_id, inst.decl_block_id);
           continue;
         }
-        case NameRef::Kind: {
-          add_inst_name_id(inst.As<NameRef>().name_id, ".ref");
+        case CARBON_KIND(NameRef inst): {
+          add_inst_name_id(inst.name_id, ".ref");
           continue;
         }
         // The namespace is specified here due to the name conflict.
-        case SemIR::Namespace::Kind: {
-          add_inst_name_id(sem_ir_.name_scopes()
-                               .Get(inst.As<SemIR::Namespace>().name_scope_id)
-                               .name_id);
+        case CARBON_KIND(SemIR::Namespace inst): {
+          add_inst_name_id(
+              sem_ir_.name_scopes().Get(inst.name_scope_id).name_id);
           continue;
         }
-        case Param::Kind: {
-          add_inst_name_id(inst.As<Param>().name_id);
+        case CARBON_KIND(Param inst): {
+          add_inst_name_id(inst.name_id);
           continue;
         }
-        case SpliceBlock::Kind: {
-          CollectNamesInBlock(scope_id, inst.As<SpliceBlock>().block_id);
+        case CARBON_KIND(SpliceBlock inst): {
+          CollectNamesInBlock(scope_id, inst.block_id);
           break;
         }
-        case VarStorage::Kind: {
-          add_inst_name_id(inst.As<VarStorage>().name_id, ".var");
+        case CARBON_KIND(VarStorage inst): {
+          add_inst_name_id(inst.name_id, ".var");
           continue;
         }
         default: {
@@ -554,7 +548,7 @@ class InstNamer {
       }
 
       // Sequentially number all remaining values.
-      if (inst.kind().value_kind() != InstValueKind::None) {
+      if (untyped_inst.kind().value_kind() != InstValueKind::None) {
         add_inst_name("");
       }
     }