[stdlib] Integrate hint_trivial_type in List.__copyinit__

stdlib/benchmarks/collections/bench_list.mojo

@@ -0,0 +1,93 @@
+# ===----------------------------------------------------------------------=== #
+# Copyright (c) 2024, Modular Inc. All rights reserved.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions:
+# https://llvm.org/LICENSE.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===----------------------------------------------------------------------=== #
+# RUN: %mojo-no-debug %s -t
+# NOTE: to test changes on the current branch using run-benchmarks.sh, remove
+# the -t flag. Remember to replace it again before pushing any code.
+
+from benchmark import Bench, BenchConfig, Bencher, BenchId, Unit, keep, run
+from random import seed
+from collections import List
+from random import *
+
+
+# ===-----------------------------------------------------------------------===#
+# Benchmark Data
+# ===-----------------------------------------------------------------------===#
+fn make_list[
+    size: Int, DT: DType, is_trivial: Bool
+]() -> List[Scalar[DT], is_trivial]:
+    alias scalar_t = Scalar[DT]
+    var d = List[Scalar[DT], is_trivial](capacity=size)
+    rand[DT](
+        d.unsafe_ptr(),
+        size,
+        min=scalar_t.MIN.cast[DType.float64](),
+        max=scalar_t.MAX.cast[DType.float64](),
+    )
+    d.size = size
+    return d
+
+
+# ===-----------------------------------------------------------------------===#
+# Benchmark `List[DT, True].__copyinit__`
+# ===-----------------------------------------------------------------------===#
+
+
+@parameter
+fn bench_list_copyinit[
+    size: Int, DT: DType, is_trivial: Bool
+](mut b: Bencher) raises:
+    var items = make_list[size, DT, is_trivial]()
+    var result = List[Scalar[DT], is_trivial]()
+    var res = 0
+
+    @always_inline
+    @parameter
+    fn call_fn() raises:
+        result = items
+        res += len(result)
+        keep(result.data)
+        keep(items.data)
+
+    b.iter[call_fn]()
+    print(res)
+    keep(bool(items))
+    keep(bool(result))
+    keep(result.data)
+    keep(items.data)
+
+
+def main():
+    var m = Bench(
+        BenchConfig(
+            num_repetitions=1,
+            max_runtime_secs=0.5,
+            min_runtime_secs=0.25,
+            min_warmuptime_secs=0,  # FIXME: adjust the values
+        )
+    )
+    alias lengths = (1, 2, 4, 8, 16, 32, 128, 256, 512, 1024, 2048, 4096)
+
+    @parameter
+    for i in range(len(lengths)):
+        alias length = lengths.get[i, Int]()
+        m.bench_function[bench_list_copyinit[length, DType.uint8, True]](
+            BenchId("List[Scalar[DT], True].__copyinit__ [" + str(length) + "]")
+        )
+        m.bench_function[bench_list_copyinit[length, DType.uint8, False]](
+            BenchId(
+                "List[Scalar[DT], False].__copyinit__ [" + str(length) + "]"
+            )
+        )
+
+    m.dump_report()

stdlib/src/collections/list.mojo

@@ -206,8 +206,14 @@ struct List[T: CollectionElement, hint_trivial_type: Bool = False](
             existing: The list to copy.
         """
         self = Self(capacity=existing.capacity)
-        for i in range(len(existing)):
-            self.append(existing[i])
+
+        @parameter
+        if hint_trivial_type:
+            memcpy(self.data, existing.data, len(existing))
+            self.size = existing.size
+        else:
+            for i in range(len(existing)):
+                self.append(existing[i])
 
     fn __del__(owned self):
         """Destroy all elements in the list and free its memory."""

stdlib/test/collections/test_list.mojo

@@ -18,6 +18,7 @@ from sys.info import sizeof
 from memory import UnsafePointer, Span
 from test_utils import CopyCounter, MoveCounter
 from testing import assert_equal, assert_false, assert_raises, assert_true
+from testing import assert_not_equal
 
 
 def test_mojo_issue_698():
@@ -935,6 +936,48 @@ def test_list_repr():
     assert_equal(empty.__repr__(), "[]")
 
 
+def test_copyinit_trivial_types[dt: DType, hint_trivial_type: Bool]():
+    alias sizes = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512)
+    assert_equal(len(sizes), 10)
+    var test_current_size = 1
+
+    @parameter
+    for sizes_index in range(len(sizes)):
+        alias current_size = sizes.get[sizes_index, Int]()
+        x = List[Scalar[dt], hint_trivial_type]()
+        for i in range(current_size):
+            x.append(i)
+        y = x
+        assert_equal(test_current_size, current_size)
+        assert_equal(len(y), current_size)
+        assert_not_equal(int(x.data), int(y.data))
+        for i in range(current_size):
+            assert_equal(i, x[i])
+            assert_equal(y[i], x[i])
+        test_current_size *= 2
+    assert_equal(test_current_size, 1024)
+
+
+def test_copyinit_trivial_types_dtypes():
+    alias dtypes = (
+        DType.int64,
+        DType.int32,
+        DType.float64,
+        DType.float32,
+        DType.uint8,
+        DType.int8,
+        DType.bool,
+    )
+    var test_index_dtype = 0
+
+    @parameter
+    for index_dtype in range(len(dtypes)):
+        test_copyinit_trivial_types[dtypes.get[index_dtype, DType](), True]()
+        test_copyinit_trivial_types[dtypes.get[index_dtype, DType](), False]()
+        test_index_dtype += 1
+    assert_equal(test_index_dtype, 7)
+
+
 # ===-------------------------------------------------------------------===#
 # main
 # ===-------------------------------------------------------------------===#
@@ -974,3 +1017,4 @@ def main():
     test_list_dtor()
     test_destructor_trivial_elements()
     test_list_repr()
+    test_copyinit_trivial_types_dtypes()

stdlib/test/collections/test_string.mojo

@@ -1593,6 +1593,26 @@ def test_reserve():
     assert_equal(s._buffer.capacity, 1)
 
 
+def test_copyinit():
+    alias sizes = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512)
+    assert_equal(len(sizes), 10)
+    var test_current_size = 1
+
+    @parameter
+    for sizes_index in range(len(sizes)):
+        alias current_size = sizes.get[sizes_index, Int]()
+        x = String("")
+        for i in range(current_size):
+            x += str(i)[0]
+        y = x
+        assert_equal(x._buffer, y._buffer)
+        assert_equal(test_current_size, current_size)
+        assert_equal(len(y), current_size)
+        assert_not_equal(int(x._buffer.data), int(y._buffer.data))
+        test_current_size *= 2
+    assert_equal(test_current_size, 1024)
+
+
 def main():
     test_constructors()
     test_copy()
@@ -1649,3 +1669,4 @@ def main():
     test_center()
     test_float_conversion()
     test_slice_contains()
+    test_copyinit()