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
20using namespace llvm;
21using namespace msgpack;
22
23// Convert this DocNode into an empty array.
24void DocNode::convertToArray() { *this = getDocument()->getArrayNode(); }
25
26// Convert this DocNode into an empty map.
27void DocNode::convertToMap() { *this = getDocument()->getMapNode(); }
28
29/// Find the key in the MapDocNode.
30DocNode::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.
36DocNode &MapDocNode::operator[](StringRef S) {
37  return (*this)[getDocument()->getNode(S)];
38}
39
40/// Member access for MapDocNode.
41DocNode &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.
52DocNode &MapDocNode::operator[](int Key) {
53  return (*this)[getDocument()->getNode(Key)];
54}
55DocNode &MapDocNode::operator[](unsigned Key) {
56  return (*this)[getDocument()->getNode(Key)];
57}
58DocNode &MapDocNode::operator[](int64_t Key) {
59  return (*this)[getDocument()->getNode(Key)];
60}
61DocNode &MapDocNode::operator[](uint64_t Key) {
62  return (*this)[getDocument()->getNode(Key)];
63}
64
65/// Array element access. This extends the array if necessary.
66DocNode &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.
79DocNode &DocNode::operator=(StringRef Val) {
80  *this = getDocument()->getNode(Val);
81  return *this;
82}
83DocNode &DocNode::operator=(bool Val) {
84  *this = getDocument()->getNode(Val);
85  return *this;
86}
87DocNode &DocNode::operator=(int Val) {
88  *this = getDocument()->getNode(Val);
89  return *this;
90}
91DocNode &DocNode::operator=(unsigned Val) {
92  *this = getDocument()->getNode(Val);
93  return *this;
94}
95DocNode &DocNode::operator=(int64_t Val) {
96  *this = getDocument()->getNode(Val);
97  return *this;
98}
99DocNode &DocNode::operator=(uint64_t Val) {
100  *this = getDocument()->getNode(Val);
101  return *this;
102}
103
104// A level in the document reading stack.
105struct 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
127bool 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
239struct 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.
247void 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}
315