1 files changed, 89 insertions, 52 deletions
diff --git a/gcc/rust/util/rust-optional.h b/gcc/rust/util/rust-optional.h
index c1b547a..5646540 100644
--- a/gcc/rust/util/rust-optional.h
+++ b/gcc/rust/util/rust-optional.h
@@ -96,6 +96,7 @@ private:
 
 public:
   Optional (const Optional &other) = default;
+  Optional &operator= (const Optional &other) = default;
   Optional (Optional &&other) = default;
 
   static Optional<T> some (T value)
@@ -167,74 +168,110 @@ public:
   }
 };
 
-} // namespace Rust
+template <typename T> class Optional<T &>
+{
+private:
+  struct Empty
+  {
+  };
 
-#ifdef CHECKING_P
+  enum Kind
+  {
+    Some,
+    None
+  } kind;
 
-static void
-rust_optional_create ()
-{
-  auto opt = Rust::Optional<int>::some (15);
+  union Content
+  {
+    Empty empty;
+    T *value;
 
-  ASSERT_TRUE (opt.is_some ());
-  ASSERT_EQ (opt.get (), 15);
+    Content () = default;
+  } content;
 
-  Rust::Optional<int> const_opt = Rust::Optional<int>::some (15);
-  const int &value = const_opt.get ();
+  Optional<T &> (Kind kind, Content content) : kind (kind), content (content) {}
 
-  ASSERT_EQ (value, 15);
-}
+public:
+  Optional (const Optional &other) = default;
+  Optional (Optional &&other) = default;
 
-static void
-rust_optional_operators ()
-{
-  auto opt = Rust::Optional<int>::some (15);
+  static Optional<T &> some (T &value)
+  {
+    Content content;
+    content.value = &value;
+
+    return Optional (Kind::Some, content);
+  }
+
+  static Optional<T &> none ()
+  {
+    Content content;
+    content.empty = Empty ();
+
+    return Optional (Kind::None, content);
+  }
+
+  bool is_some () const { return kind == Kind::Some; }
+  bool is_none () const { return !is_some (); }
+
+  // FIXME: Can we factor this in a single class?
+
+  /**
+   * Enable boolean-like comparisons.
+   */
+  operator bool () { return is_some (); }
 
-  // as bool
-  ASSERT_TRUE (opt);
+  /**
+   * Enables dereferencing to access the contained value
+   */
+  T &operator* () { return get (); }
+  const T &operator* () const { return get (); }
+  T *operator-> () { return &get (); }
+  const T *operator-> () const { return &get (); }
 
-  // deref
-  ASSERT_EQ (*opt, 15);
+  const T &get () const
+  {
+    rust_assert (is_some ());
+
+    return *content.value;
+  }
 
-  class Methodable
+  T &get ()
   {
-  public:
-    int method () { return 15; }
-  };
+    rust_assert (is_some ());
 
-  auto m_opt = Rust::Optional<Methodable>::some (Methodable ());
-  ASSERT_EQ (m_opt->method (), 15);
-}
+    return *content.value;
+  }
 
-static void
-rust_optional_take ()
-{
-  auto opt = Rust::Optional<int>::some (15);
-  auto value = opt.take ();
+  T &take ()
+  {
+    rust_assert (is_some ());
 
-  ASSERT_EQ (value, 15);
-  ASSERT_TRUE (opt.is_none ());
-}
+    auto to_return = std::move (content.value);
 
-static void
-rust_optional_map ()
-{
-  auto opt = Rust::Optional<int>::some (15);
-  auto twice = opt.map<int> ([] (int value) { return value * 2; });
+    content.empty = Empty ();
+    kind = Kind::None;
+
+    return *to_return;
+  }
 
-  ASSERT_FALSE (opt);
-  ASSERT_TRUE (twice);
-  ASSERT_EQ (*twice, 30);
-}
+  template <typename U> Optional<U &> map (std::function<U &(T &)> functor)
+  {
+    if (is_none ())
+      return Optional::none ();
 
-static void
-rust_optional_test ()
-{
-  rust_optional_create ();
-  rust_optional_operators ();
-  rust_optional_take ();
-  rust_optional_map ();
-}
+    auto value = functor (take ());
+
+    return Optional::some (value);
+  }
+};
+
+} // namespace Rust
+
+#ifdef CHECKING_P
+
+void
+rust_optional_test ();
 
 #endif // !CHECKING_P