union support for NDK backend

Generated class for AIDL union for NDK backend is similar to the one for
C++ backend. Differences are
- not inheriting Parcelable, but has similar read/write methods
- not having comparison operators.

Bug: 150948558
Bug: 170784707
Test: aidl_unittests / aidl_integration_test
Test: CtsNdkBinderTestCases
Change-Id: Ib7f9265d7664ae98b9edd4a40d50351180c9a4d2
diff --git a/aidl_to_cpp_common.cpp b/aidl_to_cpp_common.cpp
index 04e805f..fac699a 100644
--- a/aidl_to_cpp_common.cpp
+++ b/aidl_to_cpp_common.cpp
@@ -380,7 +380,7 @@
   return code.str();
 }
 
-std::string TemplateDecl(const AidlStructuredParcelable& defined_type) {
+std::string TemplateDecl(const AidlParcelable& defined_type) {
   std::string decl = "";
   if (defined_type.IsGeneric()) {
     std::vector<std::string> template_params;
diff --git a/aidl_to_cpp_common.h b/aidl_to_cpp_common.h
index f49edb0..3c6996e 100644
--- a/aidl_to_cpp_common.h
+++ b/aidl_to_cpp_common.h
@@ -73,7 +73,7 @@
 
 std::string GenerateEnumValues(const AidlEnumDeclaration& enum_decl,
                                const std::vector<std::string>& enclosing_namespaces_of_enum_decl);
-std::string TemplateDecl(const AidlStructuredParcelable& defined_type);
+std::string TemplateDecl(const AidlParcelable& defined_type);
 
 void GenerateParcelableComparisonOperators(CodeWriter& out, const AidlParcelable& parcelable);
 
diff --git a/aidl_unittest.cpp b/aidl_unittest.cpp
index 3dd4e62..9dc2af0 100644
--- a/aidl_unittest.cpp
+++ b/aidl_unittest.cpp
@@ -2831,6 +2831,132 @@
 }  // namespace a
 )";
 
+const char kUnionExampleExpectedOutputNdkHeader[] = R"(#pragma once
+#include <android/binder_interface_utils.h>
+#include <android/binder_parcelable_utils.h>
+
+#include <type_traits>
+#include <utility>
+#include <variant>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+#ifdef BINDER_STABILITY_SUPPORT
+#include <android/binder_stability.h>
+#endif  // BINDER_STABILITY_SUPPORT
+#include <aidl/a/ByteEnum.h>
+namespace aidl {
+namespace a {
+class Foo {
+public:
+  typedef std::false_type fixed_size;
+  static const char* descriptor;
+
+  enum Tag : int32_t {
+    ns = 0,  // int[] ns;
+    e,  // a.ByteEnum e;
+  };
+
+  template<typename _Tp>
+  static constexpr bool _not_self = !std::is_same_v<std::remove_cv_t<std::remove_reference_t<_Tp>>, Foo>;
+
+  Foo() : _value(std::in_place_index<ns>, std::vector<int32_t>({42})) { }
+  Foo(const Foo&) = default;
+  Foo(Foo&&) = default;
+  Foo& operator=(const Foo&) = default;
+  Foo& operator=(Foo&&) = default;
+
+  template <typename _Tp, std::enable_if_t<_not_self<_Tp>, int> = 0>
+  constexpr Foo(_Tp&& _arg)
+      : _value(std::forward<_Tp>(_arg)) {}
+
+  template <typename... _Tp>
+  constexpr explicit Foo(_Tp&&... _args)
+      : _value(std::forward<_Tp>(_args)...) {}
+
+  template <Tag _tag, typename... _Tp>
+  static Foo make(_Tp&&... _args) {
+    return Foo(std::in_place_index<_tag>, std::forward<_Tp>(_args)...);
+  }
+
+  template <Tag _tag, typename _Tp, typename... _Up>
+  static Foo make(std::initializer_list<_Tp> _il, _Up&&... _args) {
+    return Foo(std::in_place_index<_tag>, std::move(_il), std::forward<_Up>(_args)...);
+  }
+
+  Tag getTag() const {
+    return static_cast<Tag>(_value.index());
+  }
+
+  template <Tag _tag>
+  const auto& get() const {
+    if (getTag() != _tag) { abort(); }
+    return std::get<_tag>(_value);
+  }
+
+  template <Tag _tag>
+  auto& get() {
+    if (getTag() != _tag) { abort(); }
+    return std::get<_tag>(_value);
+  }
+
+  template <Tag _tag, typename... _Tp>
+  void set(_Tp&&... _args) {
+    _value.emplace<_tag>(std::forward<_Tp>(_args)...);
+  }
+
+  binder_status_t readFromParcel(const AParcel* _parcel);
+  binder_status_t writeToParcel(AParcel* _parcel) const;
+  static const ::ndk::parcelable_stability_t _aidl_stability = ::ndk::STABILITY_LOCAL;
+private:
+  std::variant<std::vector<int32_t>, ::aidl::a::ByteEnum> _value;
+};
+}  // namespace a
+}  // namespace aidl
+)";
+
+const char kUnionExampleExpectedOutputNdkSource[] = R"(#include "aidl/a/Foo.h"
+
+#include <android/binder_parcel_utils.h>
+
+namespace aidl {
+namespace a {
+const char* Foo::descriptor = "a.Foo";
+
+binder_status_t Foo::readFromParcel(const AParcel* _parcel) {
+  binder_status_t _aidl_ret_status;
+  int32_t _aidl_tag;
+  if ((_aidl_ret_status = AParcel_readInt32(_parcel, &_aidl_tag)) != STATUS_OK) return _aidl_ret_status;
+  switch (_aidl_tag) {
+  case ns: {
+    std::vector<int32_t> _aidl_value;
+    if ((_aidl_ret_status = ::ndk::AParcel_readVector(_parcel, &_aidl_value)) != STATUS_OK) return _aidl_ret_status;
+    set<ns>(std::move(_aidl_value));
+    return STATUS_OK; }
+  case e: {
+    ::aidl::a::ByteEnum _aidl_value;
+    if ((_aidl_ret_status = AParcel_readByte(_parcel, reinterpret_cast<int8_t*>(&_aidl_value))) != STATUS_OK) return _aidl_ret_status;
+    set<e>(std::move(_aidl_value));
+    return STATUS_OK; }
+  }
+  return STATUS_BAD_VALUE;
+}
+binder_status_t Foo::writeToParcel(AParcel* _parcel) const {
+  binder_status_t _aidl_ret_status = AParcel_writeInt32(_parcel, getTag());
+  if (_aidl_ret_status != STATUS_OK) return _aidl_ret_status;
+  switch (getTag()) {
+  case ns: return ::ndk::AParcel_writeVector(_parcel, get<ns>());
+  case e: return AParcel_writeByte(_parcel, static_cast<int8_t>(get<e>()));
+  }
+  abort();
+}
+
+}  // namespace a
+}  // namespace aidl
+)";
+
 const char kUnionExampleExpectedOutputJava[] = R"(/*
  * This file is auto-generated.  DO NOT MODIFY.
  */
