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