update

Author: waralexrom

rust/cubenativeutils/src/wrappers/neon/context.rs

@@ -1,4 +1,3 @@
-//use super::object::NeonObject;
 use super::{
     inner_types::NeonInnerTypes,
     object::{
@@ -7,12 +6,15 @@ use super::{
     },
 };
 use crate::wrappers::{
-    context::NativeContext, object::NativeObject, object_handle::NativeObjectHandle,
+    context::{NativeContext, NativeContextHolder},
+    object::NativeObject,
+    object_handle::NativeObjectHandle,
 };
 use cubesql::CubeError;
 use neon::prelude::*;
 use std::{
-    cell::RefCell,
+    cell::{RefCell, RefMut},
+    marker::PhantomData,
     panic::{catch_unwind, resume_unwind, AssertUnwindSafe},
     rc::{Rc, Weak},
 };
@@ -50,18 +52,42 @@ impl<C: Context<'static>> ContextWrapper<C> {
     }
 }
 
+pub struct NeonContextGuard<'cx, C: Context<'cx> + NoenContextLifetimeExpand<'cx>> {
+    context: Rc<RefCell<ContextWrapper<C::ExpandedResult>>>,
+    lifetime: PhantomData<&'cx ()>,
+}
+
+impl<'cx, C: Context<'cx> + NoenContextLifetimeExpand<'cx> + 'cx> NeonContextGuard<'cx, C> {
+    fn new(cx: C) -> Self {
+        Self {
+            context: ContextWrapper::new(cx.expand_lifetime()),
+            lifetime: PhantomData::default(),
+        }
+    }
+
+    fn context_holder(&self) -> ContextHolder<C::ExpandedResult> {
+        ContextHolder::new(Rc::downgrade(&self.context))
+    }
+
+    fn unwrap(self) {
+        if Rc::strong_count(&self.context) > 0 {
+            match Rc::try_unwrap(self.context) {
+                Ok(_) => {}
+                Err(_) => panic!("Guarded context have more then one reference"),
+            }
+        }
+    }
+}
+
 pub fn neon_run_with_guarded_lifetime<'cx, C, T, F>(cx: C, func: F) -> T
 where
-    C: Context<'cx> + NoenContextLifetimeExpand<'cx>,
+    C: Context<'cx> + NoenContextLifetimeExpand<'cx> + 'cx,
     F: FnOnce(ContextHolder<C::ExpandedResult>) -> T,
 {
-    let context = ContextWrapper::new(cx.expand_lifetime());
-    let context_holder = ContextHolder::new(Rc::downgrade(&context));
+    let guard = NeonContextGuard::new(cx);
+    let context_holder = guard.context_holder();
     let res = catch_unwind(AssertUnwindSafe(|| func(context_holder)));
-    match Rc::try_unwrap(context) {
-        Ok(_) => {}
-        Err(_) => panic!("Guarded context have more then one reference"),
-    };
+    guard.unwrap();
 
     match res {
         Ok(res) => res,