From 20d54805342687e3600800f3b1c9c358040c5b2d Mon Sep 17 00:00:00 2001
From: Shadowfacts <me@shadowfacts.net>
Date: Wed, 17 Aug 2022 21:19:53 -0400
Subject: [PATCH] Let Swift allocate as much memory as it needs and then
 transfer ownership to Rust

---
 .../highlight-swift/highlight_swift.swift     | 11 ++++++-----
 src/lib.rs                                    | 19 ++++---------------
 2 files changed, 10 insertions(+), 20 deletions(-)

diff --git a/highlight-swift/Sources/highlight-swift/highlight_swift.swift b/highlight-swift/Sources/highlight-swift/highlight_swift.swift
index 2ed3c7e..cd12e2c 100644
--- a/highlight-swift/Sources/highlight-swift/highlight_swift.swift
+++ b/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)
-    return html.withUTF8 { buf in
-        buf.copyBytes(to: UnsafeMutableRawBufferPointer(start: outPtr, count: Int(maxLen)))
-        return UInt64(buf.count)
+    let outPtr = UnsafeMutableBufferPointer<UInt8>.allocate(capacity: html.utf8.count)
+    _ = html.withUTF8 { buf in
+        buf.copyBytes(to: outPtr, count: buf.count)
     }
+    htmlLenPtr.pointee = UInt64(html.utf8.count)
+    return UnsafeMutableRawPointer(outPtr.baseAddress!)
 }
 
 struct MyOutputFormat: OutputFormat {
diff --git a/src/lib.rs b/src/lib.rs
index 7454119..f11f602 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,24 +1,13 @@
 pub fn highlight(code: &str) -> String {
     unsafe {
-        let mut buf = Vec::new();
-        buf.reserve(
-            2usize
-                .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)
+        let mut html_len: u64 = 0;
+        let html_ptr = highlight_swift(code.as_ptr(), code.len() as u64, &mut html_len);
+        String::from_raw_parts(html_ptr, html_len as usize, html_len as usize)
     }
 }
 
 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)]