Skip to main content

oxmpl_js/base/
js_state_convert.rs

1// Copyright (c) 2025 Ross Gardiner, Junior Sundar
2//
3// SPDX-License-Identifier: BSD-3-Clause
4
5use std::sync::Arc;
6
7use js_sys::Float64Array;
8use oxmpl::base::state::{
9    CompoundState, RealVectorState, SE2State, SE3State, SO2State, SO3State, State,
10};
11use wasm_bindgen::prelude::*;
12
13use crate::base::{
14    JsCompoundState, JsRealVectorState, JsSE2State, JsSE3State, JsSO2State, JsSO3State,
15};
16
17/// Trait for converting between Rust states and their JavaScript wrappers.
18pub trait JsStateConvert: Sized {
19    /// Convert the Rust state to a JavaScript wrapper object.
20    fn to_js_value(&self) -> JsValue;
21    /// Convert a JavaScript wrapper object back to the Rust state.
22    fn from_js_value(val: JsValue) -> Result<Self, String>;
23}
24
25impl JsStateConvert for RealVectorState {
26    fn to_js_value(&self) -> JsValue {
27        JsValue::from(JsRealVectorState {
28            inner: Arc::new(self.clone()),
29        })
30    }
31    fn from_js_value(val: JsValue) -> Result<Self, String> {
32        if let Ok(values_val) = js_sys::Reflect::get(&val, &JsValue::from_str("values")) {
33            if let Ok(values) = serde_wasm_bindgen::from_value::<Vec<f64>>(values_val) {
34                return Ok(RealVectorState::new(values));
35            }
36        }
37        Err("Expected RealVectorState or object with 'values' property".to_string())
38    }
39}
40
41impl JsStateConvert for SO2State {
42    fn to_js_value(&self) -> JsValue {
43        JsValue::from(JsSO2State {
44            inner: Arc::new(self.clone()),
45        })
46    }
47    fn from_js_value(val: JsValue) -> Result<Self, String> {
48        if let Ok(v_val) = js_sys::Reflect::get(&val, &JsValue::from_str("value")) {
49            if let Ok(v) = v_val.as_f64().ok_or("Expected number for SO2 value") {
50                return Ok(SO2State::new(v));
51            }
52        }
53        if let Some(v) = val.as_f64() {
54            return Ok(SO2State::new(v));
55        }
56        Err("Expected SO2State, object with 'value' property, or number".to_string())
57    }
58}
59
60impl JsStateConvert for SO3State {
61    fn to_js_value(&self) -> JsValue {
62        JsValue::from(JsSO3State {
63            inner: Arc::new(self.clone()),
64        })
65    }
66    fn from_js_value(val: JsValue) -> Result<Self, String> {
67        let x = js_sys::Reflect::get(&val, &JsValue::from_str("x"))
68            .ok()
69            .and_then(|v| v.as_f64());
70        let y = js_sys::Reflect::get(&val, &JsValue::from_str("y"))
71            .ok()
72            .and_then(|v| v.as_f64());
73        let z = js_sys::Reflect::get(&val, &JsValue::from_str("z"))
74            .ok()
75            .and_then(|v| v.as_f64());
76        let w = js_sys::Reflect::get(&val, &JsValue::from_str("w"))
77            .ok()
78            .and_then(|v| v.as_f64());
79
80        if let (Some(x), Some(y), Some(z), Some(w)) = (x, y, z, w) {
81            Ok(SO3State::new(x, y, z, w))
82        } else {
83            Err("Expected SO3State or object with x, y, z, w properties".to_string())
84        }
85    }
86}
87
88impl JsStateConvert for SE2State {
89    fn to_js_value(&self) -> JsValue {
90        JsValue::from(JsSE2State {
91            inner: Arc::new(self.clone()),
92        })
93    }
94    fn from_js_value(val: JsValue) -> Result<Self, String> {
95        let x = js_sys::Reflect::get(&val, &JsValue::from_str("x"))
96            .ok()
97            .and_then(|v| v.as_f64());
98        let y = js_sys::Reflect::get(&val, &JsValue::from_str("y"))
99            .ok()
100            .and_then(|v| v.as_f64());
101        let yaw = js_sys::Reflect::get(&val, &JsValue::from_str("yaw"))
102            .ok()
103            .and_then(|v| v.as_f64());
104
105        if let (Some(x), Some(y), Some(yaw)) = (x, y, yaw) {
106            Ok(SE2State::new(x, y, yaw))
107        } else {
108            Err("Expected SE2State or object with x, y, yaw properties".to_string())
109        }
110    }
111}
112
113impl JsStateConvert for SE3State {
114    fn to_js_value(&self) -> JsValue {
115        JsValue::from(JsSE3State {
116            inner: Arc::new(self.clone()),
117        })
118    }
119    fn from_js_value(val: JsValue) -> Result<Self, String> {
120        let x = js_sys::Reflect::get(&val, &JsValue::from_str("x"))
121            .ok()
122            .and_then(|v| v.as_f64());
123        let y = js_sys::Reflect::get(&val, &JsValue::from_str("y"))
124            .ok()
125            .and_then(|v| v.as_f64());
126        let z = js_sys::Reflect::get(&val, &JsValue::from_str("z"))
127            .ok()
128            .and_then(|v| v.as_f64());
129        let rotation_val = js_sys::Reflect::get(&val, &JsValue::from_str("rotation")).ok();
130
131        if let (Some(x), Some(y), Some(z), Some(rot_val)) = (x, y, z, rotation_val) {
132            let rotation = SO3State::from_js_value(rot_val)?;
133            Ok(SE3State::new(x, y, z, rotation))
134        } else {
135            Err("Expected SE3State or object with x, y, z, and rotation properties".to_string())
136        }
137    }
138}
139
140impl JsStateConvert for CompoundState {
141    fn to_js_value(&self) -> JsValue {
142        JsValue::from(JsCompoundState {
143            inner: Arc::new(self.clone()),
144        })
145    }
146    fn from_js_value(val: JsValue) -> Result<Self, String> {
147        let count_res = js_sys::Reflect::get(&val, &JsValue::from_str("componentCount"));
148        if let Ok(c) = count_res {
149            if let Some(count) = c.as_f64() {
150                let count = count as usize;
151                let mut components = Vec::with_capacity(count);
152
153                // Get getComponent function
154                let get_comp_fn_val =
155                    js_sys::Reflect::get(&val, &JsValue::from_str("getComponent"))
156                        .map_err(|_| "Missing getComponent")?;
157                let get_comp_fn = get_comp_fn_val
158                    .dyn_into::<js_sys::Function>()
159                    .map_err(|_| "getComponent is not a function")?;
160
161                for i in 0..count {
162                    let idx = JsValue::from(i as u32);
163                    let comp_val = get_comp_fn
164                        .call1(&val, &idx)
165                        .map_err(|e| format!("getComponent failed: {:?}", e))?;
166
167                    // Infer type of comp_val
168                    let comp = infer_and_convert_state(comp_val)?;
169                    components.push(comp);
170                }
171                return Ok(CompoundState::new(components));
172            }
173        }
174        Err("Expected CompoundState or object with componentCount".to_string())
175    }
176}
177
178fn infer_and_convert_state(val: JsValue) -> Result<Box<dyn State>, String> {
179    // Try SE3 (has rotation)
180    if let Ok(v) = js_sys::Reflect::get(&val, &JsValue::from_str("rotation")) {
181        if !v.is_undefined() {
182            let s = SE3State::from_js_value(val)?;
183            return Ok(Box::new(s));
184        }
185    }
186    // Try SE2 (has yaw)
187    if let Ok(v) = js_sys::Reflect::get(&val, &JsValue::from_str("yaw")) {
188        if !v.is_undefined() {
189            let s = SE2State::from_js_value(val)?;
190            return Ok(Box::new(s));
191        }
192    }
193    // Try SO3 (has w)
194    if let Ok(v) = js_sys::Reflect::get(&val, &JsValue::from_str("w")) {
195        if !v.is_undefined() {
196            let s = SO3State::from_js_value(val)?;
197            return Ok(Box::new(s));
198        }
199    }
200    // Try RealVector (has values)
201    if let Ok(v) = js_sys::Reflect::get(&val, &JsValue::from_str("values")) {
202        if !v.is_undefined() {
203            let s = RealVectorState::from_js_value(val)?;
204            return Ok(Box::new(s));
205        }
206    }
207    // Try SO2 (has value) or is number
208    if val.as_f64().is_some() {
209        let s = SO2State::from_js_value(val)?;
210        return Ok(Box::new(s));
211    }
212    if let Ok(v) = js_sys::Reflect::get(&val, &JsValue::from_str("value")) {
213        if !v.is_undefined() {
214            let s = SO2State::from_js_value(val)?;
215            return Ok(Box::new(s));
216        }
217    }
218    // Try Compound (recursive)
219    if let Ok(v) = js_sys::Reflect::get(&val, &JsValue::from_str("componentCount")) {
220        if !v.is_undefined() {
221            let s = CompoundState::from_js_value(val)?;
222            return Ok(Box::new(s));
223        }
224    }
225
226    Err("Could not infer state type from JS object".to_string())
227}
228
229pub fn real_vector_state_to_js_array(state: &RealVectorState) -> Float64Array {
230    let array = Float64Array::new_with_length(state.values.len() as u32);
231    for (i, &val) in state.values.iter().enumerate() {
232        array.set_index(i as u32, val);
233    }
234    array
235}
236
237pub fn js_array_to_real_vector_state(array: &Float64Array) -> RealVectorState {
238    let mut values = Vec::new();
239    for i in 0..array.length() {
240        values.push(array.get_index(i));
241    }
242    RealVectorState::new(values)
243}
244
245pub fn so2_state_to_js_array(state: &SO2State) -> Float64Array {
246    let array = Float64Array::new_with_length(1);
247    array.set_index(0, state.value);
248    array
249}
250
251pub fn js_array_to_so2_state(array: &Float64Array) -> SO2State {
252    let val = if array.length() > 0 {
253        array.get_index(0)
254    } else {
255        0.0
256    };
257    SO2State::new(val)
258}
259
260pub fn se2_state_to_js_array(state: &SE2State) -> Float64Array {
261    let array = Float64Array::new_with_length(3);
262    array.set_index(0, state.get_x());
263    array.set_index(1, state.get_y());
264    array.set_index(2, state.get_yaw());
265    array
266}
267
268pub fn js_array_to_se2_state(array: &Float64Array) -> SE2State {
269    let x = if array.length() > 0 {
270        array.get_index(0)
271    } else {
272        0.0
273    };
274    let y = if array.length() > 1 {
275        array.get_index(1)
276    } else {
277        0.0
278    };
279    let yaw = if array.length() > 2 {
280        array.get_index(2)
281    } else {
282        0.0
283    };
284    SE2State::new(x, y, yaw)
285}
286
287pub fn so3_state_to_js_array(state: &SO3State) -> Float64Array {
288    let array = Float64Array::new_with_length(4);
289    array.set_index(0, state.x);
290    array.set_index(1, state.y);
291    array.set_index(2, state.z);
292    array.set_index(3, state.w);
293    array
294}
295
296pub fn js_array_to_so3_state(array: &Float64Array) -> SO3State {
297    let x = if array.length() > 0 {
298        array.get_index(0)
299    } else {
300        0.0
301    };
302    let y = if array.length() > 1 {
303        array.get_index(1)
304    } else {
305        0.0
306    };
307    let z = if array.length() > 2 {
308        array.get_index(2)
309    } else {
310        0.0
311    };
312    let w = if array.length() > 3 {
313        array.get_index(3)
314    } else {
315        1.0
316    };
317    SO3State::new(x, y, z, w)
318}
319
320pub fn se3_state_to_js_array(state: &SE3State) -> Float64Array {
321    let array = Float64Array::new_with_length(7);
322    array.set_index(0, state.get_x());
323    array.set_index(1, state.get_y());
324    array.set_index(2, state.get_z());
325    array.set_index(3, state.get_rotation().x);
326    array.set_index(4, state.get_rotation().y);
327    array.set_index(5, state.get_rotation().z);
328    array.set_index(6, state.get_rotation().w);
329    array
330}
331
332pub fn js_array_to_se3_state(array: &Float64Array) -> SE3State {
333    let x = if array.length() > 0 {
334        array.get_index(0)
335    } else {
336        0.0
337    };
338    let y = if array.length() > 1 {
339        array.get_index(1)
340    } else {
341        0.0
342    };
343    let z = if array.length() > 2 {
344        array.get_index(2)
345    } else {
346        0.0
347    };
348    let qx = if array.length() > 3 {
349        array.get_index(3)
350    } else {
351        0.0
352    };
353    let qy = if array.length() > 4 {
354        array.get_index(4)
355    } else {
356        0.0
357    };
358    let qz = if array.length() > 5 {
359        array.get_index(5)
360    } else {
361        0.0
362    };
363    let qw = if array.length() > 6 {
364        array.get_index(6)
365    } else {
366        1.0
367    };
368    SE3State::new(x, y, z, SO3State::new(qx, qy, qz, qw))
369}
370
371pub fn compound_state_to_js_array(state: &CompoundState) -> Float64Array {
372    let mut values = Vec::new();
373    // Helper to recursively flatten
374    fn flatten(s: &dyn State, out: &mut Vec<f64>) {
375        let any_s = s.as_any();
376        if let Some(rv) = any_s.downcast_ref::<RealVectorState>() {
377            out.extend_from_slice(&rv.values);
378        } else if let Some(so2) = any_s.downcast_ref::<SO2State>() {
379            out.push(so2.value);
380        } else if let Some(se2) = any_s.downcast_ref::<SE2State>() {
381            out.push(se2.get_x());
382            out.push(se2.get_y());
383            out.push(se2.get_yaw());
384        } else if let Some(so3) = any_s.downcast_ref::<SO3State>() {
385            out.push(so3.x);
386            out.push(so3.y);
387            out.push(so3.z);
388            out.push(so3.w);
389        } else if let Some(se3) = any_s.downcast_ref::<SE3State>() {
390            out.push(se3.get_x());
391            out.push(se3.get_y());
392            out.push(se3.get_z());
393            out.push(se3.get_rotation().x);
394            out.push(se3.get_rotation().y);
395            out.push(se3.get_rotation().z);
396            out.push(se3.get_rotation().w);
397        } else if let Some(comp) = any_s.downcast_ref::<CompoundState>() {
398            for c in &comp.components {
399                flatten(c.as_ref(), out);
400            }
401        }
402    }
403
404    for component in &state.components {
405        flatten(component.as_ref(), &mut values);
406    }
407
408    let array = Float64Array::new_with_length(values.len() as u32);
409    for (i, &val) in values.iter().enumerate() {
410        array.set_index(i as u32, val);
411    }
412    array
413}
414
415pub fn state_to_js_array(state: &RealVectorState) -> Float64Array {
416    real_vector_state_to_js_array(state)
417}
418
419pub fn js_array_to_state(array: &Float64Array) -> RealVectorState {
420    js_array_to_real_vector_state(array)
421}
422
423#[wasm_bindgen(start)]
424pub fn set_panic_hook() {
425    console_error_panic_hook::set_once();
426}
427
428#[wasm_bindgen]
429extern "C" {
430    #[wasm_bindgen(js_namespace = console)]
431    fn log(s: &str);
432}