From c207972af64cc1f4247c0b2f75715307c1ff2209 Mon Sep 17 00:00:00 2001
From: Alastair Robertson <ajor@meta.com>
Date: Thu, 2 Nov 2023 06:22:57 -0700
Subject: [PATCH] TypeGraph: Calculate alignment before identifying containers

Not all containers have 8-byte alignment, so if we want to avoid lots of
manual logic for calculating container alignment on a case-by-case
basis, we must calculate alignment from the member variables before the
Class nodes have been replaced by Container nodes.
---
 oi/CodeGen.cpp        |  6 ++----
 oi/OITraceCode.cpp    |  5 +++++
 oi/type_graph/Types.h | 15 +--------------
 test/test_codegen.cpp | 12 ++++++------
 4 files changed, 14 insertions(+), 24 deletions(-)

diff --git a/oi/CodeGen.cpp b/oi/CodeGen.cpp
index 2d0c0eb..3fa0374 100644
--- a/oi/CodeGen.cpp
+++ b/oi/CodeGen.cpp
@@ -1168,6 +1168,7 @@ void CodeGen::transform(TypeGraph& typeGraph) {
   // Simplify the type graph first so there is less work for later passes
   pm.addPass(RemoveTopLevelPointer::createPass());
   pm.addPass(Flattener::createPass());
+  pm.addPass(AlignmentCalc::createPass());
   pm.addPass(IdentifyContainers::createPass(containerInfos_));
   pm.addPass(TypeIdentifier::createPass(config_.passThroughTypes));
   if (config_.features[Feature::PruneTypeGraph])
@@ -1183,16 +1184,13 @@ void CodeGen::transform(TypeGraph& typeGraph) {
 
     // Re-run passes over newly added children
     pm.addPass(Flattener::createPass());
+    pm.addPass(AlignmentCalc::createPass());
     pm.addPass(IdentifyContainers::createPass(containerInfos_));
     pm.addPass(TypeIdentifier::createPass(config_.passThroughTypes));
     if (config_.features[Feature::PruneTypeGraph])
       pm.addPass(Prune::createPass());
   }
 
-  // Calculate alignment before removing members, as those members may have an
-  // influence on the class' overall alignment.
-  pm.addPass(AlignmentCalc::createPass());
-
   pm.addPass(RemoveMembers::createPass(config_.membersToStub));
   if (!config_.features[Feature::TreeBuilderV2])
     pm.addPass(EnforceCompatibility::createPass());
diff --git a/oi/OITraceCode.cpp b/oi/OITraceCode.cpp
index 47e98d7..fa0beda 100644
--- a/oi/OITraceCode.cpp
+++ b/oi/OITraceCode.cpp
@@ -101,6 +101,8 @@ struct alignas(align) DummySizedOperator {
 // container if an empty class is passed.
 template <int32_t Id>
 struct DummySizedOperator<0, 0, Id> {};
+template <int32_t Id>
+struct DummySizedOperator<0, 1, Id> {};
 
 template <template <typename, size_t, size_t, int32_t> typename DerivedT,
           typename T,
@@ -130,6 +132,9 @@ struct alignas(Align) DummyAllocator
 template <typename T, int32_t Id>
 struct DummyAllocator<T, 0, 0, Id>
     : DummyAllocatorBase<DummyAllocator, T, 0, 0, Id> {};
+template <typename T, int32_t Id>
+struct DummyAllocator<T, 0, 1, Id>
+    : DummyAllocatorBase<DummyAllocator, T, 0, 1, Id> {};
 
 template <typename Type, size_t ExpectedSize, size_t ActualSize = 0>
 struct validate_size {
diff --git a/oi/type_graph/Types.h b/oi/type_graph/Types.h
index 03f8f81..f503886 100644
--- a/oi/type_graph/Types.h
+++ b/oi/type_graph/Types.h
@@ -421,20 +421,7 @@ class Container : public Type {
   }
 
   virtual uint64_t align() const override {
-    // XXX Big hack!
-    // We can not assume that all containers have 8-byte alignment. However, to
-    // properly calculate alignment would require significant refactoring of
-    // the container identification code:
-    //  - DrgnParser needs to return Class objects instead of Containers
-    //  - They must be flattened
-    //  - We must calculate alignment based on their members
-    //  - Then we are able to convert them into Container objects
-    //
-    // As a quick, hopefully temporary solution, hardcode some known container
-    // alignments here
-    if (containerInfo_.ctype == THRIFT_ISSET_TYPE)
-      return 1;
-    return 8;
+    return align_;
   }
 
   virtual NodeId id() const override {
diff --git a/test/test_codegen.cpp b/test/test_codegen.cpp
index 94b6dd0..87ec437 100644
--- a/test/test_codegen.cpp
+++ b/test/test_codegen.cpp
@@ -62,11 +62,11 @@ TEST(CodeGenTest, TransformContainerAllocator) {
           Function: deallocate
 )",
                 R"(
-[2] Container: std::vector<int32_t, DummyAllocator<int32_t, 8, 0, 3>> (size: 24)
+[2] Container: std::vector<int32_t, DummyAllocator<int32_t, 8, 1, 3>> (size: 24)
       Param
         Primitive: int32_t
       Param
-[3]     DummyAllocator [MyAlloc] (size: 8)
+[3]     DummyAllocator [MyAlloc] (size: 8, align: 1)
           Primitive: int32_t
 )");
 }
@@ -95,13 +95,13 @@ TEST(CodeGenTest, TransformContainerAllocatorParamInParent) {
           Function: deallocate
 )",
                 R"(
-[4] Container: std::map<int32_t, int32_t, DummyAllocator<std::pair<int32_t const, int32_t>, 0, 0, 6>> (size: 24)
+[4] Container: std::map<int32_t, int32_t, DummyAllocator<std::pair<int32_t const, int32_t>, 0, 1, 6>> (size: 24)
       Param
         Primitive: int32_t
       Param
         Primitive: int32_t
       Param
-[6]     DummyAllocator [MyAlloc<std::pair<const int, int>>] (size: 0)
+[6]     DummyAllocator [MyAlloc<std::pair<const int, int>>] (size: 0, align: 1)
 [5]       Container: std::pair<int32_t const, int32_t> (size: 8)
             Param
               Primitive: int32_t
@@ -167,11 +167,11 @@ TEST(CodeGenTest, ReplaceContainersAndDummies) {
           Function: deallocate
 )",
                 R"(
-[2] Container: std::vector<uint32_t, DummyAllocator<uint32_t, 0, 0, 3>> (size: 24)
+[2] Container: std::vector<uint32_t, DummyAllocator<uint32_t, 0, 1, 3>> (size: 24)
       Param
         Primitive: uint32_t
       Param
-[3]     DummyAllocator [allocator<int>] (size: 0)
+[3]     DummyAllocator [allocator<int>] (size: 0, align: 1)
           Primitive: uint32_t
 )");
 }