lib/BinaryFormat/MsgPackDocument.cpp

   1 //===-- MsgPackDocument.cpp - MsgPack Document --------------------------*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 ///
   9 /// This file implements a class that exposes a simple in-memory representation
  10 /// of a document of MsgPack objects, that can be read from MsgPack, written to
  11 /// MsgPack, and inspected and modified in memory. This is intended to be a
  12 /// lighter-weight (in terms of memory allocations) replacement for
  13 /// MsgPackTypes.
  14 ///
  15 //===----------------------------------------------------------------------===//
  16
  17 #include "llvm/BinaryFormat/MsgPackDocument.h"
  18 #include "llvm/BinaryFormat/MsgPackWriter.h"
  19
  20 using namespace llvm;
  21 using namespace msgpack;
  22
  23 // Convert this DocNode into an empty array.
  24 void DocNode::convertToArray() { *this = getDocument()->getArrayNode(); }
  25
  26 // Convert this DocNode into an empty map.
  27 void DocNode::convertToMap() { *this = getDocument()->getMapNode(); }
  28
  29 /// Find the key in the MapDocNode.
  30 DocNode::MapTy::iterator MapDocNode::find(StringRef S) {
  31   return find(getDocument()->getNode(S));
  32 }
  33
  34 /// Member access for MapDocNode. The string data must remain valid for the
  35 /// lifetime of the Document.
  36 DocNode &MapDocNode::operator[](StringRef S) {
  37   return (*this)[getDocument()->getNode(S)];
  38 }
  39
  40 /// Member access for MapDocNode.
  41 DocNode &MapDocNode::operator[](DocNode Key) {
  42   assert(!Key.isEmpty());
  43   MapTy::value_type Entry(Key, DocNode());
  44   auto ItAndInserted = Map->insert(Entry);
  45   if (ItAndInserted.second) {
  46     // Ensure a new element has its KindAndDoc initialized.
  47     ItAndInserted.first->second = getDocument()->getNode();
  48   }
  49   return ItAndInserted.first->second;
  50 }
  51
  52 /// Array element access. This extends the array if necessary.
  53 DocNode &ArrayDocNode::operator[](size_t Index) {
  54   if (size() <= Index) {
  55     // Ensure new elements have their KindAndDoc initialized.
  56     Array->resize(Index + 1, getDocument()->getNode());
  57   }
  58   return (*Array)[Index];
  59 }
  60
  61 // A level in the document reading stack.
  62 struct StackLevel {
  63   DocNode Node;
  64   size_t Length;
  65   // Points to map entry when we have just processed a map key.
  66   DocNode *MapEntry;
  67 };
  68
  69 // Read a document from a binary msgpack blob.
  70 // The blob data must remain valid for the lifetime of this Document (because a
  71 // string object in the document contains a StringRef into the original blob).
  72 // If Multi, then this sets root to an array and adds top-level objects to it.
  73 // If !Multi, then it only reads a single top-level object, even if there are
  74 // more, and sets root to that.
  75 // Returns false if failed due to illegal format.
  76 bool Document::readFromBlob(StringRef Blob, bool Multi) {
  77   msgpack::Reader MPReader(Blob);
  78   SmallVector<StackLevel, 4> Stack;
  79   if (Multi) {
  80     // Create the array for multiple top-level objects.
  81     Root = getArrayNode();
  82     Stack.push_back(StackLevel({Root, (size_t)-1, nullptr}));
  83   }
  84   do {
  85     // On to next element (or key if doing a map key next).
  86     // Read the value.
  87     Object Obj;
  88     if (!MPReader.read(Obj)) {
  89       if (Multi && Stack.size() == 1) {
  90         // OK to finish here as we've just done a top-level element with Multi
  91         break;
  92       }
  93       return false; // Finished too early
  94     }
  95     // Convert it into a DocNode.
  96     DocNode Node;
  97     switch (Obj.Kind) {
  98     case Type::Nil:
  99       Node = getNode();
 100       break;
 101     case Type::Int:
 102       Node = getNode(Obj.Int);
 103       break;
 104     case Type::UInt:
 105       Node = getNode(Obj.UInt);
 106       break;
 107     case Type::Boolean:
 108       Node = getNode(Obj.Bool);
 109       break;
 110     case Type::Float:
 111       Node = getNode(Obj.Float);
 112       break;
 113     case Type::String:
 114       Node = getNode(Obj.Raw);
 115       break;
 116     case Type::Map:
 117       Node = getMapNode();
 118       break;
 119     case Type::Array:
 120       Node = getArrayNode();
 121       break;
 122     default:
 123       return false; // Raw and Extension not supported
 124     }
 125
 126     // Store it.
 127     if (Stack.empty())
 128       Root = Node;
 129     else if (Stack.back().Node.getKind() == Type::Array) {
 130       // Reading an array entry.
 131       auto &Array = Stack.back().Node.getArray();
 132       Array.push_back(Node);
 133     } else {
 134       auto &Map = Stack.back().Node.getMap();
 135       if (!Stack.back().MapEntry) {
 136         // Reading a map key.
 137         Stack.back().MapEntry = &Map[Node];
 138       } else {
 139         // Reading the value for the map key read in the last iteration.
 140         *Stack.back().MapEntry = Node;
 141         Stack.back().MapEntry = nullptr;
 142       }
 143     }
 144
 145     // See if we're starting a new array or map.
 146     switch (Node.getKind()) {
 147     case msgpack::Type::Array:
 148     case msgpack::Type::Map:
 149       Stack.push_back(StackLevel({Node, Obj.Length, nullptr}));
 150       break;
 151     default:
 152       break;
 153     }
 154
 155     // Pop finished stack levels.
 156     while (!Stack.empty()) {
 157       if (Stack.back().Node.getKind() == msgpack::Type::Array) {
 158         if (Stack.back().Node.getArray().size() != Stack.back().Length)
 159           break;
 160       } else {
 161         if (Stack.back().MapEntry ||
 162             Stack.back().Node.getMap().size() != Stack.back().Length)
 163           break;
 164       }
 165       Stack.pop_back();
 166     }
 167   } while (!Stack.empty());
 168   return true;
 169 }
 170
 171 struct WriterStackLevel {
 172   DocNode Node;
 173   DocNode::MapTy::iterator MapIt;
 174   DocNode::ArrayTy::iterator ArrayIt;
 175   bool OnKey;
 176 };
 177
 178 /// Write a MsgPack document to a binary MsgPack blob.
 179 void Document::writeToBlob(std::string &Blob) {
 180   Blob.clear();
 181   raw_string_ostream OS(Blob);
 182   msgpack::Writer MPWriter(OS);
 183   SmallVector<WriterStackLevel, 4> Stack;
 184   DocNode Node = getRoot();
 185   for (;;) {
 186     switch (Node.getKind()) {
 187     case Type::Array:
 188       MPWriter.writeArraySize(Node.getArray().size());
 189       Stack.push_back(
 190           {Node, DocNode::MapTy::iterator(), Node.getArray().begin(), false});
 191       break;
 192     case Type::Map:
 193       MPWriter.writeMapSize(Node.getMap().size());
 194       Stack.push_back(
 195           {Node, Node.getMap().begin(), DocNode::ArrayTy::iterator(), true});
 196       break;
 197     case Type::Nil:
 198       MPWriter.writeNil();
 199       break;
 200     case Type::Boolean:
 201       MPWriter.write(Node.getBool());
 202       break;
 203     case Type::Int:
 204       MPWriter.write(Node.getInt());
 205       break;
 206     case Type::UInt:
 207       MPWriter.write(Node.getUInt());
 208       break;
 209     case Type::String:
 210       MPWriter.write(Node.getString());
 211       break;
 212     default:
 213       llvm_unreachable("unhandled msgpack object kind");
 214     }
 215     // Pop finished stack levels.
 216     while (!Stack.empty()) {
 217       if (Stack.back().Node.getKind() == Type::Map) {
 218         if (Stack.back().MapIt != Stack.back().Node.getMap().end())
 219           break;
 220       } else {
 221         if (Stack.back().ArrayIt != Stack.back().Node.getArray().end())
 222           break;
 223       }
 224       Stack.pop_back();
 225     }
 226     if (Stack.empty())
 227       break;
 228     // Get the next value.
 229     if (Stack.back().Node.getKind() == Type::Map) {
 230       if (Stack.back().OnKey) {
 231         // Do the key of a key,value pair in a map.
 232         Node = Stack.back().MapIt->first;
 233         Stack.back().OnKey = false;
 234       } else {
 235         Node = Stack.back().MapIt->second;
 236         ++Stack.back().MapIt;
 237         Stack.back().OnKey = true;
 238       }
 239     } else {
 240       Node = *Stack.back().ArrayIt;
 241       ++Stack.back().ArrayIt;
 242     }
 243   }
 244 }
 245