Make CompilerType safe

When a process gets restarted TypeSystem objects associated with it
may get deleted, and any CompilerType objects holding on to a
reference to that type system are a use-after-free in waiting. Because
of the SBAPI, we don't have tight control over where CompilerTypes go
and when they are used. This is particularly a problem in the Swift
plugin, where the scratch TypeSystem can be restarted while the
process is still running. The Swift plugin has a lock to prevent
abuse, but where there's a lock there can be bugs.

This patch changes CompilerType to store a std::weak_ptr<TypeSystem>.
Most of the std::weak_ptr<TypeSystem>* uglyness is hidden by
introducing a wrapper class CompilerType::WrappedTypeSystem that has a
dyn_cast_or_null() method. The only sites that need to know about the
weak pointer implementation detail are the ones that deal with
creating TypeSystems.

rdar://101505232

Differential Revision: https://reviews.llvm.org/D136650

Author: adrian-prantl

lldb/include/lldb/Core/Module.h

@@ -812,10 +812,10 @@ class Module : public std::enable_shared_from_this<Module>,
 
   bool GetIsDynamicLinkEditor();
 
-  llvm::Expected<TypeSystem &>
+  llvm::Expected<lldb::TypeSystemSP>
   GetTypeSystemForLanguage(lldb::LanguageType language);
 
-  void ForEachTypeSystem(llvm::function_ref<bool(TypeSystem *)> callback);
+  void ForEachTypeSystem(llvm::function_ref<bool(lldb::TypeSystemSP)> callback);
 
   // Special error functions that can do printf style formatting that will
   // prepend the message with something appropriate for this module (like the

lldb/include/lldb/Symbol/CompilerType.h

@@ -15,10 +15,12 @@
 
 #include "lldb/lldb-private.h"
 #include "llvm/ADT/APSInt.h"
