Create KJ_SWITCH_ONEOF to switch over OneOf.

Up until now, I've always felt kj::OneOf had a rather janky interface, which made me avoid using it in other interfaces. With this macro I'm pretty happy with it as *the* way to define variants.

c++/src/kj/one-of-test.c++

@@ -98,4 +98,41 @@ TEST(OneOf, Copy) {
   EXPECT_STREQ("foo", var2.get<const char*>());
 }
 
+TEST(OneOf, Switch) {
+  OneOf<int, float, const char*> var;
+  var = "foo";
+  uint count = 0;
+
+  {
+    KJ_SWITCH_ONEOF(var) {
+      KJ_CASE_ONEOF(i, int) {
+        KJ_FAIL_ASSERT("expected char*, got int", i);
+      }
+      KJ_CASE_ONEOF(s, const char*) {
+        KJ_EXPECT(kj::StringPtr(s) == "foo");
+        ++count;
+      }
+      KJ_CASE_ONEOF(n, float) {
+        KJ_FAIL_ASSERT("expected char*, got float", n);
+      }
+    }
+  }
+
+  KJ_EXPECT(count == 1);
+
+  {
+    KJ_SWITCH_ONEOF(kj::cp(var)) {
+      KJ_CASE_ONEOF(i, int) {
+        KJ_FAIL_ASSERT("expected char*, got int", i);
+      }
+      KJ_CASE_ONEOF(s, const char*) {
+        KJ_EXPECT(kj::StringPtr(s) == "foo");
+      }
+      KJ_CASE_ONEOF(n, float) {
+        KJ_FAIL_ASSERT("expected char*, got float", n);
+      }
+    }
+  }
+}
+
 }  // namespace kj

c++/src/kj/one-of.h

@@ -37,6 +37,31 @@ struct TypeIndex_ { static constexpr uint value = TypeIndex_<i + 1, Key, Rest...
 template <uint i, typename Key, typename... Rest>
 struct TypeIndex_<i, Key, Key, Rest...> { static constexpr uint value = i; };
 
+enum class Variants0 {};
+enum class Variants1 { _variant0 };
+enum class Variants2 { _variant0, _variant1 };
+enum class Variants3 { _variant0, _variant1, _variant2 };
+enum class Variants4 { _variant0, _variant1, _variant2, _variant3 };
+enum class Variants5 { _variant0, _variant1, _variant2, _variant3, _variant4 };
+enum class Variants6 { _variant0, _variant1, _variant2, _variant3, _variant4, _variant5 };
+enum class Variants7 { _variant0, _variant1, _variant2, _variant3, _variant4, _variant5, _variant6 };
+enum class Variants8 { _variant0, _variant1, _variant2, _variant3, _variant4, _variant5, _variant6,
+                       _variant7 };
+
+template <uint i> struct Variants_;
+template <> struct Variants_<0> { typedef Variants0 Type; };
+template <> struct Variants_<1> { typedef Variants1 Type; };
+template <> struct Variants_<2> { typedef Variants2 Type; };
+template <> struct Variants_<3> { typedef Variants3 Type; };
+template <> struct Variants_<4> { typedef Variants4 Type; };
+template <> struct Variants_<5> { typedef Variants5 Type; };
+template <> struct Variants_<6> { typedef Variants6 Type; };
+template <> struct Variants_<7> { typedef Variants7 Type; };
+template <> struct Variants_<8> { typedef Variants8 Type; };
+
+template <uint i>
+using Variants = typename Variants_<i>::Type;
+
 }  // namespace _ (private)
 
 template <typename... Variants>
