-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathjava_function.rs
More file actions
119 lines (110 loc) · 4.04 KB
/
java_function.rs
File metadata and controls
119 lines (110 loc) · 4.04 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
use jni::{
objects::{GlobalRef, JList, JMethodID, JObject, JValue},
sys::jvalue,
JNIEnv,
};
use crate::{get_jvm, to_jvalue, to_string, to_value};
#[derive(Clone)]
pub struct JavaFunction {
pub obj: GlobalRef,
pub method_id: JMethodID,
pub arity: Option<usize>,
}
impl JavaFunction {
pub fn new(env: &JNIEnv, obj: JObject) -> Result<Self, piper::PiperError> {
let obj = env.new_global_ref(obj).unwrap();
// Check from Function0 to Function'MAX_ARITY'
for a in 0..get_jvm().max_arity {
let cls = env
.find_class(format!("com/linkedin/feathr/online/Function{a}"))
.map_err(|e| piper::PiperError::ExternalError(e.to_string()))?;
if env
.is_instance_of(&obj, cls)
.map_err(|e| piper::PiperError::ExternalError(e.to_string()))?
{
let method_id = env
.get_method_id(
cls,
format!("apply{a}"),
format!("({})Ljava/lang/Object;", "Ljava/lang/Object;".repeat(a)),
)
.map_err(|e| piper::PiperError::ExternalError(e.to_string()))?;
return Ok(Self {
obj,
method_id,
arity: Some(a),
});
}
}
// Assume VarFunction
let cls = env
.find_class("com/azure/feathr/piper/VarFunction")
.map_err(|e| piper::PiperError::ExternalError(e.to_string()))?;
let method_id = env
.get_method_id(cls, "applyVar", "([Ljava/lang/Object;)Ljava/lang/Object;")
.map_err(|e| piper::PiperError::ExternalError(e.to_string()))?;
Ok(Self {
obj,
method_id,
arity: None,
})
}
}
impl piper::Function for JavaFunction {
fn get_output_type(
&self,
_argument_types: &[piper::ValueType],
) -> Result<piper::ValueType, piper::PiperError> {
Ok(piper::ValueType::Dynamic)
}
fn eval(&self, arguments: Vec<piper::Value>) -> piper::Value {
let env = match get_jvm().jvm.attach_current_thread_as_daemon() {
Ok(env) => env,
Err(e) => return piper::Value::Error(piper::PiperError::ExternalError(e.to_string())),
};
let args: Vec<jvalue> = match self.arity {
None => {
// call applyVar
let array_list_cls = &get_jvm().array_list_cls;
let new_array_list = get_jvm().new_array_list;
let l = env
.new_object_unchecked(array_list_cls, new_array_list, &[])
.unwrap();
let j = JList::from_env(&env, l).unwrap();
for arg in arguments {
j.add(to_jvalue(arg, &env)).unwrap();
}
let o: JObject = j.into();
vec![JValue::Object(o).to_jni()]
}
Some(arity) => arguments
.into_iter()
.take(arity)
.map(|a| JValue::Object(to_jvalue(a, &env)).to_jni())
.collect(),
};
let ret = env
.call_method_unchecked(
self.obj.as_obj(),
self.method_id,
jni::signature::ReturnType::Object,
&args,
)
.map_err(|e| piper::PiperError::ExternalError(e.to_string()))
.and_then(|v| {
to_value(v, &env).map_err(|e| piper::PiperError::ExternalError(e.to_string()))
});
if ret.is_err() {
match env.exception_occurred() {
Ok(ex) => {
if !ex.is_null() {
env.exception_clear().unwrap();
return piper::PiperError::ExternalError(to_string(&env, ex.into())).into();
}
}
Err(e) => return piper::PiperError::ExternalError(e.to_string()).into(),
};
}
ret.into()
}
}