@@ -3007,11 +3133,17 @@
                            {"out/a/Foo.h", kUnionExampleExpectedOutputCppHeader}}));
 }
 
+TEST_F(AidlUnionTest, Example_Ndk) {
+  Compile("ndk");
+  EXPECT_COMPILE_OUTPUTS(
+      map<string, string>({{"out/a/Foo.cpp", kUnionExampleExpectedOutputNdkSource},
+                           {"out/aidl/a/Foo.h", kUnionExampleExpectedOutputNdkHeader}}));
+}
+
 TEST_F(AidlUnionTest, Example_Java) {
   Compile("java");
   EXPECT_COMPILE_OUTPUTS(
       map<string, string>({{"out/a/Foo.java", kUnionExampleExpectedOutputJava}}));
-  // TODO(b/170784707) NDK
   // TODO(b/170689477) Rust
 }
 
diff --git a/generate_ndk.cpp b/generate_ndk.cpp
index 64334a9..6f4d52b 100644
--- a/generate_ndk.cpp
+++ b/generate_ndk.cpp
@@ -37,6 +37,17 @@
 static constexpr const char* kCachedHashMutex = "_aidl_cached_hash_mutex";
 
 using namespace internals;
+namespace internals {
+void GenerateParcelHeader(CodeWriter& out, const AidlTypenames& types,
+                          const AidlStructuredParcelable& defined_type, const Options& options);
+void GenerateParcelSource(CodeWriter& out, const AidlTypenames& types,
+                          const AidlStructuredParcelable& defined_type, const Options& options);
+void GenerateParcelHeader(CodeWriter& out, const AidlTypenames& types,
+                          const AidlUnionDecl& defined_type, const Options& options);
+void GenerateParcelSource(CodeWriter& out, const AidlTypenames& types,
+                          const AidlUnionDecl& defined_type, const Options& options);
+}  // namespace internals
+
 using cpp::ClassNames;
 
 void GenerateNdkInterface(const string& output_file, const Options& options,
@@ -64,8 +75,9 @@
   AIDL_FATAL_IF(!source_writer->Close(), output_file);
 }
 
