[LightSan] Eliminate indirection when using the Small and Large Objects

compiler-rt/lib/objsan/include/obj_encoding.h

@@ -457,19 +457,17 @@ using SmallObjectsTy = BucketSchemeTy</*EncodingNo=*/1,
                                       /*RealPtrBits=*/32>;
 using LargeObjectsTy = LedgerSchemeTy</*EncodingNo=*/2, /*ObjectBits=*/24>;
 
-extern SmallObjectsTy *SmallObjects;
-extern LargeObjectsTy *LargeObjects;
+extern SmallObjectsTy &SmallObjects;
+extern LargeObjectsTy &LargeObjects;
 
+SmallObjectsTy &getOrConstructSmallObjects();
 static inline SmallObjectsTy &getSmallObjects() {
-  if (!SmallObjects) [[unlikely]]
-    SmallObjects = new SmallObjectsTy();
-  return *SmallObjects;
+  return SmallObjects;
 }
 
+LargeObjectsTy &getOrConstructLargeObjects();
 static inline LargeObjectsTy &getLargeObjects() {
-  if (!LargeObjects) [[unlikely]]
-    LargeObjects = new LargeObjectsTy();
-  return *LargeObjects;
+  return LargeObjects;
 }
 
 } // namespace __objsan

compiler-rt/lib/objsan/objsan_ir_rt.cpp

@@ -94,6 +94,15 @@ char *__objsan_post_alloca(char *MPtr, int64_t ObjSize,
   return __objsan_register_object(MPtr, ObjSize, RequiresTemporalCheck);
 }
 
+OBJSAN_BIG_API_ATTRS
+void __objsan_pre_module(char *module_name, char *name, int32_t id) {
+  getOrConstructLargeObjects();
+  getOrConstructSmallObjects();
+}
+
+OBJSAN_SMALL_API_ATTRS
+void __objsan_post_module(char *module_name, char *name, int32_t id) {}
+
 // OBJSAN_BIG_API_ATTRS
 __attribute__((optnone, noinline)) char *
 __objsan_pre_global(char *MPtr, int32_t ObjSize, int8_t IsDefinition,

compiler-rt/lib/objsan/objsan_rt.cpp

@@ -1,30 +1,51 @@
 #include "include/obj_encoding.h"
 
-namespace __objsan {
+#include <new>
 
-template <typename T> class ObjectDeallocator {
-  T *&ObjectPtr;
+namespace __objsan {
 
-public:
-  ObjectDeallocator(T *&Ptr) : ObjectPtr(Ptr) {}
-  ~ObjectDeallocator() {
-    if (ObjectPtr == nullptr)
-      delete ObjectPtr;
+template <class T> struct LazyConstruction {
+  static bool Init;
+  LazyConstruction() {}
+  ~LazyConstruction() {
+    if (Init)
+      V.~T();
+  }
+  template <typename... ArgsTy> void construct(ArgsTy &&...Args) {
+    if (!Init) {
+      new (&V) T(std::forward<ArgsTy>(Args)...);
+      Init = true;
+    }
   }
+  union {
+    T V;
+  };
 };
 
-__attribute__((visibility("default"))) SmallObjectsTy *SmallObjects = nullptr;
-__attribute__((visibility("default"))) LargeObjectsTy *LargeObjects = nullptr;
+static LazyConstruction<SmallObjectsTy> LazySmallObjects;
+static LazyConstruction<LargeObjectsTy> LazyLargeObjects;
+template <> bool LazyConstruction<SmallObjectsTy>::Init = false;
+template <> bool LazyConstruction<LargeObjectsTy>::Init = false;
 
-// These will ensure the objects are deallocated when the program ends.
-ObjectDeallocator<SmallObjectsTy> SODeallocator(SmallObjects);
-ObjectDeallocator<LargeObjectsTy> LODeallocator(LargeObjects);
+__attribute__((visibility("default"))) SmallObjectsTy &SmallObjects =
+    LazySmallObjects.V;
+__attribute__((visibility("default"))) LargeObjectsTy &LargeObjects =
+    LazyLargeObjects.V;
+
+SmallObjectsTy &getOrConstructSmallObjects() {
+  LazySmallObjects.construct();
+  return SmallObjects;
+}
+LargeObjectsTy &getOrConstructLargeObjects() {
+  LazyLargeObjects.construct();
+  return LargeObjects;
+}
 
 __attribute((constructor)) void initialize() {
   // Ensure the globals are constructed before the program begins. If it is
   // multithreaded, we do not want multiple threads to initialize the objects.
-  getSmallObjects();
-  getLargeObjects();
+  getOrConstructSmallObjects();
+  getOrConstructLargeObjects();
 }
 
 } // namespace __objsan

llvm/lib/Transforms/IPO/LightSan.cpp

@@ -2690,7 +2690,7 @@ void LightSanInstrumentationConfig::populate(InstrumentorIRBuilderTy &IIRB) {
   ExtendedICmpIO::populate(*this, IIRB);
   ExtendedVAArgIO::populate(*this, IIRB);
   AllocatorCallIO::populate(*this, IIRB);
-  // ModuleIO::populate(*this, IIRB.Ctx);
+  ModuleIO::populate(*this, IIRB.Ctx);
 
   PtrToIntIO::ConfigTy PostP2IIOConfig(/*Enable=*/false);
   PostP2IIOConfig.set(PtrToIntIO::PassPointer);