xref: /freebsd/contrib/llvm-project/llvm/lib/BinaryFormat/MsgPackDocument.cpp (revision 5f757f3ff9144b609b3c433dfd370cc6bdc191ad) 
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 }
331