+#include "llvm/Support/Casting.h"
 
 namespace lldb_private {
 
 class DataExtractor;
+class TypeSystem;
 
 /// Generic representation of a type in a programming language.
 ///
@@ -38,38 +40,87 @@ class CompilerType {
   /// implementation.
   ///
   /// \see lldb_private::TypeSystemClang::GetType(clang::QualType)
-  CompilerType(TypeSystem *type_system, lldb::opaque_compiler_type_t type)
-      : m_type(type), m_type_system(type_system) {
+  CompilerType(lldb::TypeSystemWP type_system,
+               lldb::opaque_compiler_type_t type)
+    : m_type_system(type_system), m_type(type) {
+    assert(Verify() && "verification failed");
+  }
+
+  /// This is a minimal wrapper of a TypeSystem shared pointer as
+  /// returned by CompilerType which conventien dyn_cast support.
+  class TypeSystemSPWrapper {
+    lldb::TypeSystemSP m_typesystem_sp;
+
+  public:
+    TypeSystemSPWrapper() = default;
+    TypeSystemSPWrapper(lldb::TypeSystemSP typesystem_sp)
+        : m_typesystem_sp(typesystem_sp) {}
+
+    template <class TypeSystemType> bool isa_and_nonnull() {
+      if (auto *ts = m_typesystem_sp.get())
+        return llvm::isa<TypeSystemType>(ts);
+      return false;
+    }
+
+    /// Return a shared_ptr<TypeSystemType> if dyn_cast succeeds.
+    template <class TypeSystemType>
+    std::shared_ptr<TypeSystemType> dyn_cast_or_null() {
+      if (isa_and_nonnull<TypeSystemType>())
+        return std::shared_ptr<TypeSystemType>(
+            m_typesystem_sp, llvm::cast<TypeSystemType>(m_typesystem_sp.get()));
+      return nullptr;
+    }
+
+    explicit operator bool() const {
+      return static_cast<bool>(m_typesystem_sp);
+    }
+    bool operator==(const TypeSystemSPWrapper &other) const;
+    bool operator!=(const TypeSystemSPWrapper &other) const {
+      return !(*this == other);
+    }
+
+    /// Only to be used in a one-off situations like
+    ///    if (typesystem && typesystem->method())
+    /// Do not store this pointer!
+    TypeSystem *operator->() const;
+
+    lldb::TypeSystemSP GetSharedPointer() const { return m_typesystem_sp; }
+  };
+
+  CompilerType(TypeSystemSPWrapper type_system, lldb::opaque_compiler_type_t type)
+    : m_type_system(type_system.GetSharedPointer()), m_type(type) {
     assert(Verify() && "verification failed");
   }
 
   CompilerType(const CompilerType &rhs)
-      : m_type(rhs.m_type), m_type_system(rhs.m_type_system) {}
+      : m_type_system(rhs.m_type_system), m_type(rhs.m_type) {}
 
   CompilerType() = default;
 
   /// Operators.
   /// \{
   const CompilerType &operator=(const CompilerType &rhs) {
-    m_type = rhs.m_type;
     m_type_system = rhs.m_type_system;
+    m_type = rhs.m_type;
     return *this;
   }
 
   bool operator<(const CompilerType &rhs) const {
-    if (m_type_system == rhs.m_type_system)
+    auto lts = m_type_system.lock();
+    auto rts = rhs.m_type_system.lock();
+    if (lts.get() == rts.get())
       return m_type < rhs.m_type;
-    return m_type_system < rhs.m_type_system;
+    return lts.get() < rts.get();
   }
   /// \}
 
   /// Tests.
   /// \{
   explicit operator bool() const {
-    return m_type != nullptr && m_type_system != nullptr;
+    return m_type_system.lock() && m_type;
   }
 
-  bool IsValid() const { return m_type != nullptr && m_type_system != nullptr; }
+  bool IsValid() const { return (bool)*this; }
 
   bool IsArrayType(CompilerType *element_type = nullptr,
                    uint64_t *size = nullptr,
@@ -159,7 +210,10 @@ class CompilerType {
 
   /// Accessors.
   /// \{
-  TypeSystem *GetTypeSystem() const { return m_type_system; }
+
+  /// Returns a shared pointer to the type system. The
+  /// TypeSystem::TypeSystemSPWrapper can be compared for equality.
+  TypeSystemSPWrapper GetTypeSystem() const;
 
   ConstString GetTypeName(bool BaseOnly = false) const;
 
@@ -174,7 +228,9 @@ class CompilerType {
 
   lldb::TypeClass GetTypeClass() const;
 
-  void SetCompilerType(TypeSystem *type_system,
+  void SetCompilerType(lldb::TypeSystemWP type_system,
+                       lldb::opaque_compiler_type_t type);
+  void SetCompilerType(TypeSystemSPWrapper type_system,
                        lldb::opaque_compiler_type_t type);
 
   unsigned GetTypeQualifiers() const;
@@ -407,8 +463,8 @@ class CompilerType {
                         size_t data_byte_size, Scalar &value,
                         ExecutionContextScope *exe_scope) const;
   void Clear() {
+    m_type_system = {};
     m_type = nullptr;
-    m_type_system = nullptr;
   }
 
 private:
@@ -419,8 +475,8 @@ class CompilerType {
   bool Verify() const;
 #endif
 
+  lldb::TypeSystemWP m_type_system;
   lldb::opaque_compiler_type_t m_type = nullptr;
-  TypeSystem *m_type_system = nullptr;
 };
 
 bool operator==(const CompilerType &lhs, const CompilerType &rhs);

lldb/include/lldb/Symbol/SymbolFile.h

@@ -310,7 +310,7 @@ class SymbolFile : public PluginInterface {
 
   virtual void PreloadSymbols();
 
-  virtual llvm::Expected<lldb_private::TypeSystem &>
+  virtual llvm::Expected<lldb::TypeSystemSP>
   GetTypeSystemForLanguage(lldb::LanguageType language) = 0;
 
   virtual CompilerDeclContext
@@ -465,7 +465,7 @@ class SymbolFileCommon : public SymbolFile {
   uint32_t GetNumCompileUnits() override;
   lldb::CompUnitSP GetCompileUnitAtIndex(uint32_t idx) override;
 
-  llvm::Expected<lldb_private::TypeSystem &>
+  llvm::Expected<lldb::TypeSystemSP>
   GetTypeSystemForLanguage(lldb::LanguageType language) override;
 
   void Dump(Stream &s) override;

lldb/include/lldb/Symbol/SymbolFileOnDemand.h

@@ -167,7 +167,7 @@ class SymbolFileOnDemand : public lldb_private::SymbolFile {
                 lldb::TypeClass type_mask,
                 lldb_private::TypeList &type_list) override;
 
-  llvm::Expected<lldb_private::TypeSystem &>
+  llvm::Expected<lldb::TypeSystemSP>
   GetTypeSystemForLanguage(lldb::LanguageType language) override;
 
   lldb_private::CompilerDeclContext FindNamespace(

lldb/include/lldb/Symbol/TaggedASTType.h

@@ -20,7 +20,8 @@ template <unsigned int C> class TaggedASTType : public CompilerType {
   TaggedASTType(const CompilerType &compiler_type)
       : CompilerType(compiler_type) {}
 
-  TaggedASTType(lldb::opaque_compiler_type_t type, TypeSystem *type_system)
+  TaggedASTType(lldb::opaque_compiler_type_t type,
+                lldb::TypeSystemWP type_system)
       : CompilerType(type_system, type) {}
 
   TaggedASTType(const TaggedASTType<C> &tw) : CompilerType(tw) {}

lldb/include/lldb/Symbol/Type.h

@@ -297,7 +297,7 @@ class TypeImpl {
 
   CompilerType GetCompilerType(bool prefer_dynamic);
 
-  TypeSystem *GetTypeSystem(bool prefer_dynamic);
+  CompilerType::TypeSystemSPWrapper GetTypeSystem(bool prefer_dynamic);
 
   bool GetDescription(lldb_private::Stream &strm,
                       lldb::DescriptionLevel description_level);

lldb/include/lldb/Symbol/TypeSystem.h

@@ -10,12 +10,12 @@
 #define LLDB_SYMBOL_TYPESYSTEM_H
 
 #include <functional>
-#include <map>
 #include <mutex>
 #include <string>
 
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APSInt.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
@@ -72,7 +72,8 @@ struct LanguageSet {
 /// \see lldb_private::CompilerType
 /// \see lldb_private::CompilerDecl
 /// \see lldb_private::CompilerDeclContext
-class TypeSystem : public PluginInterface {
+class TypeSystem : public PluginInterface,
+                   public std::enable_shared_from_this<TypeSystem> {
 public:
   // Constructors and Destructors
   ~TypeSystem() override;
@@ -86,10 +87,10 @@ class TypeSystem : public PluginInterface {
   static lldb::TypeSystemSP CreateInstance(lldb::LanguageType language,
                                            Target *target);
 
-  // Free up any resources associated with this TypeSystem.  Done before
-  // removing all the TypeSystems from the TypeSystemMap.
+  /// Free up any resources associated with this TypeSystem.  Done before
+  /// removing all the TypeSystems from the TypeSystemMap.
   virtual void Finalize() {}
-
+ 
   virtual DWARFASTParser *GetDWARFParser() { return nullptr; }
   virtual PDBASTParser *GetPDBParser() { return nullptr; }
   virtual npdb::PdbAstBuilder *GetNativePDBParser() { return nullptr; }
@@ -526,39 +527,41 @@ class TypeSystemMap {
 
   // Iterate through all of the type systems that are created. Return true from
   // callback to keep iterating, false to stop iterating.
-  void ForEach(std::function<bool(TypeSystem *)> const &callback);
+  void ForEach(std::function<bool(lldb::TypeSystemSP)> const &callback);
 
-  llvm::Expected<TypeSystem &>
+  llvm::Expected<lldb::TypeSystemSP>
   GetTypeSystemForLanguage(lldb::LanguageType language, Module *module,
                            bool can_create);
 
-  llvm::Expected<TypeSystem &>
+  llvm::Expected<lldb::TypeSystemSP>
   GetTypeSystemForLanguage(lldb::LanguageType language, Target *target,
                            bool can_create);
 
 protected:
-  typedef std::map<lldb::LanguageType, lldb::TypeSystemSP> collection;
+  typedef llvm::DenseMap<uint16_t, lldb::TypeSystemSP> collection;
   mutable std::mutex m_mutex; ///< A mutex to keep this object happy in
-                              ///multi-threaded environments.
+                              /// multi-threaded environments.
   collection m_map;
   bool m_clear_in_progress = false;
 
 private:
   typedef llvm::function_ref<lldb::TypeSystemSP()> CreateCallback;
   /// Finds the type system for the given language. If no type system could be
-  /// found for a language and a CreateCallback was provided, the value returned
-  /// by the callback will be treated as the TypeSystem for the language.
+  /// found for a language and a CreateCallback was provided, the value
+  /// returned by the callback will be treated as the TypeSystem for the
+  /// language.
   ///
   /// \param language The language for which the type system should be found.
   /// \param create_callback A callback that will be called if no previously
   ///                        created TypeSystem that fits the given language
   ///                        could found. Can be omitted if a non-existent
-  ///                        type system should be treated as an error instead.
+  ///                        type system should be treated as an error
+  ///                        instead.
   /// \return The found type system or an error.
-  llvm::Expected<TypeSystem &> GetTypeSystemForLanguage(
+  llvm::Expected<lldb::TypeSystemSP> GetTypeSystemForLanguage(
       lldb::LanguageType language,
       llvm::Optional<CreateCallback> create_callback = llvm::None);
-};
+  };
 
 } // namespace lldb_private
 

lldb/include/lldb/Target/Target.h

@@ -1119,11 +1119,12 @@ class Target : public std::enable_shared_from_this<Target>,
 
   PathMappingList &GetImageSearchPathList();
 
-  llvm::Expected<TypeSystem &>
+  llvm::Expected<lldb::TypeSystemSP>
   GetScratchTypeSystemForLanguage(lldb::LanguageType language,
                                   bool create_on_demand = true);
 
-  std::vector<TypeSystem *> GetScratchTypeSystems(bool create_on_demand = true);
+  std::vector<lldb::TypeSystemSP>
+  GetScratchTypeSystems(bool create_on_demand = true);
 
   PersistentExpressionState *
   GetPersistentExpressionStateForLanguage(lldb::LanguageType language);

lldb/include/lldb/lldb-enumerations.h

@@ -9,6 +9,7 @@
 #ifndef LLDB_LLDB_ENUMERATIONS_H
 #define LLDB_LLDB_ENUMERATIONS_H
 
+#include <cstdint>
 #include <type_traits>
 
 #ifndef SWIG
@@ -433,7 +434,7 @@ FLAGS_ENUM(WatchpointEventType){
 /// specification for ease of use and consistency.
 /// The enum -> string code is in Language.cpp, don't change this
 /// table without updating that code as well.
-enum LanguageType {
+enum LanguageType : uint16_t {
   eLanguageTypeUnknown = 0x0000,        ///< Unknown or invalid language value.
   eLanguageTypeC89 = 0x0001,            ///< ISO C:1989.
   eLanguageTypeC = 0x0002,              ///< Non-standardized C, such as K&R.

lldb/include/lldb/lldb-forward.h

@@ -432,6 +432,7 @@ typedef std::shared_ptr<lldb_private::TypeMemberFunctionImpl>
 typedef std::shared_ptr<lldb_private::TypeEnumMemberImpl> TypeEnumMemberImplSP;
 typedef std::shared_ptr<lldb_private::TypeFilterImpl> TypeFilterImplSP;
 typedef std::shared_ptr<lldb_private::TypeSystem> TypeSystemSP;
+typedef std::weak_ptr<lldb_private::TypeSystem> TypeSystemWP;
 typedef std::shared_ptr<lldb_private::TypeFormatImpl> TypeFormatImplSP;
 typedef std::shared_ptr<lldb_private::TypeNameSpecifierImpl>
     TypeNameSpecifierImplSP;