llvm/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   DocNode &N = (*Map)[Key];
  44   if (N.isEmpty()) {
  45     // Ensure a new element has its KindAndDoc initialized.
  46     N = getDocument()->getEmptyNode();
  47   }
  48   return N;
  49 }
  50
  51 /// Member access for MapDocNode for integer key.
  52 DocNode &MapDocNode::operator[](int Key) {
  53   return (*this)[getDocument()->getNode(Key)];
  54 }
  55 DocNode &MapDocNode::operator[](unsigned Key) {
  56   return (*this)[getDocument()->getNode(Key)];
  57 }
  58 DocNode &MapDocNode::operator[](int64_t Key) {
  59   return (*this)[getDocument()->getNode(Key)];
  60 }
  61 DocNode &MapDocNode::operator[](uint64_t Key) {
  62   return (*this)[getDocument()->getNode(Key)];
  63 }
  64
  65 /// Array element access. This extends the array if necessary.
  66 DocNode &ArrayDocNode::operator[](size_t Index) {
  67   if (size() <= Index) {
  68     // Ensure new elements have their KindAndDoc initialized.
  69     Array->resize(Index + 1, getDocument()->getEmptyNode());
  70   }
  71   return (*Array)[Index];
  72 }
  73
  74 // Convenience assignment operators. This only works if the destination
  75 // DocNode has an associated Document, i.e. it was not constructed using the
  76 // default constructor. The string one does not copy, so the string must
  77 // remain valid for the lifetime of the Document. Use fromString to avoid
  78 // that restriction.
  79 DocNode &DocNode::operator=(StringRef Val) {
  80   *this = getDocument()->getNode(Val);
  81   return *this;
  82 }
  83 DocNode &DocNode::operator=(MemoryBufferRef Val) {
  84   *this = getDocument()->getNode(Val);
  85   return *this;
  86 }
  87 DocNode &DocNode::operator=(bool Val) {
  88   *this = getDocument()->getNode(Val);
  89   return *this;
  90 }
  91 DocNode &DocNode::operator=(int Val) {
  92   *this = getDocument()->getNode(Val);
  93   return *this;
  94 }
  95 DocNode &DocNode::operator=(unsigned Val) {
  96   *this = getDocument()->getNode(Val);
  97   return *this;
  98 }
  99 DocNode &DocNode::operator=(int64_t Val) {
 100   *this = getDocument()->getNode(Val);
 101   return *this;
 102 }
 103 DocNode &DocNode::operator=(uint64_t Val) {
 104   *this = getDocument()->getNode(Val);
 105   return *this;
 106 }
 107
 108 // A level in the document reading stack.
 109 struct StackLevel {
 110   StackLevel(DocNode Node, size_t StartIndex, size_t Length,
 111              DocNode *MapEntry = nullptr)
 112       : Node(Node), Index(StartIndex), End(StartIndex + Length),
 113         MapEntry(MapEntry) {}
 114   DocNode Node;
 115   size_t Index;
 116   size_t End;
 117   // Points to map entry when we have just processed a map key.
 118   DocNode *MapEntry;
 119   DocNode MapKey;
 120 };
 121
 122 // Read a document from a binary msgpack blob, merging into anything already in
 123 // the Document.
 124 // The blob data must remain valid for the lifetime of this Document (because a
 125 // string object in the document contains a StringRef into the original blob).
 126 // If Multi, then this sets root to an array and adds top-level objects to it.
 127 // If !Multi, then it only reads a single top-level object, even if there are
 128 // more, and sets root to that.
 129 // Returns false if failed due to illegal format or merge error.
 130
 131 bool Document::readFromBlob(
 132     StringRef Blob, bool Multi,
 133     function_ref<int(DocNode *DestNode, DocNode SrcNode, DocNode MapKey)>
 134         Merger) {
 135   msgpack::Reader MPReader(Blob);
 136   SmallVector<StackLevel, 4> Stack;
 137   if (Multi) {
 138     // Create the array for multiple top-level objects.
 139     Root = getArrayNode();
 140     Stack.push_back(StackLevel(Root, 0, (size_t)-1));
 141   }
 142   do {
 143     // On to next element (or key if doing a map key next).
 144     // Read the value.
 145     Object Obj;
 146     Expected<bool> ReadObj = MPReader.read(Obj);
 147     if (!ReadObj) {
 148       // FIXME: Propagate the Error to the caller.
 149       consumeError(ReadObj.takeError());
 150       return false;
 151     }
 152     if (!ReadObj.get()) {
 153       if (Multi && Stack.size() == 1) {
 154         // OK to finish here as we've just done a top-level element with Multi
 155         break;
 156       }
 157       return false; // Finished too early
 158     }
 159     // Convert it into a DocNode.
 160     DocNode Node;
 161     switch (Obj.Kind) {
 162     case Type::Nil:
 163       Node = getNode();
 164       break;
 165     case Type::Int:
 166       Node = getNode(Obj.Int);
 167       break;
 168     case Type::UInt:
 169       Node = getNode(Obj.UInt);
 170       break;
 171     case Type::Boolean:
 172       Node = getNode(Obj.Bool);
 173       break;
 174     case Type::Float:
 175       Node = getNode(Obj.Float);
 176       break;
 177     case Type::String:
 178       Node = getNode(Obj.Raw);
 179       break;
 180     case Type::Binary:
 181       Node = getNode(MemoryBufferRef(Obj.Raw, ""));
 182       break;
 183     case Type::Map:
 184       Node = getMapNode();
 185       break;
 186     case Type::Array:
 187       Node = getArrayNode();
 188       break;
 189     default:
 190       return false; // Raw and Extension not supported
 191     }
 192
 193     // Store it.
 194     DocNode *DestNode = nullptr;
 195     if (Stack.empty())
 196       DestNode = &Root;
 197     else if (Stack.back().Node.getKind() == Type::Array) {
 198       // Reading an array entry.
 199       auto &Array = Stack.back().Node.getArray();
 200       DestNode = &Array[Stack.back().Index++];
 201     } else {
 202       auto &Map = Stack.back().Node.getMap();
 203       if (!Stack.back().MapEntry) {
 204         // Reading a map key.
 205         Stack.back().MapKey = Node;
 206         Stack.back().MapEntry = &Map[Node];
 207         continue;
 208       }
 209       // Reading the value for the map key read in the last iteration.
 210       DestNode = Stack.back().MapEntry;
 211       Stack.back().MapEntry = nullptr;
 212       ++Stack.back().Index;
 213     }
 214     int MergeResult = 0;
 215     if (!DestNode->isEmpty()) {
 216       // In a merge, there is already a value at this position. Call the
 217       // callback to attempt to resolve the conflict. The resolution must result
 218       // in an array or map if Node is an array or map respectively.
 219       DocNode MapKey = !Stack.empty() && !Stack.back().MapKey.isEmpty()
 220                            ? Stack.back().MapKey
 221                            : getNode();
 222       MergeResult = Merger(DestNode, Node, MapKey);
 223       if (MergeResult < 0)
 224         return false; // Merge conflict resolution failed
 225       assert(!((Node.isMap() && !DestNode->isMap()) ||
 226                (Node.isArray() && !DestNode->isArray())));
 227     } else
 228       *DestNode = Node;
 229
 230     // See if we're starting a new array or map.
 231     switch (DestNode->getKind()) {
 232     case msgpack::Type::Array:
 233     case msgpack::Type::Map:
 234       Stack.push_back(StackLevel(*DestNode, MergeResult, Obj.Length, nullptr));
 235       break;
 236     default:
 237       break;
 238     }
 239
 240     // Pop finished stack levels.
 241     while (!Stack.empty()) {
 242       if (Stack.back().MapEntry)
 243         break;
 244       if (Stack.back().Index != Stack.back().End)
 245         break;
 246       Stack.pop_back();
 247     }
 248   } while (!Stack.empty());
 249   return true;
 250 }
 251
 252 struct WriterStackLevel {
 253   DocNode Node;
 254   DocNode::MapTy::iterator MapIt;
 255   DocNode::ArrayTy::iterator ArrayIt;
 256   bool OnKey;
 257 };
 258
 259 /// Write a MsgPack document to a binary MsgPack blob.
 260 void Document::writeToBlob(std::string &Blob) {
 261   Blob.clear();
 262   raw_string_ostream OS(Blob);
 263   msgpack::Writer MPWriter(OS);
 264   SmallVector<WriterStackLevel, 4> Stack;
 265   DocNode Node = getRoot();
 266   for (;;) {
 267     switch (Node.getKind()) {
 268     case Type::Array:
 269       MPWriter.writeArraySize(Node.getArray().size());
 270       Stack.push_back(
 271           {Node, DocNode::MapTy::iterator(), Node.getArray().begin(), false});
 272       break;
 273     case Type::Map:
 274       MPWriter.writeMapSize(Node.getMap().size());
 275       Stack.push_back(
 276           {Node, Node.getMap().begin(), DocNode::ArrayTy::iterator(), true});
 277       break;
 278     case Type::Nil:
 279       MPWriter.writeNil();
 280       break;
 281     case Type::Boolean:
 282       MPWriter.write(Node.getBool());
 283       break;
 284     case Type::Int:
 285       MPWriter.write(Node.getInt());
 286       break;
 287     case Type::UInt:
 288       MPWriter.write(Node.getUInt());
 289       break;
 290     case Type::String:
 291       MPWriter.write(Node.getString());
 292       break;
 293     case Type::Binary:
 294       MPWriter.write(Node.getBinary());
 295       break;
 296     case Type::Empty:
 297       llvm_unreachable("unhandled empty msgpack node");
 298     default:
 299       llvm_unreachable("unhandled msgpack object kind");
 300     }
 301     // Pop finished stack levels.
 302     while (!Stack.empty()) {
 303       if (Stack.back().Node.getKind() == Type::Map) {
 304         if (Stack.back().MapIt != Stack.back().Node.getMap().end())
 305           break;
 306       } else {
 307         if (Stack.back().ArrayIt != Stack.back().Node.getArray().end())
 308           break;
 309       }
 310       Stack.pop_back();
 311     }
 312     if (Stack.empty())
 313       break;
 314     // Get the next value.
 315     if (Stack.back().Node.getKind() == Type::Map) {
 316       if (Stack.back().OnKey) {
 317         // Do the key of a key,value pair in a map.
 318         Node = Stack.back().MapIt->first;
 319         Stack.back().OnKey = false;
 320       } else {
 321         Node = Stack.back().MapIt->second;
 322         ++Stack.back().MapIt;
 323         Stack.back().OnKey = true;
 324       }
 325     } else {
 326       Node = *Stack.back().ArrayIt;
 327       ++Stack.back().ArrayIt;
 328     }
 329   }
 330 }