From 889b75bf902237ddf039935684f86e3d7a692580 Mon Sep 17 00:00:00 2001
From: "K.J. Valencik" <kjvalencik@gmail.com>
Date: Wed, 24 Jul 2024 11:27:41 -0400
Subject: [PATCH] Add methods to get Cx to ContextInternal

---
 crates/neon/src/context/internal.rs        |  8 +++++--
 crates/neon/src/context/mod.rs             | 28 +++++++++++++++-------
 crates/neon/src/types_impl/extract/with.rs | 24 +++++++++++++++++++
 3 files changed, 50 insertions(+), 10 deletions(-)
 create mode 100644 crates/neon/src/types_impl/extract/with.rs

diff --git a/crates/neon/src/context/internal.rs b/crates/neon/src/context/internal.rs
index bf872d2e5..52164a5d6 100644
--- a/crates/neon/src/context/internal.rs
+++ b/crates/neon/src/context/internal.rs
@@ -1,7 +1,7 @@
 use std::{cell::RefCell, ffi::c_void, mem::MaybeUninit};
 
 use crate::{
-    context::ModuleContext,
+    context::{Cx, ModuleContext},
     handle::Handle,
     result::NeonResult,
     sys::{self, raw},
@@ -47,7 +47,11 @@ impl Env {
 }
 
 pub trait ContextInternal<'cx>: Sized {
-    fn env(&self) -> Env;
+    fn cx(&self) -> &Cx<'cx>;
+    fn cx_mut(&mut self) -> &mut Cx<'cx>;
+    fn env(&self) -> Env {
+        self.cx().env
+    }
 }
 
 fn default_main(mut cx: ModuleContext) -> NeonResult<()> {
diff --git a/crates/neon/src/context/mod.rs b/crates/neon/src/context/mod.rs
index 036d65589..7bad93f4e 100644
--- a/crates/neon/src/context/mod.rs
+++ b/crates/neon/src/context/mod.rs
@@ -266,8 +266,12 @@ impl<'cx> Cx<'cx> {
 }
 
 impl<'cx> ContextInternal<'cx> for Cx<'cx> {
-    fn env(&self) -> Env {
-        self.env
+    fn cx(&self) -> &Cx<'cx> {
+        self
+    }
+
+    fn cx_mut(&mut self) -> &mut Cx<'cx> {
+        self
     }
 }
 
@@ -664,13 +668,13 @@ impl<'cx> Deref for ModuleContext<'cx> {
     type Target = Cx<'cx>;
 
     fn deref(&self) -> &Self::Target {
-        &self.cx
+        self.cx()
     }
 }
 
 impl<'cx> DerefMut for ModuleContext<'cx> {
     fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.cx
+        self.cx_mut()
     }
 }
 
@@ -727,8 +731,12 @@ impl<'cx> ModuleContext<'cx> {
 }
 
 impl<'cx> ContextInternal<'cx> for ModuleContext<'cx> {
-    fn env(&self) -> Env {
-        self.cx.env
+    fn cx(&self) -> &Cx<'cx> {
+        &self.cx
+    }
+
+    fn cx_mut(&mut self) -> &mut Cx<'cx> {
+        &mut self.cx
     }
 }
 
@@ -873,8 +881,12 @@ impl<'cx> FunctionContext<'cx> {
 }
 
 impl<'cx> ContextInternal<'cx> for FunctionContext<'cx> {
-    fn env(&self) -> Env {
-        self.cx.env
+    fn cx(&self) -> &Cx<'cx> {
+        &self.cx
+    }
+
+    fn cx_mut(&mut self) -> &mut Cx<'cx> {
+        &mut self.cx
     }
 }
 
diff --git a/crates/neon/src/types_impl/extract/with.rs b/crates/neon/src/types_impl/extract/with.rs
new file mode 100644
index 000000000..f28f6652a
--- /dev/null
+++ b/crates/neon/src/types_impl/extract/with.rs
@@ -0,0 +1,24 @@
+use crate::{
+    context::{Context, Cx},
+    result::JsResult,
+    types::extract::TryIntoJs,
+};
+
+pub struct With<F>(pub F);
+
+impl<'cx, F, O> TryIntoJs<'cx> for With<F>
+where
+    F: FnOnce(&mut Cx) -> O,
+    O: TryIntoJs<'cx>,
+{
+    type Value = O::Value;
+
+    fn try_into_js<C>(self, cx: &mut C) -> JsResult<'cx, Self::Value>
+    where
+        C: Context<'cx>,
+    {
+        (self.0)(cx).try_into_js(cx)
+    }
+}
+
+impl<F> super::private::Sealed for With<F> {}