[MsgPack] Add code for floating point assignment and writes (#153544)

Allow assignment of float to DocType and support output of float in
writeToBlob method.
Expand tests coverage to various missing basic I/O operations.

Co-authored-by: Xavi Zhang <[email protected]>

Author: perlfu

llvm/include/llvm/BinaryFormat/MsgPackDocument.h

@@ -213,6 +213,7 @@ class DocNode {
   LLVM_ABI DocNode &operator=(unsigned Val);
   LLVM_ABI DocNode &operator=(int64_t Val);
   LLVM_ABI DocNode &operator=(uint64_t Val);
+  LLVM_ABI DocNode &operator=(double Val);
 
 private:
   // Private constructor setting KindAndDoc, used by methods in Document.

llvm/lib/BinaryFormat/MsgPackDocument.cpp

@@ -104,6 +104,10 @@ DocNode &DocNode::operator=(uint64_t Val) {
   *this = getDocument()->getNode(Val);
   return *this;
 }
+DocNode &DocNode::operator=(double Val) {
+  *this = getDocument()->getNode(Val);
+  return *this;
+}
 
 // A level in the document reading stack.
 struct StackLevel {
@@ -293,6 +297,9 @@ void Document::writeToBlob(std::string &Blob) {
     case Type::Binary:
       MPWriter.write(Node.getBinary());
       break;
+    case Type::Float:
+      MPWriter.write(Node.getFloat());
+      break;
     case Type::Empty:
       llvm_unreachable("unhandled empty msgpack node");
     default:

llvm/unittests/BinaryFormat/MsgPackDocumentTest.cpp

@@ -22,12 +22,58 @@ TEST(MsgPackDocument, DocNodeTest) {
   ASSERT_TRUE(Str1 == Str2);
 }
 
+TEST(MsgPackDocument, TestReadBoolean) {
+  Document Doc1;
+  bool Ok = Doc1.readFromBlob(StringRef("\xC2", 1), /*Multi=*/false);
+  ASSERT_TRUE(Ok);
+  ASSERT_EQ(Doc1.getRoot().getKind(), Type::Boolean);
+  ASSERT_EQ(Doc1.getRoot().getBool(), false);
+  Document Doc2;
+  Ok = Doc2.readFromBlob(StringRef("\xC3", 1), /*Multi=*/false);
+  ASSERT_TRUE(Ok);
+  ASSERT_EQ(Doc2.getRoot().getKind(), Type::Boolean);
+  ASSERT_EQ(Doc2.getRoot().getBool(), true);
+}
+
 TEST(MsgPackDocument, TestReadInt) {
-  Document Doc;
-  bool Ok = Doc.readFromBlob(StringRef("\xd0\x00", 2), /*Multi=*/false);
+  Document Doc1;
+  bool Ok = Doc1.readFromBlob(StringRef("\xD0\x00", 2), /*Multi=*/false);
+  ASSERT_TRUE(Ok);
+  ASSERT_EQ(Doc1.getRoot().getKind(), Type::Int);
+  ASSERT_EQ(Doc1.getRoot().getInt(), 0);
+  Document Doc2;
+  Ok = Doc2.readFromBlob(StringRef("\xFF", 1), /*Multi=*/false);
+  ASSERT_TRUE(Ok);
+  ASSERT_EQ(Doc2.getRoot().getKind(), Type::Int);
+  ASSERT_EQ(Doc2.getRoot().getInt(), -1);
+}
+
+TEST(MsgPackDocument, TestReadUInt) {
+  Document Doc1;
+  bool Ok = Doc1.readFromBlob(StringRef("\xCC\x00", 2), /*Multi=*/false);
+  ASSERT_TRUE(Ok);
+  ASSERT_EQ(Doc1.getRoot().getKind(), Type::UInt);
+  ASSERT_EQ(Doc1.getRoot().getUInt(), 0U);
+  Document Doc2;
+  Ok = Doc2.readFromBlob(StringRef("\x01", 1), /*Multi=*/false);
   ASSERT_TRUE(Ok);
-  ASSERT_EQ(Doc.getRoot().getKind(), Type::Int);
-  ASSERT_EQ(Doc.getRoot().getInt(), 0);
+  ASSERT_EQ(Doc2.getRoot().getKind(), Type::UInt);
+  ASSERT_EQ(Doc2.getRoot().getUInt(), 1U);
+}
+
+TEST(MsgPackDocument, TestReadFloat) {
+  Document Doc1;
+  bool Ok =
+      Doc1.readFromBlob(StringRef("\xCA\x3F\x80\x00\x00", 5), /*Multi=*/false);
+  ASSERT_TRUE(Ok);
+  ASSERT_EQ(Doc1.getRoot().getKind(), Type::Float);
+  ASSERT_EQ(Doc1.getRoot().getFloat(), 1.0);
+  Document Doc2;
+  Ok = Doc2.readFromBlob(StringRef("\xCB\x48\x3D\x63\x29\xF1\xC3\x5C\xA5", 9),
+                         /*Multi=*/false);
+  ASSERT_TRUE(Ok);
+  ASSERT_EQ(Doc2.getRoot().getKind(), Type::Float);
+  ASSERT_EQ(Doc2.getRoot().getFloat(), 1e40);
 }
 
 TEST(MsgPackDocument, TestReadBinary) {
@@ -192,12 +238,54 @@ TEST(MsgPackDocument, TestReadMergeMap) {
   ASSERT_EQ(BayS.getInt(), 8);
 }
 
+TEST(MsgPackDocument, TestWriteBoolean) {
+  Document Doc;
+  Doc.getRoot() = true;
+  std::string Buffer;
+  Doc.writeToBlob(Buffer);
+  ASSERT_EQ(Buffer, "\xc3");
+  Doc.getRoot() = false;
+  Doc.writeToBlob(Buffer);
+  ASSERT_EQ(Buffer, "\xc2");
+}
+
 TEST(MsgPackDocument, TestWriteInt) {
   Document Doc;
   Doc.getRoot() = 1;
   std::string Buffer;
   Doc.writeToBlob(Buffer);
   ASSERT_EQ(Buffer, "\x01");
+  Doc.getRoot() = -1;
+  Doc.writeToBlob(Buffer);
+  ASSERT_EQ(Buffer, "\xFF");
+  Doc.getRoot() = -4096;
+  Doc.writeToBlob(Buffer);
+  ASSERT_EQ(Buffer, StringRef("\xD1\xF0\x00", 3));
+}
+
+TEST(MsgPackDocument, TestWriteUInt) {
+  Document Doc;
+  Doc.getRoot() = 1U;
+  std::string Buffer;
+  Doc.writeToBlob(Buffer);
+  ASSERT_EQ(Buffer, "\x01");
+  Doc.getRoot() = 4096U;
+  Doc.writeToBlob(Buffer);
+  ASSERT_EQ(Buffer, StringRef("\xCD\x10\x00", 3));
+}
+
+TEST(MsgPackDocument, TestWriteFloat) {
+  Document Doc;
+  Doc.getRoot() = 1.0;
+  std::string Buffer;
+  Doc.writeToBlob(Buffer);
+  ASSERT_EQ(Buffer, StringRef("\xCA\x3F\x80\x00\x00", 5));
+  Doc.getRoot() = 1.0f;
+  Doc.writeToBlob(Buffer);
+  ASSERT_EQ(Buffer, StringRef("\xCA\x3F\x80\x00\x00", 5));
+  Doc.getRoot() = 1e40;
+  Doc.writeToBlob(Buffer);
+  ASSERT_EQ(Buffer, "\xCB\x48\x3D\x63\x29\xF1\xC3\x5C\xA5");
 }
 
 TEST(MsgPackDocument, TestWriteBinary) {