Add cord data structure (#2273)

Summary:
Pull Request resolved: #2273

Introduce cord data structure to store bytes/bytearrays during serialization.

This allows us to manipulate bytes/bytearrays without copying data.

bypass-github-export-checks

Reviewed By: dbort

Differential Revision: D54514244

fbshipit-source-id: 65397dcdea93054d54feea1b9f3ebfb0940c8513

Author: lucylq

exir/_serialize/TARGETS

@@ -29,6 +29,7 @@ runtime.python_library(
     name = "lib",
     srcs = [
         "__init__.py",
+        "_cord.py",
         "_dataclass.py",
         "_flatbuffer.py",
         "_program.py",

exir/_serialize/_cord.py

@@ -0,0 +1,49 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import io
+from typing import List, Optional, Union
+
+
+class Cord:
+    """A `bytes`-like sequence of bytes, stored non-contiguously.
+
+    Users can use a Cord to assemble large files and data blobs using references
+    to and slices of other data, instead of copying and appending that data to a
+    `bytes` or `bytearray` object.
+    """
+
+    def __init__(self, data: Optional[Union[bytes, "Cord"]] = None) -> None:
+        """Initialize Cord data structure."""
+        self._buffers: List[bytes] = []
+        self._byte_size: int = 0
+
+        if data is not None:
+            self.append(data)
+
+    def __len__(self):
+        """Number of bytes in the Cord."""
+        return self._byte_size
+
+    def __bytes__(self) -> bytes:
+        """Return the contents of the Cord as a single `bytes` object."""
+        return b"".join(self._buffers)
+
+    def append(self, data: Union[bytes, "Cord"]) -> None:
+        """Append a bytes or Cord to the current Cord."""
+        if isinstance(data, bytes):
+            self._buffers.append(data)
+            self._byte_size += len(data)
+        elif isinstance(data, Cord):
+            self._buffers.extend(data._buffers)
+            self._byte_size += len(data)
+        else:
+            raise TypeError(f"Can only append bytes or Cords, received {type(data)}")
+
+    def write_to_file(self, outfile: io.BufferedIOBase) -> None:
+        """Write the Cord to a file."""
+        for item in self._buffers:
+            outfile.write(item)

exir/_serialize/test/TARGETS

@@ -23,3 +23,13 @@ python_unittest(
         "//executorch/exir/_serialize:lib",
     ],
 )
+
+python_unittest(
+    name = "cord",
+    srcs = [
+        "test_cord.py",
+    ],
+    deps = [
+        "//executorch/exir/_serialize:lib",
+    ],
+)

exir/_serialize/test/test_cord.py

@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import io
+import unittest
+
+from executorch.exir._serialize._cord import Cord
+
+
+class TestCord(unittest.TestCase):
+    def test_cord_init(self) -> None:
+        cord_empty = Cord()
+        self.assertEqual(0, len(cord_empty))
+
+        cord = Cord(b"HelloWorld")
+        self.assertEqual(10, len(cord))
+        self.assertEqual(b"HelloWorld", bytes(cord))
+
+        cord2 = Cord(cord)
+        self.assertEqual(10, len(cord2))
+        self.assertEqual(b"HelloWorld", bytes(cord))
+
+        # Confirm no copies were made.
+        self.assertEqual(id(cord._buffers[0]), id(cord2._buffers[0]))
+
+    def test_cord_append(self) -> None:
+        cord = Cord()
+        cord.append(b"Hello")
+        self.assertEqual(5, len(cord))
+        self.assertEqual(b"Hello", bytes(cord))
+
+        cord.append(b"World")
+        self.assertEqual(10, len(cord))
+        self.assertEqual(b"HelloWorld", bytes(cord))
+
+    def test_cord_append_cord(self) -> None:
+        cord = Cord()
+        cord.append(b"Hello")
+        cord.append((b"World"))
+
+        cord2 = Cord()
+        cord2.append(b"Prefix")
+        cord2.append(cord)
+
+        self.assertEqual(16, len(cord2))
+        self.assertEqual(b"PrefixHelloWorld", bytes(cord2))
+
+        # Confirm that no copies were made when appending a Cord.
+        self.assertEqual(id(cord2._buffers[1]), id(cord._buffers[0]))
+        self.assertEqual(id(cord2._buffers[2]), id(cord._buffers[1]))
+
+    def test_cord_write_to_file(self) -> None:
+        cord = Cord()
+        cord.append(b"Hello")
+        cord.append(b"World")
+
+        outfile = io.BytesIO()
+        cord.write_to_file(outfile)
+        self.assertEqual(b"HelloWorld", outfile.getvalue())