@@ -48,7 +73,12 @@ class OneOf {
 public:
   inline OneOf(): tag(0) {}
   OneOf(const OneOf& other) { copyFrom(other); }
+  OneOf(OneOf& other) { copyFrom(other); }
   OneOf(OneOf&& other) { moveFrom(other); }
+  template <typename T>
+  OneOf(T&& other): tag(typeIndex<Decay<T>>()) {
+    ctor(*reinterpret_cast<Decay<T>*>(space), kj::fwd<T>(other));
+  }
   ~OneOf() { destroy(); }
 
   OneOf& operator=(const OneOf& other) { if (tag != 0) destroy(); copyFrom(other); return *this; }
@@ -96,6 +126,22 @@ public:
   // block call allHandled<n>() where n is the number of variants. This will fail to compile
   // if new variants are added in the future.
 
+  typedef _::Variants<sizeof...(Variants)> Tag;
+
+  Tag which() {
+    KJ_IREQUIRE(tag != 0, "Can't KJ_SWITCH_ONEOF() on uninitialized value.");
+    return static_cast<Tag>(tag - 1);
+  }
+
+  template <typename T>
+  static constexpr Tag tagFor() {
+    return static_cast<Tag>(typeIndex<T>() - 1);
+  }
+
+  OneOf* _switchSubject() & { return this; }
+  const OneOf* _switchSubject() const& { return this; }
+  _::NullableValue<OneOf> _switchSubject() && { return kj::mv(*this); }
+
 private:
   uint tag;
 
@@ -150,6 +196,20 @@ private:
     doAll(copyVariantFrom<Variants>(other)...);
   }
 
+  template <typename T>
+  inline bool copyVariantFrom(OneOf& other) {
+    if (other.is<T>()) {
+      ctor(*reinterpret_cast<T*>(space), other.get<T>());
+    }
+    return false;
+  }
+  void copyFrom(OneOf& other) {
+    // Initialize as a copy of `other`.  Expects that `this` starts out uninitialized, so the tag
+    // is invalid.
+    tag = other.tag;
+    doAll(copyVariantFrom<Variants>(other)...);
+  }
+
   template <typename T>
   inline bool moveVariantFrom(OneOf& other) {
     if (other.is<T>()) {
@@ -176,6 +236,53 @@ void OneOf<Variants...>::allHandled() {
   KJ_UNREACHABLE;
 }
 
+#if __cplusplus > 201402L
+#define KJ_SWITCH_ONEOF(value) \
+  switch (auto _kj_switch_subject = value._switchSubject(); _kj_switch_subject->which())
+#else
+#define KJ_SWITCH_ONEOF(value) \
+  // Without C++17, we can only support one switch per containing block. Deal with it.
+  auto _kj_switch_subject = value._switchSubject(); \
+  switch (_kj_switch_subject->which())
+#endif
+#define KJ_CASE_ONEOF(name, ...) \
+    break; \
+  case ::kj::Decay<decltype(*_kj_switch_subject)>::tagFor<__VA_ARGS__>(): \
+    for (auto& name = _kj_switch_subject->get<__VA_ARGS__>(), *_kj_switch_done = &name; \
+         _kj_switch_done; _kj_switch_done = nullptr)
+#define KJ_CASE_ONEOF_DEFAULT break; default:
+// Allows switching over a OneOf.
+//
+// Example:
+//
+//     kj::OneOf<int, float, const char*> variant;
+//     KJ_SWITCH_ONEOF(variant) {
+//       KJ_CASE_ONEOF(i, int) {
+//         doSomethingWithInt(i);
+//       }
+//       KJ_CASE_ONEOF(s, const char*) {
+//         doSomethingWithString(s);
+//       }
+//       KJ_CASE_ONEOF_DEFAULT {
+//         doSomethingElse();
+//       }
+//     }
+//
+// Notes:
+// - If you don't handle all possible types and don't include a default branch, you'll get a
+//   compiler warning, just like a regular switch() over an enum where one of the enum values is
+//   missing.
+// - There's no need for a `break` statement in a KJ_CASE_ONEOF; it is implied.
+// - Under C++11 and C++14, only one KJ_SWITCH_ONEOF() can appear in a block. Wrap the switch in
+//   a pair of braces if you need a second switch in the same block. If C++17 is enabled, this is
+//   not an issue.
+//
+// Implementation notes:
+// - The use of __VA_ARGS__ is to account for template types that have commas separating type
+//   parameters, since macros don't recognize <> as grouping.
+// - _kj_switch_done is really used as a boolean flag to prevent the for() loop from actually
+//   looping, but it's defined as a pointer since that's all we can define in this context.
+
 }  // namespace kj
 
 #endif  // KJ_ONE_OF_H_