From e47f2c704fad3f6bec7387f333f1c88cedf74d20 Mon Sep 17 00:00:00 2001
From: Lukas Hermann <lukashermann28@gmail.com>
Date: Thu, 7 Nov 2024 16:01:47 -0800
Subject: [PATCH] implement Writer for List[Byte]

Signed-off-by: Lukas Hermann <lukashermann28@gmail.com>
---
 stdlib/src/collections/list.mojo       | 41 ++++++++++++++++++++++++++
 stdlib/test/collections/test_list.mojo | 13 ++++++++
 2 files changed, 54 insertions(+)

diff --git a/stdlib/src/collections/list.mojo b/stdlib/src/collections/list.mojo
index b849766a9e..3496bfebb0 100644
--- a/stdlib/src/collections/list.mojo
+++ b/stdlib/src/collections/list.mojo
@@ -419,6 +419,35 @@ struct List[T: CollectionElement, hint_trivial_type: Bool = False](
         self.write_to(output)
         return output^
 
+    fn write_bytes(inout self, bytes: Span[Byte, _]):
+        """
+        Write a byte span to this List. `T` **must** be of type `Byte`.
+
+        Args:
+            bytes: The byte span to write to this String. Must NOT be
+              null terminated.
+        """
+        constrained[_type_is_eq[T, Byte]()]()
+        self.extend(rebind[Span[T, bytes.origin]](bytes))
+
+    fn write[*Ts: Writable](inout self, *args: *Ts):
+        """Write a sequence of Writable arguments to the provided Writer.
+        `T` **must** be of type `Byte`.
+
+        Parameters:
+            Ts: Types of the provided argument sequence.
+
+        Args:
+            args: Sequence of arguments to write to this Writer.
+        """
+        constrained[_type_is_eq[T, Byte]()]()
+
+        @parameter
+        fn write_arg[T: Writable](arg: T):
+            arg.write_to[__type_of(self)](self)
+
+        args.each[write_arg]()
+
     @no_inline
     fn write_to[
         W: Writer, U: RepresentableCollectionElement, //
@@ -592,6 +621,18 @@ struct List[T: CollectionElement, hint_trivial_type: Bool = False](
         # list.
         self.size = final_size
 
+    fn extend(inout self, other: Span[T]):
+        """Extends this list by copying the elements of `other`.
+
+        Args:
+            other: Span whose elements will be added in order at the end of this list.
+        """
+        var final_size = len(self) + len(other)
+        self.reserve(final_size)
+
+        for item in other:
+            self.append(item[])
+
     fn pop(inout self, i: Int = -1) -> T:
         """Pops a value from the list at the given index.
 
diff --git a/stdlib/test/collections/test_list.mojo b/stdlib/test/collections/test_list.mojo
index 96b88b8a83..8aab693ec1 100644
--- a/stdlib/test/collections/test_list.mojo
+++ b/stdlib/test/collections/test_list.mojo
@@ -922,6 +922,19 @@ def test_list_repr():
     assert_equal(empty.__repr__(), "[]")
 
 
+def test_writer():
+    var l = List[Byte]()
+    var s = "abc".as_bytes()
+    l.write_bytes(s)
+    assert_equal(l[0], ord("a"))
+    assert_equal(l[1], ord("b"))
+    assert_equal(l[2], ord("c"))
+    var l2 = List[Byte]()
+    assert_equal(l2[0], 0)
+    assert_equal(l2[1], 1)
+    assert_equal(l2[2], 2)
+
+
 # ===-------------------------------------------------------------------===#
 # main
 # ===-------------------------------------------------------------------===#