Clean up LocIdAndInst::UncheckedLoc uses (#5151)

Adds subset conversion of `NodeIdOneOf` due to the choice usage, plus
the pre-existing TODO. Fixes incorrect information about nodes on
StructLiteral and TupleLiteral.

After this change, `UncheckedLoc` is only used in a couple import
contexts (hard to verify) plus `InstStore::GetWithLocId`.

toolchain/check/context.h

@@ -177,7 +177,9 @@ class Context {
   // known. This represents each binding to be done at the end of checking the
   // Choice type.
   struct ChoiceDeferredBinding {
-    Parse::NodeId node_id;
+    Parse::NodeIdOneOf<Parse::ChoiceAlternativeListCommaId,
+                       Parse::ChoiceDefinitionId>
+        node_id;
     NameComponent name_component;
   };
   auto choice_deferred_bindings() -> llvm::SmallVector<ChoiceDeferredBinding>& {

toolchain/check/handle_choice.cpp

@@ -138,8 +138,10 @@ auto HandleParseNode(Context& context, Parse::ChoiceDefinitionStartId node_id)
   return true;
 }
 
-static auto AddChoiceAlternative(Context& context, Parse::NodeId node_id)
-    -> void {
+static auto AddChoiceAlternative(
+    Context& context, Parse::NodeIdOneOf<Parse::ChoiceAlternativeListCommaId,
+                                         Parse::ChoiceDefinitionId>
+                          node_id) -> void {
   // Note, there is nothing like a ChoiceAlternativeIntroducer node, so no parse
   // node to pop here.
   auto name_component = PopNameComponent(context);
@@ -191,7 +193,7 @@ static auto MakeLetBinding(Context& context, const ChoiceInfo& choice_info,
     -> void {
   SemIR::InstId discriminant_value_id = [&] {
     if (choice_info.num_alternative_bits == 0) {
-      return AddInst(context, SemIR::LocIdAndInst::UncheckedLoc(
+      return AddInst(context, SemIR::LocIdAndInst(
                                   binding.node_id,
                                   SemIR::TupleLiteral{
                                       .type_id = GetTupleType(context, {}),
@@ -208,7 +210,7 @@ static auto MakeLetBinding(Context& context, const ChoiceInfo& choice_info,
 
   auto self_value_id = ConvertToValueOfType(
       context, binding.node_id,
-      AddInst(context, SemIR::LocIdAndInst::UncheckedLoc(
+      AddInst(context, SemIR::LocIdAndInst(
                            binding.node_id,
                            SemIR::StructLiteral{
                                .type_id = choice_info.self_struct_type_id,
@@ -226,12 +228,12 @@ static auto MakeLetBinding(Context& context, const ChoiceInfo& choice_info,
       {.name_id = binding.name_component.name_id,
        .parent_scope_id = choice_info.name_scope_id});
   auto bind_name_id = AddInst(
-      context, SemIR::LocIdAndInst::UncheckedLoc(
-                   binding.node_id, SemIR::BindName{
-                                        .type_id = choice_info.self_type_id,
-                                        .entity_name_id = entity_name_id,
-                                        .value_id = self_value_id,
-                                    }));
+      context, SemIR::LocIdAndInst(binding.node_id,
+                                   SemIR::BindName{
+                                       .type_id = choice_info.self_type_id,
+                                       .entity_name_id = entity_name_id,
+                                       .value_id = self_value_id,
+                                   }));
   context.name_scopes()
       .Get(choice_info.name_scope_id)
       .AddRequired({.name_id = binding.name_component.name_id,

toolchain/check/handle_let_and_var.cpp

@@ -121,7 +121,7 @@ static auto GetOrAddStorage(Context& context, SemIR::InstId var_pattern_id)
 
   return AddInst(
       context,
-      SemIR::LocIdAndInst::UncheckedLoc(
+      SemIR::LocIdAndInst(
           pattern.loc_id,
           SemIR::VarStorage{
               .type_id = pattern.inst.type_id(),

toolchain/check/import_ref.cpp

@@ -3277,12 +3277,9 @@ static auto HasCompatibleImportedNodeKind(SemIR::InstKind imported_kind,
   }
   if (imported_kind == SemIR::ImportDecl::Kind &&
       local_kind == SemIR::Namespace::Kind) {
-    // TODO: Consider supporting NodeIdOneOf conversions to make the below work.
-    // But this extra validation is the primary use-case at present.
-    // static_assert(
-    //     std::is_convertible_v<
-    //         decltype(SemIR::ImportDecl::Kind)::TypedNodeId,
-    //         decltype(SemIR::Namespace::Kind)::TypedNodeId>);
+    static_assert(
+        std::is_convertible_v<decltype(SemIR::ImportDecl::Kind)::TypedNodeId,
+                              decltype(SemIR::Namespace::Kind)::TypedNodeId>);
     return true;
   }
   return false;

toolchain/parse/node_ids.h

@@ -107,6 +107,23 @@ using AnyPackageNameId = NodeIdInCategory<NodeCategory::PackageName>;
 template <typename... T>
   requires(sizeof...(T) >= 2)
 struct NodeIdOneOf : public NodeId {
+ private:
+  // True if `OtherT` is one of `T`.
+  template <typename OtherT>
+  static constexpr bool Contains = (std::is_same<OtherT, T>{} || ...);
+
+  // True if `NodeIdOneOf<SubsetT...>` is a subset of this `NodeIdOneOf`.
+  template <typename Unused>
+  static constexpr bool IsSubset = false;
+  template <typename... SubsetT>
+  static constexpr bool IsSubset<NodeIdOneOf<SubsetT...>> =
+      (Contains<SubsetT> && ...);
+
+  // Private to prevent accidental explicit construction from an untyped
+  // NodeId.
+  explicit constexpr NodeIdOneOf(NodeId node_id) : NodeId(node_id) {}
+
+ public:
   // Provide a factory function for construction from `NodeId`. This doesn't
   // validate the type, so it's unsafe.
   static constexpr auto UnsafeMake(NodeId node_id) -> NodeIdOneOf {
@@ -114,17 +131,17 @@ struct NodeIdOneOf : public NodeId {
   }
 
   template <const NodeKind& Kind>
-    requires((T::Kind == Kind) || ...)
+    requires(Contains<NodeIdForKind<Kind>>)
   // NOLINTNEXTLINE(google-explicit-constructor)
   NodeIdOneOf(NodeIdForKind<Kind> node_id) : NodeId(node_id) {}
 
+  template <typename OtherNodeIdOneOf>
+    requires(IsSubset<OtherNodeIdOneOf>)
   // NOLINTNEXTLINE(google-explicit-constructor)
-  constexpr NodeIdOneOf(NoneNodeId /*none*/) : NodeId(NoneIndex) {}
+  NodeIdOneOf(OtherNodeIdOneOf node_id) : NodeId(node_id) {}
 
- private:
-  // Private to prevent accidental explicit construction from an untyped
-  // NodeId.
-  explicit constexpr NodeIdOneOf(NodeId node_id) : NodeId(node_id) {}
+  // NOLINTNEXTLINE(google-explicit-constructor)
+  constexpr NodeIdOneOf(NoneNodeId /*none*/) : NodeId(NoneIndex) {}
 };
 
 using AnyClassDeclId =

toolchain/parse/typed_nodes_test.cpp

@@ -8,6 +8,7 @@
 #include <gtest/gtest.h>
 
 #include <forward_list>
+#include <type_traits>
 
 #include "toolchain/lex/lex.h"
 #include "toolchain/lex/tokenized_buffer.h"
@@ -399,5 +400,23 @@ TEST_F(TypedNodeTest, CategoryMatches) {
 #include "toolchain/parse/node_kind.def"
 }
 
+TEST_F(TypedNodeTest, NodeIdOneOf) {
+  auto id1 = Parse::ChoiceDefinitionId::None;
+  Parse::NodeIdOneOf<Parse::ChoiceDefinitionId, Parse::ChoiceIntroducer> id2 =
+      id1;
+  Parse::NodeIdOneOf<Parse::ChoiceDefinitionId, Parse::ChoiceIntroducer,
+                     Parse::ChoiceDefinitionStartId>
+      id3 = id1;
+  Parse::NodeIdOneOf<Parse::ChoiceDefinitionId, Parse::ChoiceIntroducer,
+                     Parse::ChoiceSignature>
+      id4 = id1;
+  id3 = id2;
+  id4 = id2;
+  static_assert(!std::is_assignable_v<decltype(id2), decltype(id3)>);
+  static_assert(!std::is_assignable_v<decltype(id2), decltype(id4)>);
+  static_assert(!std::is_assignable_v<decltype(id4), decltype(id3)>);
+  static_assert(!std::is_assignable_v<decltype(id3), decltype(id4)>);
+}
+
 }  // namespace
 }  // namespace Carbon::Parse

toolchain/sem_ir/typed_insts.h

@@ -1503,10 +1503,10 @@ struct StructInit {
 
 // A literal struct value, such as `{.a = 1, .b = 2}`.
 struct StructLiteral {
-  static constexpr auto Kind =
-      InstKind::StructLiteral.Define<Parse::StructLiteralId>(
-          {.ir_name = "struct_literal",
-           .constant_kind = InstConstantKind::Never});
+  static constexpr auto Kind = InstKind::StructLiteral.Define<
+      Parse::NodeIdOneOf<Parse::ChoiceAlternativeListCommaId,
+                         Parse::ChoiceDefinitionId, Parse::StructLiteralId>>(
+      {.ir_name = "struct_literal", .constant_kind = InstConstantKind::Never});
 
   TypeId type_id;
   InstBlockId elements_id;
@@ -1580,10 +1580,10 @@ struct TupleInit {
 
 // A literal tuple value.
 struct TupleLiteral {
-  static constexpr auto Kind =
-      InstKind::TupleLiteral.Define<Parse::TupleLiteralId>(
-          {.ir_name = "tuple_literal",
-           .constant_kind = InstConstantKind::Never});
+  static constexpr auto Kind = InstKind::TupleLiteral.Define<
+      Parse::NodeIdOneOf<Parse::ChoiceAlternativeListCommaId,
+                         Parse::ChoiceDefinitionId, Parse::TupleLiteralId>>(
+      {.ir_name = "tuple_literal", .constant_kind = InstConstantKind::Never});
 
   TypeId type_id;
   InstBlockId elements_id;