onnxruntime/js/react_native/ios/OnnxruntimeJSIHelper.mm
Hans 2a17958b34
[js/rn] Fix some bugs (#20242)
### Description
<!-- Describe your changes. -->
- Fix `logSeverityLevel`
- Correct get RCTCxxBridge, old method for some cases will got wrong
bridge


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

---------

Co-authored-by: Yulong Wang <7679871+fs-eire@users.noreply.github.com>
2024-05-15 10:32:08 -07:00

90 lines
3.5 KiB
Text

#import "OnnxruntimeJSIHelper.h"
#import <React/RCTBlobManager.h>
#import <React/RCTBridge+Private.h>
#import <jsi/jsi.h>
@implementation OnnxruntimeJSIHelper
@synthesize bridge = _bridge;
RCT_EXPORT_MODULE()
- (void)setBridge:(RCTBridge *)bridge {
_bridge = bridge;
}
RCT_EXPORT_BLOCKING_SYNCHRONOUS_METHOD(install) {
RCTCxxBridge *cxxBridge = (RCTCxxBridge *)_bridge;
if (cxxBridge == nil) {
return @false;
}
using namespace facebook;
auto jsiRuntime = (jsi::Runtime *)cxxBridge.runtime;
if (jsiRuntime == nil) {
return @false;
}
auto &runtime = *jsiRuntime;
auto resolveArrayBuffer = jsi::Function::createFromHostFunction(
runtime, jsi::PropNameID::forUtf8(runtime, "jsiOnnxruntimeResolveArrayBuffer"), 1,
[](jsi::Runtime &runtime, const jsi::Value &thisArg, const jsi::Value *args, size_t count) -> jsi::Value {
if (count != 1) {
throw jsi::JSError(runtime, "jsiOnnxruntimeResolveArrayBuffer(..) expects one argument (object)!");
}
auto data = args[0].asObject(runtime);
auto blobId = data.getProperty(runtime, "blobId").asString(runtime).utf8(runtime);
auto size = data.getProperty(runtime, "size").asNumber();
auto offset = data.getProperty(runtime, "offset").asNumber();
RCTBlobManager *blobManager = [[RCTBridge currentBridge] moduleForClass:RCTBlobManager.class];
if (blobManager == nil) {
throw jsi::JSError(runtime, "RCTBlobManager is not initialized");
}
NSString *blobIdStr = [NSString stringWithUTF8String:blobId.c_str()];
auto blob = [blobManager resolve:blobIdStr offset:(long)offset size:(long)size];
jsi::Function arrayBufferCtor = runtime.global().getPropertyAsFunction(runtime, "ArrayBuffer");
jsi::Object o = arrayBufferCtor.callAsConstructor(runtime, (int)blob.length).getObject(runtime);
jsi::ArrayBuffer buf = o.getArrayBuffer(runtime);
memcpy(buf.data(runtime), blob.bytes, blob.length);
[blobManager remove:blobIdStr];
return buf;
});
runtime.global().setProperty(runtime, "jsiOnnxruntimeResolveArrayBuffer", resolveArrayBuffer);
auto storeArrayBuffer = jsi::Function::createFromHostFunction(
runtime, jsi::PropNameID::forUtf8(runtime, "jsiOnnxruntimeStoreArrayBuffer"), 1,
[](jsi::Runtime &runtime, const jsi::Value &thisArg, const jsi::Value *args, size_t count) -> jsi::Value {
if (count != 1) {
throw jsi::JSError(runtime, "jsiOnnxruntimeStoreArrayBuffer(..) expects one argument (object)!");
}
auto arrayBuffer = args[0].asObject(runtime).getArrayBuffer(runtime);
auto size = arrayBuffer.length(runtime);
NSData *data = [NSData dataWithBytesNoCopy:arrayBuffer.data(runtime) length:size freeWhenDone:NO];
RCTBlobManager *blobManager = [[RCTBridge currentBridge] moduleForClass:RCTBlobManager.class];
if (blobManager == nil) {
throw jsi::JSError(runtime, "RCTBlobManager is not initialized");
}
NSString *blobId = [blobManager store:data];
jsi::Object result(runtime);
auto blobIdString = jsi::String::createFromUtf8(runtime, [blobId cStringUsingEncoding:NSUTF8StringEncoding]);
result.setProperty(runtime, "blobId", blobIdString);
result.setProperty(runtime, "offset", jsi::Value(0));
result.setProperty(runtime, "size", jsi::Value(static_cast<double>(size)));
return result;
});
runtime.global().setProperty(runtime, "jsiOnnxruntimeStoreArrayBuffer", storeArrayBuffer);
return @true;
}
@end