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