+template <typename ParcelableType>
 void GenerateNdkParcel(const string& output_file, const Options& options,
-                       const AidlTypenames& types, const AidlStructuredParcelable& defined_type,
+                       const AidlTypenames& types, const ParcelableType& defined_type,
                        const IoDelegate& io_delegate) {
   const string header_path =
       options.OutputHeaderDir() + NdkHeaderFile(defined_type, ClassNames::RAW);
@@ -133,7 +145,13 @@
                  const AidlDefinedType& defined_type, const IoDelegate& io_delegate) {
   if (const AidlStructuredParcelable* parcelable = defined_type.AsStructuredParcelable();
       parcelable != nullptr) {
-    GenerateNdkParcel(output_file, options, types, *parcelable, io_delegate);
+    GenerateNdkParcel<AidlStructuredParcelable>(output_file, options, types, *parcelable,
+                                                io_delegate);
+    return;
+  }
+
+  if (const AidlUnionDecl* union_decl = defined_type.AsUnionDeclaration(); union_decl != nullptr) {
+    GenerateNdkParcel<AidlUnionDecl>(output_file, options, types, *union_decl, io_delegate);
     return;
   }
 
@@ -199,6 +217,8 @@
       return NdkHeaderFile(*type, ClassNames::RAW, false /*use_os_sep*/);
     } else if (type->AsStructuredParcelable() != nullptr) {
       return NdkHeaderFile(*type, ClassNames::RAW, false /*use_os_sep*/);
+    } else if (type->AsUnionDeclaration() != nullptr) {
+      return NdkHeaderFile(*type, ClassNames::RAW, false /*use_os_sep*/);
     } else if (type->AsParcelable() != nullptr) {
       return type->AsParcelable()->GetCppHeader();
     } else if (type->AsEnumDeclaration() != nullptr) {
@@ -231,9 +251,8 @@
     }
   }
 
-  const AidlStructuredParcelable* parcelable = defined_type.AsStructuredParcelable();
-  if (parcelable != nullptr) {
-    for (const auto& field : parcelable->GetFields()) {
+  auto visit_parcelable = [&](const auto& parcelable) {
+    for (const auto& field : parcelable.GetFields()) {
       visit(field->GetType());
       // Check the fields for generic type arguments
       if (field->GetType().IsGeneric()) {
@@ -242,6 +261,16 @@
         }
       }
     }
+  };
+
+  const AidlStructuredParcelable* parcelable = defined_type.AsStructuredParcelable();
+  if (parcelable != nullptr) {
+    visit_parcelable(*parcelable);
+  }
+
+  const AidlUnionDecl* union_decl = defined_type.AsUnionDeclaration();
+  if (union_decl != nullptr) {
+    visit_parcelable(*union_decl);
   }
 
   const AidlEnumDeclaration* enum_decl = defined_type.AsEnumDeclaration();
@@ -1067,6 +1096,110 @@
   LeaveNdkNamespace(out, defined_type);
 }
 
