Let Swift allocate as much memory as it needs and then transfer ownership to Rust

highlight-swift/Sources/highlight-swift/highlight_swift.swift

@@ -2,17 +2,18 @@ import Foundation
 import Splash
 
 @_cdecl("highlight_swift")
-public func highlight(codePtr: UnsafeRawPointer, codeLen: UInt64, outPtr: UnsafeMutableRawPointer, maxLen: UInt64) -> UInt64 {
+public func highlight(codePtr: UnsafeRawPointer, codeLen: UInt64, htmlLenPtr: UnsafeMutablePointer<UInt64>) -> UnsafeMutableRawPointer {
     // don't free, the underlying data is owned by rust
     let code = String(bytesNoCopy: UnsafeMutableRawPointer(mutating: codePtr), length: Int(codeLen), encoding: .utf8, freeWhenDone: false)!
 
     let highligher = SyntaxHighlighter(format: MyOutputFormat())
     var html = highligher.highlight(code)
-    precondition(html.utf8.count <= maxLen)
+    let outPtr = UnsafeMutableBufferPointer<UInt8>.allocate(capacity: html.utf8.count)
-    return html.withUTF8 { buf in
+    _ = html.withUTF8 { buf in
-        buf.copyBytes(to: UnsafeMutableRawBufferPointer(start: outPtr, count: Int(maxLen)))
+        buf.copyBytes(to: outPtr, count: buf.count)
-        return UInt64(buf.count)
     }
+    htmlLenPtr.pointee = UInt64(html.utf8.count)
+    return UnsafeMutableRawPointer(outPtr.baseAddress!)
 }
 
 struct MyOutputFormat: OutputFormat {

src/lib.rs

@@ -1,24 +1,13 @@
 pub fn highlight(code: &str) -> String {
     unsafe {
-        let mut buf = Vec::new();
+        let mut html_len: u64 = 0;
-        buf.reserve(
+        let html_ptr = highlight_swift(code.as_ptr(), code.len() as u64, &mut html_len);
-            2usize
+        String::from_raw_parts(html_ptr, html_len as usize, html_len as usize)
-                .pow(1 + (code.len() as f32).log2().ceil() as u32)
-                .max(1024),
-        );
-        let used = highlight_swift(
-            code.as_ptr(),
-            code.len() as u64,
-            buf.as_mut_ptr(),
-            buf.capacity() as u64,
-        );
-        buf.set_len(used as usize);
-        String::from_utf8_unchecked(buf)
     }
 }
 
 extern "C" {
-    fn highlight_swift(code_ptr: *const u8, code_len: u64, out_ptr: *mut u8, max_len: u64) -> u64;
+    fn highlight_swift(code_ptr: *const u8, code_len: u64, html_len_ptr: *mut u64) -> *mut u8;
 }
 
 #[cfg(test)]