Use OnceCell to hold the jitted code.

Cargo.lock

@@ -1500,6 +1500,7 @@ dependencies = [
  "num-bigint",
  "num-complex",
  "num-traits",
+ "once_cell",
  "parking_lot",
  "rand",
 ]

common/Cargo.toml

@@ -15,6 +15,7 @@ num-bigint = "0.3"
 lexical-core = "0.7"
 hexf-parse = "0.1.0"
 cfg-if = "0.1"
+once_cell = "1.4.1"
 
 [dev-dependencies]
 rand = "0.7.3"

common/src/cell.rs

@@ -1,4 +1,8 @@
 #[cfg(feature = "threading")]
+pub use once_cell::sync::{Lazy, OnceCell};
+#[cfg(not(feature = "threading"))]
+pub use once_cell::unsync::{Lazy, OnceCell};
+#[cfg(feature = "threading")]
 use parking_lot::{
     MappedRwLockReadGuard, MappedRwLockWriteGuard, Mutex, MutexGuard, RwLock, RwLockReadGuard,
     RwLockWriteGuard,

jit/src/lib.rs

@@ -130,6 +130,7 @@ impl CompiledCode {
 }
 
 unsafe impl Send for CompiledCode {}
+unsafe impl Sync for CompiledCode {}
 
 impl Drop for CompiledCode {
     fn drop(&mut self) {

vm/src/obj/objfunction.rs

@@ -12,15 +12,15 @@ use crate::obj::objgenerator::PyGenerator;
 #[cfg(feature = "jit")]
 use crate::pyobject::IntoPyObject;
 use crate::pyobject::{
-    BorrowValue, IdProtocol, ItemProtocol, PyClassImpl, PyContext, PyObjectRef,
-    PyRef, PyResult, PyValue, TypeProtocol,
+    BorrowValue, IdProtocol, ItemProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult,
+    PyValue, TypeProtocol,
 };
 use crate::scope::Scope;
 use crate::slots::{SlotCall, SlotDescriptor};
 use crate::VirtualMachine;
 use itertools::Itertools;
 #[cfg(feature = "jit")]
-use rustpython_common::cell::PyMutex;
+use rustpython_common::cell::OnceCell;
 #[cfg(feature = "jit")]
 use rustpython_jit::CompiledCode;
 
@@ -31,7 +31,7 @@ pub type PyFunctionRef = PyRef<PyFunction>;
 pub struct PyFunction {
     code: PyCodeRef,
     #[cfg(feature = "jit")]
-    jitted_code: PyMutex<Option<CompiledCode>>,
+    jitted_code: OnceCell<CompiledCode>,
     scope: Scope,
     defaults: Option<PyTupleRef>,
     kw_only_defaults: Option<PyDictRef>,
@@ -63,7 +63,7 @@ impl PyFunction {
         PyFunction {
             code,
             #[cfg(feature = "jit")]
-            jitted_code: PyMutex::new(None),
+            jitted_code: OnceCell::new(),
             scope,
             defaults,
             kw_only_defaults,
@@ -234,7 +234,7 @@ impl PyFunction {
         vm: &VirtualMachine,
     ) -> PyResult {
         #[cfg(feature = "jit")]
-        if let Some(jitted_code) = &*self.jitted_code.lock() {
+        if let Some(jitted_code) = self.jitted_code.get() {
             return Ok(jitted_code.invoke().into_pyobject(vm));
         }
 
@@ -307,19 +307,12 @@ impl PyFunction {
         // TODO move cfg onto function when https://github.com/RustPython/RustPython/pull/2120 lands
         #[cfg(feature = "jit")]
         {
-            let mut guard = self.jitted_code.lock();
-            match rustpython_jit::compile(&self.code.code) {
-                Ok(code) => {
-                    *guard = Some(code);
-                    Ok(())
-                }
-                Err(err) => Err(vm.new_runtime_error(err.to_string())),
-            }
-        }
-        #[cfg(not(feature = "jit"))]
-        {
-            Ok(())
+            self.jitted_code.get_or_try_init(|| {
+                rustpython_jit::compile(&self.code.code)
+                    .map_err(|err| vm.new_runtime_error(err.to_string()))
+            })?;
         }
+        Ok(())
     }
 }