+void GenerateParcelHeader(CodeWriter& out, const AidlTypenames& types,
+                          const AidlUnionDecl& defined_type, const Options& /*options*/) {
+  const std::string clazz = ClassName(defined_type, ClassNames::RAW);
+  cpp::UnionWriter uw{defined_type, types,
+                      [&](const AidlTypeSpecifier& type, const AidlTypenames& types) {
+                        return NdkNameOf(types, type, StorageMode::STACK);
+                      },
+                      &ConstantValueDecorator};
+
+  out << "#pragma once\n";
+  out << "#include <android/binder_interface_utils.h>\n";
+  out << "#include <android/binder_parcelable_utils.h>\n";
+  out << "\n";
+
+  for (const auto& header : cpp::UnionWriter::headers) {
+    out << "#include <" << header << ">\n";
+  }
+  GenerateHeaderIncludes(out, types, defined_type);
+
+  EnterNdkNamespace(out, defined_type);
+  out << cpp::TemplateDecl(defined_type);
+  out << "class " << clazz << " {\n";
+  out << "public:\n";
+  out.Indent();
+  if (defined_type.IsFixedSize()) {
+    out << "typedef std::true_type fixed_size;\n";
+  } else {
+    out << "typedef std::false_type fixed_size;\n";
+  }
+  out << "static const char* descriptor;\n";
+  out << "\n";
+  uw.PublicFields(out);
+
+  out << "binder_status_t readFromParcel(const AParcel* _parcel);\n";
+  out << "binder_status_t writeToParcel(AParcel* _parcel) const;\n";
+
+  out << "static const ::ndk::parcelable_stability_t _aidl_stability = ::ndk::"
+      << (defined_type.IsVintfStability() ? "STABILITY_VINTF" : "STABILITY_LOCAL") << ";\n";
+  out.Dedent();
+  out << "private:\n";
+  out.Indent();
+  uw.PrivateFields(out);
+  out.Dedent();
+  out << "};\n";
+  LeaveNdkNamespace(out, defined_type);
+}
+void GenerateParcelSource(CodeWriter& out, const AidlTypenames& types,
+                          const AidlUnionDecl& defined_type, const Options& /*options*/) {
+  std::string clazz = ClassName(defined_type, ClassNames::RAW);
+  if (defined_type.IsGeneric()) {
+    std::vector<std::string> template_params;
+    for (const auto& parameter : defined_type.GetTypeParameters()) {
+      template_params.push_back(parameter);
+    }
+    clazz += base::StringPrintf("<%s>", base::Join(template_params, ", ").c_str());
+  }
+
+  cpp::UnionWriter uw{defined_type, types,
+                      [&](const AidlTypeSpecifier& type, const AidlTypenames& types) {
+                        return NdkNameOf(types, type, StorageMode::STACK);
+                      },
+                      &ConstantValueDecorator};
+  cpp::ParcelWriterContext ctx{
+      .status_type = "binder_status_t",
+      .status_ok = "STATUS_OK",
+      .status_bad = "STATUS_BAD_VALUE",
+      .read_func =
+          [&](CodeWriter& out, const std::string& var, const AidlTypeSpecifier& type) {
+            ReadFromParcelFor({out, types, type, "_parcel", "&" + var});
+          },
+      .write_func =
+          [&](CodeWriter& out, const std::string& value, const AidlTypeSpecifier& type) {
+            WriteToParcelFor({out, types, type, "_parcel", value});
+          },
+  };
+
+  out << "#include \"" << NdkHeaderFile(defined_type, ClassNames::RAW, false /*use_os_sep*/)
+      << "\"\n";
+  out << "\n";
+  GenerateSourceIncludes(out, types, defined_type);
+  out << "\n";
+  EnterNdkNamespace(out, defined_type);
+  out << cpp::TemplateDecl(defined_type);
+  out << "const char* " << clazz << "::" << kDescriptor << " = \""
+      << defined_type.GetCanonicalName() << "\";\n";
+  out << "\n";
+
+  out << cpp::TemplateDecl(defined_type);
+  out << "binder_status_t " << clazz << "::readFromParcel(const AParcel* _parcel) {\n";
+  out.Indent();
+  uw.ReadFromParcel(out, ctx);
+  out.Dedent();
+  out << "}\n";
+
+  out << cpp::TemplateDecl(defined_type);
+  out << "binder_status_t " << clazz << "::writeToParcel(AParcel* _parcel) const {\n";
+  out.Indent();
+  uw.WriteToParcel(out, ctx);
+  out.Dedent();
+  out << "}\n";
+  out << "\n";
+  LeaveNdkNamespace(out, defined_type);
+}
+
 std::string GenerateEnumToString(const AidlTypenames& typenames,
                                  const AidlEnumDeclaration& enum_decl) {
   std::ostringstream code;
diff --git a/generate_ndk.h b/generate_ndk.h
index 94a7669..6b893a3 100644
--- a/generate_ndk.h
+++ b/generate_ndk.h
@@ -47,11 +47,6 @@
 void GenerateInterfaceHeader(CodeWriter& out, const AidlTypenames& types,
                              const AidlInterface& defined_type, const Options& options);
 
-void GenerateParcelHeader(CodeWriter& out, const AidlTypenames& types,
-                          const AidlStructuredParcelable& defined_type, const Options& options);
-void GenerateParcelSource(CodeWriter& out, const AidlTypenames& types,
-                          const AidlStructuredParcelable& defined_type, const Options& options);
-
 void GenerateEnumHeader(CodeWriter& out, const AidlTypenames& types,
                         const AidlEnumDeclaration& enum_decl, const Options& options);