2 files changed, 71 insertions, 3 deletions

diff --git a/llvm/include/llvm/Object/WindowsResource.h b/llvm/include/llvm/Object/WindowsResource.h
index f94ad09ce0c..6867798db09 100644
--- a/llvm/include/llvm/Object/WindowsResource.h
+++ b/llvm/include/llvm/Object/WindowsResource.h
@@ -30,11 +30,18 @@
 #define LLVM_INCLUDE_LLVM_OBJECT_RESFILE_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Object/Binary.h"
+#include "llvm/Object/Error.h"
 #include "llvm/Support/BinaryByteStream.h"
 #include "llvm/Support/BinaryStreamReader.h"
+#include "llvm/Support/COFF.h"
+#include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/ScopedPrinter.h"
+
+#include <map>
 
 namespace llvm {
 namespace object {
@@ -44,23 +51,44 @@ class WindowsResource;
 class ResourceEntryRef {
 public:
   Error moveNext(bool &End);
+  bool checkTypeString() const { return IsStringType; }
+  ArrayRef<UTF16> getTypeString() const { return Type; }
+  uint16_t getTypeID() const { return TypeID; }
+  bool checkNameString() const { return IsStringName; }
+  ArrayRef<UTF16> getNameString() const { return Name; }
+  uint16_t getNameID() const { return NameID; }
+  uint16_t getLanguage() const { return Suffix->Language; }
 
 private:
   friend class WindowsResource;
 
   ResourceEntryRef(BinaryStreamRef Ref, const WindowsResource *Owner,
                    Error &Err);
+
   Error loadNext();
 
+  struct HeaderSuffix {
+    support::ulittle32_t DataVersion;
+    support::ulittle16_t MemoryFlags;
+    support::ulittle16_t Language;
+    support::ulittle32_t Version;
+    support::ulittle32_t Characteristics;
+  };
+
   BinaryStreamReader Reader;
-  BinaryStreamRef HeaderBytes;
-  BinaryStreamRef DataBytes;
+  bool IsStringType;
+  ArrayRef<UTF16> Type;
+  uint16_t TypeID;
+  bool IsStringName;
+  ArrayRef<UTF16> Name;
+  uint16_t NameID;
+  const HeaderSuffix *Suffix = nullptr;
+  ArrayRef<uint8_t> Data;
   const WindowsResource *OwningRes = nullptr;
 };
 
 class WindowsResource : public Binary {
 public:
-  ~WindowsResource() override;
   Expected<ResourceEntryRef> getHeadEntry();
 
   static bool classof(const Binary *V) { return V->isWinRes(); }
@@ -76,6 +104,38 @@ private:
   BinaryByteStream BBS;
 };
 
+class WindowsResourceParser {
+public:
+  WindowsResourceParser();
+
+  Error parse(WindowsResource *WR);
+
+  void printTree() const;
+
+private:
+  class TreeNode {
+  public:
+    TreeNode() = default;
+    explicit TreeNode(uint32_t ID);
+    explicit TreeNode(ArrayRef<UTF16> Ref);
+    void addEntry(const ResourceEntryRef &Entry);
+    void print(ScopedPrinter &Writer, StringRef Name) const;
+
+  private:
+    TreeNode &addTypeNode(const ResourceEntryRef &Entry);
+    TreeNode &addNameNode(const ResourceEntryRef &Entry);
+    TreeNode &addLanguageNode(const ResourceEntryRef &Entry);
+    TreeNode &addChild(uint32_t ID);
+    TreeNode &addChild(ArrayRef<UTF16> NameRef);
+    uint16_t ID;
+    std::vector<UTF16> Name;
+    std::map<uint32_t, std::unique_ptr<TreeNode>> IDChildren;
+    std::map<std::string, std::unique_ptr<TreeNode>> StringChildren;
+  };
+
+  TreeNode Root;
+};
+
 } // namespace object
 } // namespace llvm
 
diff --git a/llvm/include/llvm/Support/BinaryStreamReader.h b/llvm/include/llvm/Support/BinaryStreamReader.h
index 56375f41d2c..29e8a2ab08a 100644
--- a/llvm/include/llvm/Support/BinaryStreamReader.h
+++ b/llvm/include/llvm/Support/BinaryStreamReader.h
@@ -14,6 +14,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/BinaryStreamArray.h"
 #include "llvm/Support/BinaryStreamRef.h"
+#include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/type_traits.h"
@@ -104,6 +105,13 @@ public:
   /// returns an appropriate error code.
   Error readCString(StringRef &Dest);
 
+  /// Similar to readCString, however read a null-terminated UTF16 string
+  /// instead.
+  ///
+  /// \returns a success error code if the data was successfully read, otherwise
+  /// returns an appropriate error code.
+  Error readWideString(ArrayRef<UTF16> &Dest);
+
   /// Read a \p Length byte string into \p Dest.  Whether a copy occurs depends
   /// on the implementation of the underlying stream.  Updates the stream's
   /// offset to point after the newly read data.