x

Author: soburi

rust/src/string.rs

@@ -295,6 +295,139 @@ unsafe fn write_replaced_bytes(
     }
 }
 
+enum ReplaceBytesFlow {
+    Return(bool),
+    Proceed {
+        current_len: c_uint,
+        hay_len: usize,
+        needle_len: usize,
+        replace_len: usize,
+    },
+}
+
+struct ReplaceBytesScratch {
+    ptr: *mut c_char,
+    written: usize,
+    new_len: c_uint,
+}
+
+#[inline]
+fn validate_replace_bytes(
+    buffer: *mut c_char,
+    current_len: c_uint,
+    find: *const c_char,
+    find_len: c_uint,
+    replace: *const c_char,
+    replace_len: c_uint,
+) -> ReplaceBytesFlow {
+    if current_len == 0 || find_len == 0 {
+        return ReplaceBytesFlow::Return(true);
+    }
+    if buffer.is_null() || find.is_null() {
+        return ReplaceBytesFlow::Return(false);
+    }
+    if replace.is_null() && replace_len != 0 {
+        return ReplaceBytesFlow::Return(false);
+    }
+
+    let hay_len = current_len as usize;
+    let needle_len = find_len as usize;
+    if needle_len > hay_len {
+        return ReplaceBytesFlow::Return(true);
+    }
+
+    ReplaceBytesFlow::Proceed {
+        current_len,
+        hay_len,
+        needle_len,
+        replace_len: replace_len as usize,
+    }
+}
+
+#[inline]
+fn prepare_replace_scratch(
+    buffer: *const c_char,
+    find: *const c_char,
+    replace: *const c_char,
+    flow: &ReplaceBytesFlow,
+) -> Result<Option<ReplaceBytesScratch>, ()> {
+    let ReplaceBytesFlow::Proceed {
+        current_len,
+        hay_len,
+        needle_len,
+        replace_len,
+    } = flow
+    else {
+        return Err(());
+    };
+
+    // SAFETY: `flow` is validated by `validate_replace_bytes`, and pointers/lengths here are coherent.
+    unsafe {
+        let match_count = count_non_overlapping_matches(buffer, *hay_len, find, *needle_len);
+        if match_count == 0 {
+            return Ok(None);
+        }
+
+        let Some(new_len) = compute_replaced_len(
+            *current_len,
+            *needle_len as c_uint,
+            *replace_len as c_uint,
+            match_count,
+        ) else {
+            return Err(());
+        };
+
+        let scratch_size = (new_len as usize).saturating_add(1);
+        let scratch_ptr = malloc(scratch_size) as *mut c_char;
+        if scratch_ptr.is_null() {
+            return Err(());
+        }
+
+        let written = write_replaced_bytes(
+            buffer,
+            *hay_len,
+            find,
+            *needle_len,
+            replace,
+            *replace_len,
+            scratch_ptr,
+            scratch_size,
+        );
+
+        Ok(Some(ReplaceBytesScratch {
+            ptr: scratch_ptr,
+            written,
+            new_len,
+        }))
+    }
+}
+
+#[inline]
+fn commit_replace_scratch(
+    buffer: &mut *mut c_char,
+    capacity: &mut c_uint,
+    len: &mut c_uint,
+    scratch: ReplaceBytesScratch,
+) -> bool {
+    if !ensure_capacity(buffer, capacity, len, scratch.new_len) {
+        // SAFETY: pointer came from malloc in `prepare_replace_scratch`.
+        unsafe { free(scratch.ptr.cast()) }
+        return false;
+    }
+
+    // SAFETY: destination has enough capacity; source points to scratch allocation.
+    unsafe {
+        ptr::copy_nonoverlapping(
+            scratch.ptr.cast::<u8>(),
+            (*buffer).cast::<u8>(),
+            scratch.written + 1,
+        );
+        free(scratch.ptr.cast());
+    }
+    *len = scratch.new_len;
+    true
+}
+
 #[inline]
 fn ensure_capacity(
     buffer: &mut *mut c_char,
@@ -1133,71 +1266,16 @@ pub extern "C" fn arduino_string_replace_bytes(
         return false;
     };
 
-    let current_len = *len;
-    if current_len == 0 || find_len == 0 {
-        return true;
-    }
-    if (*buffer).is_null() || find.is_null() {
-        return false;
-    }
-    if replace.is_null() && replace_len != 0 {
-        return false;
+    let flow = validate_replace_bytes(*buffer, *len, find, find_len, replace, replace_len);
+    if let ReplaceBytesFlow::Return(ok) = &flow {
+        return *ok;
     }
 
-    let hay_len = current_len as usize;
-    let needle_len = find_len as usize;
-    if needle_len > hay_len {
-        return true;
+    match prepare_replace_scratch(*buffer, find, replace, &flow) {
+        Ok(None) => true,
+        Ok(Some(scratch)) => commit_replace_scratch(buffer, capacity, len, scratch),
+        Err(()) => false,
     }
-
-    // SAFETY: inputs were validated above.
-    let match_count = unsafe { count_non_overlapping_matches(*buffer, hay_len, find, needle_len) };
-    if match_count == 0 {
-        return true;
-    }
-
-    let Some(new_len) = compute_replaced_len(current_len, find_len, replace_len, match_count)
-    else {
-        return false;
-    };
-    let tmp_size = (new_len as usize).saturating_add(1);
-    // SAFETY: malloc returns either null or writable memory.
-    let tmp = unsafe { malloc(tmp_size) as *mut c_char };
-    if tmp.is_null() {
-        return false;
-    }
-
-    // SAFETY: inputs and destination capacity are validated above.
-    let written = unsafe {
-        write_replaced_bytes(
-            *buffer,
-            hay_len,
-            find,
-            needle_len,
-            replace,
-            replace_len as usize,
-            tmp,
-            tmp_size,
-        )
-    };
-
-    let ok = if !ensure_capacity(buffer, capacity, len, new_len) {
-        false
-    } else {
-        // SAFETY: destination has at least written+1 bytes.
-        unsafe {
-            ptr::copy_nonoverlapping(tmp.cast::<u8>(), (*buffer).cast::<u8>(), written + 1);
-        }
-        *len = new_len;
-        true
-    };
-
-    // SAFETY: tmp was allocated above with malloc.
-    unsafe {
-        free(tmp.cast());
-    }
-
-    ok
 }
 
 #[unsafe(no_mangle)]