Merge pull request #4237 from snprajwal/str-withcapacity

builtin(str): implement Str.withCapacity

crates/compiler/builtins/bitcode/src/main.zig

@@ -144,6 +144,7 @@ comptime {
     exportStrFn(str.strTrimLeft, "trim_left");
     exportStrFn(str.strTrimRight, "trim_right");
     exportStrFn(str.strCloneTo, "clone_to");
+    exportStrFn(str.withCapacity, "with_capacity");
 
     inline for (INTEGERS) |T| {
         str.exportFromInt(T, ROC_BUILTINS ++ "." ++ STR ++ ".from_int.");

crates/compiler/builtins/bitcode/src/str.zig

@@ -2596,6 +2596,10 @@ pub fn reserve(string: RocStr, capacity: usize) callconv(.C) RocStr {
     }
 }
 
+pub fn withCapacity(capacity: usize) callconv(.C) RocStr {
+    return RocStr.allocate(0, capacity);
+}
+
 pub fn getScalarUnsafe(string: RocStr, index: usize) callconv(.C) extern struct { bytesParsed: usize, scalar: u32 } {
     const slice = string.asSlice();
     const bytesParsed = @intCast(usize, std.unicode.utf8ByteSequenceLength(slice[index]) catch unreachable);

crates/compiler/builtins/roc/Str.roc

@@ -43,6 +43,7 @@ interface Str
         appendScalar,
         walkScalars,
         walkScalarsUntil,
+        withCapacity,
     ]
     imports [
         Bool.{ Bool },
@@ -144,6 +145,9 @@ Utf8Problem : { byteIndex : Nat, problem : Utf8ByteProblem }
 isEmpty : Str -> Bool
 concat : Str, Str -> Str
 
+## Returns a string of the specified capacity without any content
+withCapacity : Nat -> Str
+
 ## Combine a list of strings into a single string, with a separator
 ## string in between each.
 ##

crates/compiler/builtins/src/bitcode.rs

@@ -361,6 +361,7 @@ pub const STR_RESERVE: &str = "roc_builtins.str.reserve";
 pub const STR_APPEND_SCALAR: &str = "roc_builtins.str.append_scalar";
 pub const STR_GET_SCALAR_UNSAFE: &str = "roc_builtins.str.get_scalar_unsafe";
 pub const STR_CLONE_TO: &str = "roc_builtins.str.clone_to";
+pub const STR_WITH_CAPACITY: &str = "roc_builtins.str.with_capacity";
 
 pub const LIST_MAP: &str = "roc_builtins.list.map";
 pub const LIST_MAP2: &str = "roc_builtins.list.map2";

crates/compiler/can/src/builtins.rs

@@ -123,6 +123,7 @@ map_symbol_to_lowlevel_and_arity! {
     StrGetScalarUnsafe; STR_GET_SCALAR_UNSAFE; 2,
     StrToNum; STR_TO_NUM; 1,
     StrGetCapacity; STR_CAPACITY; 1,
+    StrWithCapacity; STR_WITH_CAPACITY; 1,
 
     ListLen; LIST_LEN; 1,
     ListWithCapacity; LIST_WITH_CAPACITY; 1,

crates/compiler/gen_llvm/src/llvm/build.rs

@@ -6035,6 +6035,20 @@ fn run_low_level<'a, 'ctx, 'env>(
                 bitcode::STR_TRIM_RIGHT,
             )
         }
+        StrWithCapacity => {
+            // Str.withCapacity : Nat -> Str
+            debug_assert_eq!(args.len(), 1);
+
+            let str_len = load_symbol(scope, &args[0]);
+
+            call_str_bitcode_fn(
+                env,
+                &[],
+                &[str_len],
+                BitcodeReturns::Str,
+                bitcode::STR_WITH_CAPACITY,
+            )
+        }
         ListLen => {
             // List.len : List * -> Nat
             debug_assert_eq!(args.len(), 1);
@@ -6157,7 +6171,7 @@ fn run_low_level<'a, 'ctx, 'env>(
             list_prepend(env, original_wrapper, elem, elem_layout)
         }
         StrGetUnsafe => {
-            // List.getUnsafe : Str, Nat -> u8
+            // Str.getUnsafe : Str, Nat -> u8
             debug_assert_eq!(args.len(), 2);
 
             let wrapper_struct = load_symbol(scope, &args[0]);

crates/compiler/gen_wasm/src/low_level.rs

@@ -302,6 +302,7 @@ impl<'a> LowLevelCall<'a> {
             StrSubstringUnsafe => {
                 self.load_args_and_call_zig(backend, bitcode::STR_SUBSTRING_UNSAFE)
             }
+            StrWithCapacity => self.load_args_and_call_zig(backend, bitcode::STR_WITH_CAPACITY),
 
             // List
             ListLen => match backend.storage.get(&self.arguments[0]) {

crates/compiler/module/src/low_level.rs

@@ -30,6 +30,7 @@ pub enum LowLevel {
     StrAppendScalar,
     StrGetScalarUnsafe,
     StrGetCapacity,
+    StrWithCapacity,
     ListLen,
     ListWithCapacity,
     ListReserve,
@@ -249,6 +250,7 @@ map_symbol_to_lowlevel! {
     StrGetScalarUnsafe <= STR_GET_SCALAR_UNSAFE,
     StrToNum <= STR_TO_NUM,
     StrGetCapacity <= STR_CAPACITY,
+    StrWithCapacity <= STR_WITH_CAPACITY,
     ListLen <= LIST_LEN,
     ListGetCapacity <= LIST_CAPACITY,
     ListWithCapacity <= LIST_WITH_CAPACITY,

crates/compiler/module/src/symbol.rs

@@ -1296,6 +1296,7 @@ define_builtins! {
         50 STR_REPLACE_EACH: "replaceEach"
         51 STR_REPLACE_FIRST: "replaceFirst"
         52 STR_REPLACE_LAST: "replaceLast"
+        53 STR_WITH_CAPACITY: "withCapacity"
     }
     6 LIST: "List" => {
         0 LIST_LIST: "List" exposed_apply_type=true // the List.List type alias

crates/compiler/mono/src/borrow.rs

@@ -881,7 +881,7 @@ pub fn lowlevel_borrow_signature(arena: &Bump, op: LowLevel) -> &[bool] {
         Unreachable => arena.alloc_slice_copy(&[irrelevant]),
         ListLen | StrIsEmpty | StrToScalars | StrCountGraphemes | StrCountUtf8Bytes
         | StrGetCapacity | ListGetCapacity => arena.alloc_slice_copy(&[borrowed]),
-        ListWithCapacity => arena.alloc_slice_copy(&[irrelevant]),
+        ListWithCapacity | StrWithCapacity => arena.alloc_slice_copy(&[irrelevant]),
         ListReplaceUnsafe => arena.alloc_slice_copy(&[owned, irrelevant, irrelevant]),
         StrGetUnsafe | ListGetUnsafe => arena.alloc_slice_copy(&[borrowed, irrelevant]),
         ListConcat => arena.alloc_slice_copy(&[owned, owned]),

crates/compiler/test_gen/src/gen_str.rs

@@ -1930,3 +1930,31 @@ fn when_on_strings() {
         i64
     );
 }
+
+#[test]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
+fn with_capacity() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            Str.withCapacity 10
+            "#
+        ),
+        RocStr::from(""),
+        RocStr
+    );
+}
+
+#[test]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm"))]
+fn with_capacity_concat() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            Str.withCapacity 10 |> Str.concat "Forty-two"
+            "#
+        ),
+        RocStr::from("Forty-two"),
+        RocStr
+    );
+}