Refactor code into tidier classes

New changes:

* RecordFieldOrganizer is an interface that
  can be used to manipulate a RecordDecl's
  field order.
* RecordFieldOrganizer is a "friend" to the
  DeclContext (which maintains the linked list)
* The Randstruct class is a subclass that performs
  cache friendly randomization

Author: connorkuehl

clang/include/clang/AST/Decl.h

@@ -65,6 +65,7 @@ class Module;
 class NamespaceDecl;
 class ParmVarDecl;
 class RecordDecl;
+class RecordFieldReorganizer;
 class Stmt;
 class StringLiteral;
 class TagDecl;

clang/include/clang/AST/DeclBase.h

@@ -1270,6 +1270,8 @@ class DeclContext {
   friend class ExternalASTSource;
   /// For CreateStoredDeclsMap
   friend class DependentDiagnostic;
+  /// For fine-grained control of field order
+  friend class RecordFieldReorganizer;
   /// For hasNeedToReconcileExternalVisibleStorage,
   /// hasLazyLocalLexicalLookups, hasLazyExternalLexicalLookups
   friend class ASTWriter;

clang/include/clang/AST/RecordFieldReorganizer.h

@@ -0,0 +1,47 @@
+//===-- RecordFieldReorganizer.h - Interface for manipulating field order --*-
+// C++
+//-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file contains the base class that defines an interface for
+// manipulating a RecordDecl's field layouts.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_LIB_AST_RECORDFIELDREORGANIZER_H
+#define LLVM_CLANG_LIB_AST_RECORDFIELDREORGANIZER_H
+
+#include "Decl.h"
+
+namespace clang {
+
+// FIXME: Find a better alternative to SmallVector with hardcoded size!
+
+class RecordFieldReorganizer {
+public:
+  virtual ~RecordFieldReorganizer() = default;
+  void reorganizeFields(const ASTContext &C, const RecordDecl *D) const;
+
+protected:
+  virtual void reorganize(const ASTContext &C, const RecordDecl *D,
+                          SmallVector<Decl *, 64> &NewOrder) const = 0;
+
+private:
+  void commit(const RecordDecl *D,
+              SmallVectorImpl<Decl *> &NewFieldOrder) const;
+};
+
+class Randstruct : public RecordFieldReorganizer {
+protected:
+  virtual void reorganize(const ASTContext &C, const RecordDecl *D,
+                          SmallVector<Decl *, 64> &NewOrder) const override;
+};
+
+} // namespace clang
+
+#endif

clang/lib/AST/CMakeLists.txt

@@ -44,7 +44,7 @@ add_clang_library(clangAST
   InheritViz.cpp
   ItaniumCXXABI.cpp
   ItaniumMangle.cpp
-  LayoutFieldRandomizer.cpp
+  RecordFieldReorganizer.cpp
   Mangle.cpp
   MicrosoftCXXABI.cpp
   MicrosoftMangle.cpp

clang/lib/AST/LayoutFieldRandomizer.h

@@ -1,4 +1,4 @@
-//===----- LayoutFieldRandomizer.cpp - Randstruct Implementation -*- C++
+//===----- RecordFieldReorganizer.cpp - Implementation for field reorder -*- C++
 //-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
@@ -7,20 +7,59 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// Cache line best-effort field randomization
+// Contains the implementation for RecordDecl field reordering.
 //
 //===----------------------------------------------------------------------===//
 
-#include "LayoutFieldRandomizer.h"
-#include "llvm/ADT/SmallVector.h"
+#include "clang/AST/RecordFieldReorganizer.h"
+#include "clang/AST/ASTContext.h"
 
 #include <algorithm>
 #include <cstdint>
 #include <random>
+#include <set>
 #include <vector>
 
+// FIXME: Find a better alternative to SmallVector with hardcoded size!
+
 namespace clang {
 
+void RecordFieldReorganizer::reorganizeFields(const ASTContext &C,
+                                              const RecordDecl *D) const {
+  // Save original fields for asserting later that a subclass hasn't
+  // sabotaged the RecordDecl by removing or adding fields
+  std::set<Decl *> mutateGuard;
+
+  SmallVector<Decl *, 64> fields;
+  for (auto f : D->fields()) {
+    mutateGuard.insert(f);
+    fields.push_back(f);
+  }
+
+  // Now allow subclass implementations to reorder the fields
+  reorganize(C, D, fields);
+
+  // Assert all fields are still present
+  assert(mutateGuard.size() == fields.size() &&
+         "Field count altered after reorganization");
+  for (auto f : fields) {
+    auto found = std::find(std::begin(mutateGuard), std::end(mutateGuard), f);
+    assert(found != std::end(mutateGuard) &&
+           "Unknown field encountered after reorganization");
+  }
+
+  commit(D, fields);
+}
+
+void RecordFieldReorganizer::commit(
+    const RecordDecl *D, SmallVectorImpl<Decl *> &NewFieldOrder) const {
+  Decl *First, *Last;
+  std::tie(First, Last) = DeclContext::BuildDeclChain(
+      NewFieldOrder, D->hasLoadedFieldsFromExternalStorage());
+  D->FirstDecl = First;
+  D->LastDecl = Last;
+}
+
 /// Bucket to store fields up to size of a cache line during randomization.
 class Bucket {
 public:
@@ -51,11 +90,11 @@ class BitfieldRun : public Bucket {
   virtual bool isBitfieldRun() const override;
 };
 
-// TODO: Is there a way to detect this? (i.e. on 32bit system vs 64?)
+// FIXME: Is there a way to detect this? (i.e. on 32bit system vs 64?)
 const size_t CACHE_LINE = 64;
 
 SmallVector<FieldDecl *, 64> Bucket::randomize() {
-  // TODO use seed
+  // FIXME use seed
   auto rng = std::default_random_engine{};
   std::shuffle(std::begin(fields), std::end(fields), rng);
   return fields;
@@ -100,10 +139,7 @@ bool BitfieldRun::canFit(size_t size) const {
   return true;
 }
 
-bool BitfieldRun::isBitfieldRun() const {
-  // Yes.
-  return true;
-}
+bool BitfieldRun::isBitfieldRun() const { return true; }
 
 SmallVector<Decl *, 64> randomize(SmallVector<Decl *, 64> fields) {
   auto rng = std::default_random_engine{};
@@ -161,7 +197,6 @@ SmallVector<Decl *, 64> perfrandomize(const ASTContext &ctx,
         currentBucket = llvm::make_unique<Bucket>();
       }
 
-      // FIXME get access to AST Context
       auto width = ctx.getTypeInfo(f->getType()).Width;
 
       // If we can fit, add it.
@@ -209,9 +244,10 @@ SmallVector<Decl *, 64> perfrandomize(const ASTContext &ctx,
   return finalOrder;
 }
 
-SmallVector<Decl *, 64> rearrange(const ASTContext &ctx,
-                                  SmallVector<Decl *, 64> fields) {
-  return perfrandomize(ctx, fields);
+void Randstruct::reorganize(const ASTContext &C, const RecordDecl *D,
+                            SmallVector<Decl *, 64> &NewOrder) const {
+  SmallVector<Decl *, 64> randomized = perfrandomize(C, NewOrder);
+  NewOrder = randomized;
 }
 
 } // namespace clang

clang/lib/AST/LayoutFieldRandomizer.cpp renamed to clang/lib/AST/RecordFieldReorganizer.cpp

@@ -6,7 +6,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "LayoutFieldRandomizer.h"
 #include "clang/AST/RecordLayout.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTDiagnostic.h"
@@ -16,6 +15,7 @@
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/Expr.h"
+#include "clang/AST/RecordFieldReorganizer.h"
 #include "clang/Basic/TargetInfo.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/Format.h"
@@ -2989,6 +2989,8 @@ ASTContext::getASTRecordLayout(const RecordDecl *D) const {
 
   bool ShouldBeRandomized = D->getAttr<RandomizeLayoutAttr>() != nullptr;
   if (ShouldBeRandomized) {
+	  Randstruct randstruct;
+	  randstruct.reorganizeFields(*this, D);
   }
 
   if (isMsLayout(*this)) {