diff --git a/CinderX/Jit/jit_context.cpp b/CinderX/Jit/jit_context.cpp
index f5d0dc0523f..757d88563c7 100644
--- a/CinderX/Jit/jit_context.cpp
+++ b/CinderX/Jit/jit_context.cpp
@@ -7,275 +7,102 @@
 #include "Jit/log.h"
 #include "Jit/pyjit.h"
 
-static void deopt_func(_PyJITContext* ctx, BorrowedRef<PyFunctionObject> func) {
-  if (ctx->compiled_funcs.erase(func) == 0) {
-    return;
-  }
-
-  // Reset the entry point.
-  func->vectorcall = (vectorcallfunc)PyEntry_LazyInit;
-}
-
-_PyJITContext::~_PyJITContext() {
-  /* De-optimize any remaining compiled functions. */
-  for (auto it = compiled_funcs.begin(); it != compiled_funcs.end();) {
-    PyFunctionObject* func = *it;
-    ++it;
-    deopt_func(this, func);
-  }
-}
-
-void _PyJITContext_ClearCache(_PyJITContext* ctx) {
-  for (auto& entry : ctx->compiled_codes) {
-    ctx->orphaned_compiled_codes.emplace_back(std::move(entry.second));
-  }
-  ctx->compiled_codes.clear();
-}
+#include <unordered_set>
 
-// Record per-function metadata and set the function's entrypoint.
-static _PyJIT_Result finalizeCompiledFunc(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func,
-    const jit::CompiledFunction& compiled) {
-  jit::ThreadedCompileSerialize guard;
-  if (!ctx->compiled_funcs.emplace(func).second) {
-    // Someone else compiled the function between when our caller checked and
-    // called us.
-    return PYJIT_RESULT_OK;
-  }
-
-  func->vectorcall = compiled.vectorcallEntry();
-  jit::Runtime* rt = jit::Runtime::get();
-  if (rt->hasFunctionEntryCache(func)) {
-    void** indirect = rt->findFunctionEntryCache(func);
-    *indirect = compiled.staticEntry();
-  }
-  return PYJIT_RESULT_OK;
-}
+namespace jit {
 
 namespace {
-struct CompilationResult {
-  jit::CompiledFunction* compiled;
-  _PyJIT_Result result;
-};
 
-jit::CompiledFunction* lookupCompiledCode(
-    _PyJITContext* ctx,
-    BorrowedRef<PyCodeObject> code,
-    BorrowedRef<PyDictObject> builtins,
-    BorrowedRef<PyDictObject> globals) {
-  jit::ThreadedCompileSerialize guard;
-  auto it = ctx->compiled_codes.find(CompilationKey{code, builtins, globals});
-  return it == ctx->compiled_codes.end() ? nullptr : it->second.get();
-}
-
-jit::CompiledFunction* lookupCompiledFunction(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  return lookupCompiledCode(
-      ctx, func->func_code, func->func_builtins, func->func_globals);
-}
-
-CompilationResult compilePreloader(
-    _PyJITContext* ctx,
-    const jit::hir::Preloader& preloader) {
-  BorrowedRef<PyCodeObject> code = preloader.code();
-  BorrowedRef<PyDictObject> globals = preloader.globals();
-  BorrowedRef<PyDictObject> builtins = preloader.builtins();
-
-  int required_flags = CO_OPTIMIZED | CO_NEWLOCALS;
-  int prohibited_flags = CO_SUPPRESS_JIT;
-  // Don't care flags: CO_NOFREE, CO_FUTURE_* (the only still-relevant future
-  // is "annotations" which doesn't impact bytecode execution.)
-  if (code == nullptr ||
-      ((code->co_flags & required_flags) != required_flags) ||
-      (code->co_flags & prohibited_flags) != 0) {
-    return {nullptr, PYJIT_RESULT_CANNOT_SPECIALIZE};
-  }
-
-  // We maintain a set of compilations that are active in all threads, as well
-  // as a per-thread recursion limit (since the JIT can invoke itself to try
-  // and statically bind calls).
-  static std::unordered_set<CompilationKey> active_compiles;
-  static thread_local int compile_depth = 0;
-  const int kMaxCompileDepth = 10;
-  if (compile_depth == kMaxCompileDepth) {
-    return {nullptr, PYJIT_RESULT_RETRY};
-  }
+// We maintain a set of compilations that are active in all threads, as well
+// as a per-thread recursion limit (since the JIT can invoke itself to try
+// and statically bind calls).
+constexpr int kMaxCompileDepth = 10;
+std::unordered_set<CompilationKey> active_compiles;
+thread_local int compile_depth = 0;
 
-  CompilationKey key{code, builtins, globals};
-  {
-    // Attempt to atomically transition the code from "not compiled" to "in
-    // progress".
-    jit::ThreadedCompileSerialize guard;
-    if (jit::CompiledFunction* compiled =
-            lookupCompiledCode(ctx, code, builtins, globals)) {
-      return {compiled, PYJIT_RESULT_OK};
-    }
-    if (!active_compiles.insert(key).second) {
-      return {nullptr, PYJIT_RESULT_RETRY};
-    }
-  }
-
-  compile_depth++;
-  std::unique_ptr<jit::CompiledFunction> compiled =
-      ctx->jit_compiler.Compile(preloader);
-  compile_depth--;
-
-  jit::ThreadedCompileSerialize guard;
-  active_compiles.erase(key);
-  if (compiled == nullptr) {
-    return {nullptr, PYJIT_RESULT_UNKNOWN_ERROR};
-  }
-
-  register_pycode_debug_symbol(
-      code, preloader.fullname().c_str(), compiled.get());
-
-  // Store the compiled code.
-  auto pair = ctx->compiled_codes.emplace(key, std::move(compiled));
-  JIT_CHECK(pair.second == true, "CompilationKey already present");
-  return {pair.first->second.get(), PYJIT_RESULT_OK};
-}
+} // namespace
 
-// Compile the given code object.
-//
-// Returns the CompiledFunction* and PYJIT_RESULT_OK if successful, or nullptr
-// and a failure reason if not.
-template <class... Args>
-CompilationResult compileCode(_PyJITContext* ctx, Args&&... args) {
-  JIT_CHECK(
-      !jit::g_threaded_compile_context.compileRunning(),
-      "multi-thread compile must preload first");
-  auto preloader =
-      jit::hir::Preloader::getPreloader(std::forward<Args>(args)...);
-  if (!preloader) {
-    return {nullptr, PYJIT_RESULT_PYTHON_EXCEPTION};
+Context::~Context() {
+  /* De-optimize any remaining compiled functions. */
+  for (auto it = compiled_funcs_.begin(); it != compiled_funcs_.end();) {
+    PyFunctionObject* func = *it;
+    ++it;
+    deoptFunc(func);
   }
-  return compilePreloader(ctx, *preloader);
 }
 
-} // namespace
-
-_PyJIT_Result _PyJITContext_CompileFunction(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  if (_PyJITContext_DidCompile(ctx, func) == 1) {
+_PyJIT_Result Context::compileFunc(BorrowedRef<PyFunctionObject> func) {
+  if (didCompile(func)) {
     return PYJIT_RESULT_OK;
   }
-  CompilationResult result = compileCode(ctx, func);
-  if (result.compiled == nullptr) {
-    return result.result;
-  }
-
-  return finalizeCompiledFunc(ctx, func, *result.compiled);
+  auto preloader = hir::Preloader::getPreloader(func);
+  return preloader != nullptr ? compilePreloader(func, *preloader)
+                              : PYJIT_RESULT_PYTHON_EXCEPTION;
 }
 
-_PyJIT_Result _PyJITContext_CompileCode(
-    _PyJITContext* ctx,
+_PyJIT_Result Context::compileCode(
     BorrowedRef<> module,
     BorrowedRef<PyCodeObject> code,
     BorrowedRef<PyDictObject> builtins,
     BorrowedRef<PyDictObject> globals) {
-  std::string fullname = jit::codeFullname(module, code);
-  return compileCode(ctx, code, builtins, globals, fullname).result;
+  std::string fullname = codeFullname(module, code);
+  auto preloader =
+      hir::Preloader::getPreloader(code, builtins, globals, fullname);
+  return preloader != nullptr ? compilePreloader(*preloader).result
+                              : PYJIT_RESULT_PYTHON_EXCEPTION;
 }
 
-_PyJIT_Result _PyJITContext_CompilePreloader(
-    _PyJITContext* ctx,
+_PyJIT_Result Context::compilePreloader(
     BorrowedRef<PyFunctionObject> func,
-    const jit::hir::Preloader& preloader) {
-  CompilationResult result = compilePreloader(ctx, preloader);
+    const hir::Preloader& preloader) {
+  CompilationResult result = compilePreloader(preloader);
   if (result.compiled == nullptr) {
     return result.result;
   }
   if (func != nullptr) {
-    return finalizeCompiledFunc(ctx, func, *result.compiled);
+    finalizeFunc(func, *result.compiled);
   }
   return PYJIT_RESULT_OK;
 }
 
-_PyJIT_Result _PyJITContext_AttachCompiledCode(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  JIT_DCHECK(
-      _PyJITContext_DidCompile(ctx, func) == 0, "Function is already compiled");
+_PyJIT_Result Context::attachCompiledCode(BorrowedRef<PyFunctionObject> func) {
+  JIT_DCHECK(!didCompile(func), "Function is already compiled");
 
-  if (jit::CompiledFunction* compiled = lookupCompiledFunction(ctx, func)) {
-    return finalizeCompiledFunc(ctx, func, *compiled);
+  if (CompiledFunction* compiled = lookupFunc(func)) {
+    finalizeFunc(func, *compiled);
+    return PYJIT_RESULT_OK;
   }
   return PYJIT_RESULT_CANNOT_SPECIALIZE;
 }
 
-void _PyJITContext_FuncModified(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  deopt_func(ctx, func);
+void Context::funcModified(BorrowedRef<PyFunctionObject> func) {
+  deoptFunc(func);
 }
 
-void _PyJITContext_FuncDestroyed(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  ctx->compiled_funcs.erase(func);
+void Context::funcDestroyed(BorrowedRef<PyFunctionObject> func) {
+  compiled_funcs_.erase(func);
 }
 
-int _PyJITContext_DidCompile(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  jit::ThreadedCompileSerialize guard;
-
-  return ctx->compiled_funcs.count(func) != 0;
+bool Context::didCompile(BorrowedRef<PyFunctionObject> func) {
+  ThreadedCompileSerialize guard;
+  return compiled_funcs_.count(func) != 0;
 }
 
-int _PyJITContext_GetCodeSize(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  jit::CompiledFunction* jitfunc = lookupCompiledFunction(ctx, func);
-  if (jitfunc == nullptr) {
-    return -1;
-  }
-
-  int size = jitfunc->codeSize();
-  return size;
-}
-
-int _PyJITContext_GetStackSize(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  jit::CompiledFunction* jitfunc = lookupCompiledFunction(ctx, func);
-  if (jitfunc == nullptr) {
-    return -1;
-  }
-
-  return jitfunc->stackSize();
-}
-
-int _PyJITContext_GetSpillStackSize(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  jit::CompiledFunction* jitfunc = lookupCompiledFunction(ctx, func);
-  if (jitfunc == nullptr) {
-    return -1;
-  }
-
-  return jitfunc->spillStackSize();
+CompiledFunction* Context::lookupFunc(BorrowedRef<PyFunctionObject> func) {
+  return lookupCode(func->func_code, func->func_builtins, func->func_globals);
 }
 
 // TODO(T142228417): Deprecate this once callsites have been updated to use
 // get_inlined_functions_stats
-int _PyJITContext_GetNumInlinedFunctions(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  jit::CompiledFunction* jitfunc = lookupCompiledFunction(ctx, func);
-  if (jitfunc == nullptr) {
-    return -1;
-  }
-  return jitfunc->inlinedFunctionsStats().num_inlined_functions;
+int Context::numInlinedFunctions(BorrowedRef<PyFunctionObject> func) {
+  CompiledFunction* jitfunc = lookupFunc(func);
+  return jitfunc != nullptr
+      ? jitfunc->inlinedFunctionsStats().num_inlined_functions
+      : -1;
 }
 
-PyObject* _PyJITContext_GetInlinedFunctionsStats(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  jit::CompiledFunction* jitfunc = lookupCompiledFunction(ctx, func);
+Ref<> Context::inlinedFunctionsStats(BorrowedRef<PyFunctionObject> func) {
+  CompiledFunction* jitfunc = lookupFunc(func);
   if (jitfunc == nullptr) {
     return nullptr;
   }
@@ -321,53 +148,133 @@ PyObject* _PyJITContext_GetInlinedFunctionsStats(
   if (PyDict_SetItemString(py_stats, "failure_stats", failure_stats) < 0) {
     return nullptr;
   }
-  return py_stats.release();
+  return py_stats;
 }
 
-const jit::hir::OpcodeCounts* _PyJITContext_GetHIROpcodeCounts(
-    _PyJITContext* ctx,
+const hir::OpcodeCounts* Context::hirOpcodeCounts(
     BorrowedRef<PyFunctionObject> func) {
-  jit::CompiledFunction* jit_func = lookupCompiledFunction(ctx, func);
+  CompiledFunction* jit_func = lookupFunc(func);
+  return jit_func != nullptr ? &jit_func->hirOpcodeCounts() : nullptr;
+}
+
+int Context::printHIR(BorrowedRef<PyFunctionObject> func) {
+  CompiledFunction* jit_func = lookupFunc(func);
   if (jit_func == nullptr) {
-    return nullptr;
+    return -1;
   }
-  return &jit_func->hirOpcodeCounts();
+  jit_func->printHIR();
+  return 0;
 }
 
-PyObject* _PyJITContext_GetCompiledFunctions(_PyJITContext* ctx) {
-  auto funcs = Ref<>::steal(PyList_New(0));
-  if (funcs == nullptr) {
-    return nullptr;
+int Context::disassemble(BorrowedRef<PyFunctionObject> func) {
+  CompiledFunction* jit_func = lookupFunc(func);
+  if (jit_func == nullptr) {
+    return -1;
   }
+  jit_func->disassemble();
+  return 0;
+}
 
-  for (BorrowedRef<PyFunctionObject> func : ctx->compiled_funcs) {
-    if (PyList_Append(funcs, func) < 0) {
-      return nullptr;
-    }
+const UnorderedSet<BorrowedRef<PyFunctionObject>>& Context::compiledFuncs() {
+  return compiled_funcs_;
+}
+
+void Context::setCinderJitModule(Ref<> mod) {
+  cinderjit_module_ = std::move(mod);
+}
+
+void Context::clearCache() {
+  for (auto& entry : compiled_codes_) {
+    orphaned_compiled_codes_.emplace_back(std::move(entry.second));
   }
-  return funcs.release();
+  compiled_codes_.clear();
 }
 
-int _PyJITContext_PrintHIR(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  jit::CompiledFunction* jit_func = lookupCompiledFunction(ctx, func);
-  if (jit_func == nullptr) {
-    return -1;
+Context::CompilationResult Context::compilePreloader(
+    const hir::Preloader& preloader) {
+  BorrowedRef<PyCodeObject> code = preloader.code();
+  BorrowedRef<PyDictObject> builtins = preloader.builtins();
+  BorrowedRef<PyDictObject> globals = preloader.globals();
+
+  int required_flags = CO_OPTIMIZED | CO_NEWLOCALS;
+  int prohibited_flags = CO_SUPPRESS_JIT;
+  // Don't care flags: CO_NOFREE, CO_FUTURE_* (the only still-relevant future
+  // is "annotations" which doesn't impact bytecode execution.)
+  if (code == nullptr ||
+      ((code->co_flags & required_flags) != required_flags) ||
+      (code->co_flags & prohibited_flags) != 0) {
+    return {nullptr, PYJIT_RESULT_CANNOT_SPECIALIZE};
   }
-  jit_func->printHIR();
 
-  return 0;
+  if (compile_depth == kMaxCompileDepth) {
+    return {nullptr, PYJIT_RESULT_RETRY};
+  }
+
+  CompilationKey key{code, builtins, globals};
+  {
+    // Attempt to atomically transition the code from "not compiled" to "in
+    // progress".
+    ThreadedCompileSerialize guard;
+    if (CompiledFunction* compiled = lookupCode(code, builtins, globals)) {
+      return {compiled, PYJIT_RESULT_OK};
+    }
+    if (!active_compiles.insert(key).second) {
+      return {nullptr, PYJIT_RESULT_RETRY};
+    }
+  }
+
+  compile_depth++;
+  std::unique_ptr<CompiledFunction> compiled = jit_compiler_.Compile(preloader);
+  compile_depth--;
+
+  ThreadedCompileSerialize guard;
+  active_compiles.erase(key);
+  if (compiled == nullptr) {
+    return {nullptr, PYJIT_RESULT_UNKNOWN_ERROR};
+  }
+
+  register_pycode_debug_symbol(
+      code, preloader.fullname().c_str(), compiled.get());
+
+  // Store the compiled code.
+  auto pair = compiled_codes_.emplace(key, std::move(compiled));
+  JIT_CHECK(pair.second, "CompilationKey already present");
+  return {pair.first->second.get(), PYJIT_RESULT_OK};
 }
 
-int _PyJITContext_Disassemble(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func) {
-  jit::CompiledFunction* jit_func = lookupCompiledFunction(ctx, func);
-  if (jit_func == nullptr) {
-    return -1;
+CompiledFunction* Context::lookupCode(
+    BorrowedRef<PyCodeObject> code,
+    BorrowedRef<PyDictObject> builtins,
+    BorrowedRef<PyDictObject> globals) {
+  ThreadedCompileSerialize guard;
+  auto it = compiled_codes_.find(CompilationKey{code, builtins, globals});
+  return it == compiled_codes_.end() ? nullptr : it->second.get();
+}
+
+void Context::deoptFunc(BorrowedRef<PyFunctionObject> func) {
+  if (compiled_funcs_.erase(func) != 0) {
+    // Reset the entry point.
+    func->vectorcall = (vectorcallfunc)PyEntry_LazyInit;
   }
-  jit_func->disassemble();
+}
 
-  return 0;
+void Context::finalizeFunc(
+    BorrowedRef<PyFunctionObject> func,
+    const CompiledFunction& compiled) {
+  ThreadedCompileSerialize guard;
+  if (!compiled_funcs_.emplace(func).second) {
+    // Someone else compiled the function between when our caller checked and
+    // called us.
+    return;
+  }
+
+  func->vectorcall = compiled.vectorcallEntry();
+  Runtime* rt = Runtime::get();
+  if (rt->hasFunctionEntryCache(func)) {
+    void** indirect = rt->findFunctionEntryCache(func);
+    *indirect = compiled.staticEntry();
+  }
+  return;
 }
+
+} // namespace jit
diff --git a/CinderX/Jit/jit_context.h b/CinderX/Jit/jit_context.h
index 36ce6aafbee..55118d39d6b 100644
--- a/CinderX/Jit/jit_context.h
+++ b/CinderX/Jit/jit_context.h
@@ -12,13 +12,13 @@
 #include "Jit/ref.h"
 #include "Jit/util.h"
 
-#include <functional>
 #include <memory>
-#include <unordered_map>
 #include <vector>
 
-// Lookup key for _PyJITContext::compiled_codes: a code object and a globals
-// dict it was JIT-compiled with.
+namespace jit {
+
+// Lookup key for compiled functions in Context: a code object and the globals
+// and builtins dicts it was JIT-compiled with.
 struct CompilationKey {
   // These three are borrowed references; the values are kept alive by strong
   // references in the corresponding jit::CodeRuntime.
@@ -29,15 +29,14 @@ struct CompilationKey {
   CompilationKey(PyObject* code, PyObject* builtins, PyObject* globals)
       : code(code), builtins(builtins), globals(globals) {}
 
-  bool operator==(const CompilationKey& other) const {
-    return code == other.code && globals == other.globals &&
-        builtins == other.builtins;
-  }
+  constexpr bool operator==(const CompilationKey& other) const = default;
 };
 
+} // namespace jit
+
 template <>
-struct std::hash<CompilationKey> {
-  std::size_t operator()(const CompilationKey& key) const {
+struct std::hash<jit::CompilationKey> {
+  std::size_t operator()(const jit::CompilationKey& key) const {
     std::hash<PyObject*> hasher;
     return jit::combineHash(
         jit::combineHash(hasher(key.code), hasher(key.globals)),
@@ -45,198 +44,182 @@ struct std::hash<CompilationKey> {
   }
 };
 
+namespace jit {
+
 /*
- * A JIT context encapsulates all the state managed by an instance of the JIT.
+ * A jit::Context encapsulates all the state managed by an instance of the JIT.
  */
-struct _PyJITContext {
-  ~_PyJITContext();
+class Context {
+ public:
+  ~Context();
 
-  /* General purpose jit compiler */
-  jit::Compiler jit_compiler;
+  /*
+   * JIT compile func and patch its entry point.
+   *
+   * On success, positional only calls to func will use the JIT compiled
+   * version.
+   *
+   * Will return PYJIT_RESULT_OK if the function was already compiled.
+   */
+  _PyJIT_Result compileFunc(BorrowedRef<PyFunctionObject> func);
 
   /*
-   * Set of which functions have JIT-compiled entrypoints.
+   * JIT compile code and store the result in the context.
+   *
+   * This does not patch the entry point of any functions; it is primarily
+   * useful to pre-compile the code object for a nested function so it's
+   * available for use after disabling the JIT.
+   *
+   * Will return PYJIT_RESULT_OK if the code was already compiled.
    */
-  jit::UnorderedSet<BorrowedRef<PyFunctionObject>> compiled_funcs;
+  _PyJIT_Result compileCode(
+      BorrowedRef<> module,
+      BorrowedRef<PyCodeObject> code,
+      BorrowedRef<PyDictObject> builtins,
+      BorrowedRef<PyDictObject> globals);
 
   /*
-   * Compiled code objects, keyed by PyCodeObject* and the globals dict they
-   * were compiled with.
+   * JIT compile function/code-object from a Preloader.
+   *
+   * Patches func entrypoint if the Preloader contains a func.
+   *
+   * Will return PYJIT_RESULT_OK if the function/code object was already
+   * compiled.
    */
-  jit::UnorderedMap<CompilationKey, std::unique_ptr<jit::CompiledFunction>>
-      compiled_codes;
+  _PyJIT_Result compilePreloader(
+      BorrowedRef<PyFunctionObject> func,
+      const hir::Preloader& preloader);
 
   /*
-   * Code which is being kept alive in case it was in use when
-   * _PyJITContext_ClearCache was called. Only intended to be used during
-   * multithreaded_compile_test.
+   * Attach already-compiled code to the given function, if it exists.
+   *
+   * Intended for (but not limited to) use with nested functions after the JIT
+   * is disabled.
+   *
+   * Will return PYJIT_RESULT_OK if the given function already had compiled code
+   * attached.
    */
-  std::vector<std::unique_ptr<jit::CompiledFunction>> orphaned_compiled_codes;
+  _PyJIT_Result attachCompiledCode(BorrowedRef<PyFunctionObject> func);
 
-  Ref<> cinderjit_module;
-};
+  /*
+   * Callbacks invoked by the runtime when a PyFunctionObject is modified or
+   * destroyed.
+   */
+  void funcModified(BorrowedRef<PyFunctionObject> func);
+  void funcDestroyed(BorrowedRef<PyFunctionObject> func);
 
-/*
- * Clear cache of compiled code such that subsequent compilations are always
- * full rather than just re-binding pre-compiled code. Only intended to be used
- * during multithreaded_compile_test.
- */
-void _PyJITContext_ClearCache(_PyJITContext* ctx);
+  /*
+   * Return whether or not this context compiled the supplied function.
+   */
+  bool didCompile(BorrowedRef<PyFunctionObject> func);
 
-/*
- * JIT compile func and patch its entry point.
- *
- * On success, positional only calls to func will use the JIT compiled version.
- *
- * Returns PYJIT_RESULT_OK on success, or if the function was already compiled.
- */
-_PyJIT_Result _PyJITContext_CompileFunction(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  /*
+   * Look up the compiled function object for a given Python function object.
+   */
+  CompiledFunction* lookupFunc(BorrowedRef<PyFunctionObject> func);
 
-/*
- * JIT compile code and store the result in ctx.
- *
- * This does not patch the entry point of any functions; it is primarily useful
- * to pre-compile the code object for a nested function so it's available for
- * use after disabling the JIT.
- *
- * Returns 1 on success, 0 on failure.
- */
-_PyJIT_Result _PyJITContext_CompileCode(
-    _PyJITContext* ctx,
-    BorrowedRef<> module,
-    BorrowedRef<PyCodeObject> code,
-    BorrowedRef<PyDictObject> builtins,
-    BorrowedRef<PyDictObject> globals);
+  /*
+   * Returns the number of functions inlined into a specified JIT-compiled
+   * function.
+   *
+   * Returns -1 if an error occurred.
+   */
+  int numInlinedFunctions(BorrowedRef<PyFunctionObject> func);
 
-/*
- * JIT compile function/code-object from Preloader.
- *
- * Patches func entrypoint if Preloader contains a func.
- */
-_PyJIT_Result _PyJITContext_CompilePreloader(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func,
-    const jit::hir::Preloader& preloader);
+  /*
+   * Return a stats object on the functions inlined into a specified
+   * JIT-compiled function.
+   *
+   * Will return nullptr if the supplied function has not been JIT-compiled.
+   */
+  Ref<> inlinedFunctionsStats(BorrowedRef<PyFunctionObject> func);
 
-/*
- * Attach already-compiled code to the given function, if it exists.
- *
- * Intended for (but not limited to) use with nested functions after the JIT is
- * disabled.
- *
- * Returns PYJIT_RESULT_OK on success or if the given function already had
- * compiled code attached.
- */
-_PyJIT_Result _PyJITContext_AttachCompiledCode(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  /*
+   * Return the HIR opcode counts for a JIT-compiled function, or nullptr if the
+   * function has not been JIT-compiled.
+   */
+  const hir::OpcodeCounts* hirOpcodeCounts(BorrowedRef<PyFunctionObject> func);
 
-/*
- * Callbacks invoked by the runtime when a PyFunctionObject is modified or
- * destroyed.
- */
-void _PyJITContext_FuncModified(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
-void _PyJITContext_FuncDestroyed(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  /*
+   * Print the HIR for func to stdout if it was JIT-compiled.
+   * This function is a no-op if func was not JIT-compiled.
+   *
+   * Returns -1 if an error occurred or 0 otherwise.
+   */
+  int printHIR(BorrowedRef<PyFunctionObject> func);
 
-/*
- * Return whether or not this context compiled the supplied function.
- *
- * Return 1 if so, 0 if not, and -1 if an error occurred.
- */
-int _PyJITContext_DidCompile(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  /*
+   * Print the disassembled code for func to stdout if it was JIT-compiled.
+   * This function is a no-op if func was not JIT-compiled.
+   *
+   * Returns -1 if an error occurred or 0 otherwise.
+   */
+  int disassemble(BorrowedRef<PyFunctionObject> func);
 
-/*
- * Returns the code size in bytes for a specified JIT-compiled function.
- *
- * Returns 0 if the function is not JIT-compiled.
- *
- * Returns -1 if an error occurred.
- */
+  /*
+   * Get a range over all function objects that have been compiled.
+   */
+  const UnorderedSet<BorrowedRef<PyFunctionObject>>& compiledFuncs();
 
-int _PyJITContext_GetCodeSize(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  /*
+   * Set and hold a reference to the cinderjit Python module.
+   */
+  void setCinderJitModule(Ref<> mod);
 
-/*
- * Returns the stack size in bytes for a specified JIT-compiled function.
- *
- * Returns -1 if an error occurred.
- */
+  /*
+   * Clear cache of compiled code such that subsequent compilations are always
+   * full rather than just re-binding pre-compiled code. Only intended to be
+   * used during multithreaded_compile_test.
+   */
+  void clearCache();
 
-int _PyJITContext_GetStackSize(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+ private:
+  struct CompilationResult {
+    CompiledFunction* compiled;
+    _PyJIT_Result result;
+  };
 
-/*
- * Returns the stack size used for spills in bytes for a specified JIT-compiled
- * function.
- *
- * Returns -1 if an error occurred.
- */
-int _PyJITContext_GetSpillStackSize(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  CompilationResult compilePreloader(const hir::Preloader& preloader);
 
-/*
- * Returns the number of functions inlined into a specified JIT-compiled
- * function.
- *
- * Returns -1 if an error occurred.
- */
-int _PyJITContext_GetNumInlinedFunctions(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  CompiledFunction* lookupCode(
+      BorrowedRef<PyCodeObject> code,
+      BorrowedRef<PyDictObject> builtins,
+      BorrowedRef<PyDictObject> globals);
 
-/*
- * Returns the number of functions inlined into a specified JIT-compiled
- * function.
- *
- * Returns -1 if an error occurred.
- */
-PyObject* _PyJITContext_GetInlinedFunctionsStats(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  /*
+   * Reset a function's entry point if it was JIT-compiled.
+   */
+  void deoptFunc(BorrowedRef<PyFunctionObject> func);
 
-/*
- * Returns the HIR opcode counts for a JIT-compiled function, or nullptr if the
- * function has not been JIT-compiled.
- */
-const jit::hir::OpcodeCounts* _PyJITContext_GetHIROpcodeCounts(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  /*
+   * Record per-function metadata for a newly compiled function and set the
+   * function's entrypoint.
+   */
+  void finalizeFunc(
+      BorrowedRef<PyFunctionObject> func,
+      const CompiledFunction& compiled);
 
-/*
- * Return a list of functions that are currently JIT-compiled.
- *
- * Returns a new reference.
- *
- */
-PyObject* _PyJITContext_GetCompiledFunctions(_PyJITContext* ctx);
+  /* General purpose jit compiler */
+  Compiler jit_compiler_;
 
-/*
- * Print the HIR for func to stdout if it was JIT-compiled.
- * This function is a no-op if func was not JIT-compiled.
- *
- * Returns -1 if an error occurred or 0 otherwise.
- */
-int _PyJITContext_PrintHIR(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  /* Set of which functions have JIT-compiled entrypoints. */
+  UnorderedSet<BorrowedRef<PyFunctionObject>> compiled_funcs_;
 
-/*
- * Print the disassembled code for func to stdout if it was JIT-compiled.
- * This function is a no-op if func was not JIT-compiled.
- *
- * Returns -1 if an error occurred or 0 otherwise.
- */
-int _PyJITContext_Disassemble(
-    _PyJITContext* ctx,
-    BorrowedRef<PyFunctionObject> func);
+  /*
+   * Map of all compiled code objects, keyed by their address and also their
+   * builtins and globals objects.
+   */
+  UnorderedMap<CompilationKey, std::unique_ptr<CompiledFunction>>
+      compiled_codes_;
+
+  /*
+   * Code which is being kept alive in case it was in use when
+   * clearCache was called. Only intended to be used during
+   * multithreaded_compile_test.
+   */
+  std::vector<std::unique_ptr<CompiledFunction>> orphaned_compiled_codes_;
+
+  Ref<> cinderjit_module_;
+};
+
+} // namespace jit
diff --git a/CinderX/Jit/pyjit.cpp b/CinderX/Jit/pyjit.cpp
index d9765e4599b..9bfc25357c6 100644
--- a/CinderX/Jit/pyjit.cpp
+++ b/CinderX/Jit/pyjit.cpp
@@ -72,7 +72,7 @@ long g_batch_compilation_time_ms = 0;
 
 } // namespace
 
-static _PyJITContext* jit_ctx;
+static Context* jit_ctx{nullptr};
 static JITList* g_jit_list{nullptr};
 
 // Function and code objects ("units") registered for compilation.
@@ -689,13 +689,12 @@ static _PyJIT_Result compileUnit(BorrowedRef<> unit) {
   if (PyFunction_Check(unit)) {
     BorrowedRef<PyFunctionObject> func(unit);
     CompilationTimer t{func};
-    return _PyJITContext_CompileFunction(jit_ctx, func);
+    return jit_ctx->compileFunc(func);
   }
   JIT_CHECK(PyCode_Check(unit), "Expected function or code object");
   BorrowedRef<PyCodeObject> code(unit);
   const CodeData& data = map_get(jit_code_data, code);
-  return _PyJITContext_CompileCode(
-      jit_ctx, data.module, code, data.builtins, data.globals);
+  return jit_ctx->compileCode(data.module, code, data.builtins, data.globals);
 }
 
 // Compile the given function or code object with a preloader from the global
@@ -705,7 +704,7 @@ static _PyJIT_Result compilePreloaded(BorrowedRef<> unit) {
   if (PyFunction_Check(unit)) {
     func = BorrowedRef<PyFunctionObject>{unit};
   }
-  return _PyJITContext_CompilePreloader(jit_ctx, func, getPreloader(unit));
+  return jit_ctx->compilePreloader(func, getPreloader(unit));
 }
 
 static void compile_worker_thread() {
@@ -839,7 +838,7 @@ static PyObject* multithreaded_compile_test(PyObject*, PyObject*) {
   g_compile_workers_attempted = 0;
   g_compile_workers_retries = 0;
   JIT_LOG("(Re)compiling {} units", jit_reg_units.size());
-  _PyJITContext_ClearCache(jit_ctx);
+  jit_ctx->clearCache();
   std::chrono::time_point time_start = std::chrono::steady_clock::now();
   multithread_compile_all();
   std::chrono::time_point time_end = std::chrono::steady_clock::now();
@@ -945,7 +944,7 @@ int _PyJIT_IsCompiled(PyObject* func) {
       "Expected PyFunctionObject, got '{:.200}'",
       Py_TYPE(func)->tp_name);
 
-  return _PyJITContext_DidCompile(jit_ctx, func);
+  return int{jit_ctx->didCompile(func)};
 }
 
 static PyObject* is_jit_compiled(PyObject* /* self */, PyObject* func) {
@@ -966,15 +965,12 @@ static PyObject* print_hir(PyObject* /* self */, PyObject* func) {
     return nullptr;
   }
 
-  int st = _PyJITContext_DidCompile(jit_ctx, func);
-  if (st == -1) {
-    return nullptr;
-  } else if (st == 0) {
+  if (!jit_ctx->didCompile(func)) {
     PyErr_SetString(PyExc_ValueError, "function is not jit compiled");
     return nullptr;
   }
 
-  if (_PyJITContext_PrintHIR(jit_ctx, func) < 0) {
+  if (jit_ctx->printHIR(func) < 0) {
     return nullptr;
   } else {
     Py_RETURN_NONE;
@@ -987,15 +983,12 @@ static PyObject* disassemble(PyObject* /* self */, PyObject* func) {
     return nullptr;
   }
 
-  int st = _PyJITContext_DidCompile(jit_ctx, func);
-  if (st == -1) {
-    return nullptr;
-  } else if (st == 0) {
+  if (!jit_ctx->didCompile(func)) {
     PyErr_SetString(PyExc_ValueError, "function is not jit compiled");
     return nullptr;
   }
 
-  if (_PyJITContext_Disassemble(jit_ctx, func) < 0) {
+  if (jit_ctx->disassemble(func) < 0) {
     return nullptr;
   } else {
     Py_RETURN_NONE;
@@ -1024,7 +1017,16 @@ static PyObject* jit_list_append(PyObject* /* self */, PyObject* line) {
 }
 
 static PyObject* get_compiled_functions(PyObject* /* self */, PyObject*) {
-  return _PyJITContext_GetCompiledFunctions(jit_ctx);
+  auto funcs = Ref<>::steal(PyList_New(0));
+  if (funcs == nullptr) {
+    return nullptr;
+  }
+  for (BorrowedRef<PyFunctionObject> func : jit_ctx->compiledFuncs()) {
+    if (PyList_Append(funcs, func) < 0) {
+      return nullptr;
+    }
+  }
+  return funcs.release();
 }
 
 static PyObject* get_compilation_time(PyObject* /* self */, PyObject*) {
@@ -1050,23 +1052,19 @@ static PyObject* get_inlined_functions_stats(PyObject*, PyObject* func) {
   if (jit_ctx == nullptr) {
     Py_RETURN_NONE;
   }
-  return _PyJITContext_GetInlinedFunctionsStats(jit_ctx, func);
+  return jit_ctx->inlinedFunctionsStats(func).release();
 }
 
 static PyObject* get_num_inlined_functions(PyObject*, PyObject* func) {
-  if (jit_ctx == nullptr) {
-    return PyLong_FromLong(0);
-  }
-  int size = _PyJITContext_GetNumInlinedFunctions(jit_ctx, func);
+  int size = jit_ctx != nullptr ? jit_ctx->numInlinedFunctions(func) : 0;
   return PyLong_FromLong(size);
 }
 
 static PyObject* get_function_hir_opcode_counts(PyObject*, PyObject* func) {
-  if (jit_ctx == NULL) {
+  if (jit_ctx == nullptr) {
     Py_RETURN_NONE;
   }
-  const hir::OpcodeCounts* counts =
-      _PyJITContext_GetHIROpcodeCounts(jit_ctx, func);
+  const hir::OpcodeCounts* counts = jit_ctx->hirOpcodeCounts(func);
   if (counts == nullptr) {
     Py_RETURN_NONE;
   }
@@ -1217,20 +1215,18 @@ static PyObject* get_compiled_size(PyObject* /* self */, PyObject* func) {
   if (jit_ctx == nullptr) {
     return PyLong_FromLong(0);
   }
-
-  long size = _PyJITContext_GetCodeSize(jit_ctx, func);
-  PyObject* res = PyLong_FromLong(size);
-  return res;
+  CompiledFunction* compiled_func = jit_ctx->lookupFunc(func);
+  int size = compiled_func != nullptr ? compiled_func->codeSize() : -1;
+  return PyLong_FromLong(size);
 }
 
 static PyObject* get_compiled_stack_size(PyObject* /* self */, PyObject* func) {
   if (jit_ctx == nullptr) {
     return PyLong_FromLong(0);
   }
-
-  long size = _PyJITContext_GetStackSize(jit_ctx, func);
-  PyObject* res = PyLong_FromLong(size);
-  return res;
+  CompiledFunction* compiled_func = jit_ctx->lookupFunc(func);
+  int size = compiled_func != nullptr ? compiled_func->stackSize() : -1;
+  return PyLong_FromLong(size);
 }
 
 static PyObject* get_compiled_spill_stack_size(
@@ -1239,10 +1235,9 @@ static PyObject* get_compiled_spill_stack_size(
   if (jit_ctx == nullptr) {
     return PyLong_FromLong(0);
   }
-
-  long size = _PyJITContext_GetSpillStackSize(jit_ctx, func);
-  PyObject* res = PyLong_FromLong(size);
-  return res;
+  CompiledFunction* compiled_func = jit_ctx->lookupFunc(func);
+  int size = compiled_func != nullptr ? compiled_func->spillStackSize() : -1;
+  return PyLong_FromLong(size);
 }
 
 static PyObject* jit_frame_mode(PyObject* /* self */, PyObject*) {
@@ -1855,14 +1850,14 @@ int _PyJIT_Initialize() {
 
   CodeAllocator::makeGlobalCodeAllocator();
 
-  jit_ctx = new _PyJITContext();
+  jit_ctx = new Context();
 
   PyObject* mod = PyModule_Create(&jit_module);
   if (mod == nullptr) {
     return -1;
   }
 
-  jit_ctx->cinderjit_module = Ref<PyObject>::steal(mod);
+  jit_ctx->setCinderJitModule(Ref<PyObject>::steal(mod));
 
   PyObject* modname = PyUnicode_InternFromString("cinderjit");
   if (modname == nullptr) {
@@ -1935,7 +1930,7 @@ _PyJIT_Result _PyJIT_CompileFunction(PyFunctionObject* raw_func) {
     if (preloader == nullptr) {
       return PYJIT_RESULT_CANNOT_SPECIALIZE;
     }
-    return _PyJITContext_CompilePreloader(jit_ctx, func, *preloader);
+    return jit_ctx->compilePreloader(func, *preloader);
   }
 
   if (!shouldCompile(func)) {
@@ -1944,7 +1939,7 @@ _PyJIT_Result _PyJIT_CompileFunction(PyFunctionObject* raw_func) {
 
   CompilationTimer timer(func);
   jit_reg_units.erase(func);
-  return _PyJITContext_CompileFunction(jit_ctx, func);
+  return jit_ctx->compileFunc(func);
 }
 
 // Recursively search the given co_consts tuple for any code objects that are
@@ -1981,7 +1976,7 @@ int _PyJIT_RegisterFunction(PyFunctionObject* func) {
   // Attempt to attach already-compiled code even if the JIT is disabled, as
   // long as it hasn't been finalized.
   if (jit_ctx != nullptr &&
-      _PyJITContext_AttachCompiledCode(jit_ctx, func) == PYJIT_RESULT_OK) {
+      jit_ctx->attachCompiledCode(func) == PYJIT_RESULT_OK) {
     return 1;
   }
 
@@ -2055,7 +2050,7 @@ void _PyJIT_TypeDestroyed(PyTypeObject* type) {
 
 void _PyJIT_FuncModified(PyFunctionObject* func) {
   if (jit_ctx) {
-    _PyJITContext_FuncModified(jit_ctx, func);
+    jit_ctx->funcModified(func);
   }
 }
 
@@ -2071,7 +2066,7 @@ void _PyJIT_FuncDestroyed(PyFunctionObject* func) {
     perf_trampoline_reg_units.erase(reinterpret_cast<PyObject*>(func));
   }
   if (jit_ctx) {
-    _PyJITContext_FuncDestroyed(jit_ctx, func);
+    jit_ctx->funcDestroyed(func);
   }
 }
 
@@ -2105,7 +2100,7 @@ static void dump_jit_compiled_functions(const std::string& filename) {
     JIT_LOG("Failed to open {} when dumping jit compiled functions", filename);
     return;
   }
-  for (BorrowedRef<PyFunctionObject> func : jit_ctx->compiled_funcs) {
+  for (BorrowedRef<PyFunctionObject> func : jit_ctx->compiledFuncs()) {
     file << funcFullname(func) << std::endl;
   }
 }
@@ -2137,7 +2132,7 @@ int _PyJIT_Finalize() {
   }
 
   // Always release references from Runtime objects: C++ clients may have
-  // invoked the JIT directly without initializing a full _PyJITContext.
+  // invoked the JIT directly without initializing a full jit::Context.
   jit::Runtime::get()->clearDeoptStats();
   jit::Runtime::get()->releaseReferences();
 
diff --git a/CinderX/RuntimeTests/jit_context_test.cpp b/CinderX/RuntimeTests/jit_context_test.cpp
index 6b9ce8b6238..81ea1467736 100644
--- a/CinderX/RuntimeTests/jit_context_test.cpp
+++ b/CinderX/RuntimeTests/jit_context_test.cpp
@@ -9,24 +9,23 @@
 
 #include <memory>
 
-class PyJITContextTest : public RuntimeTest {
+class JITContextTest : public RuntimeTest {
  public:
   void SetUp() override {
     RuntimeTest::SetUp();
-    jit_ctx_ = new _PyJITContext();
+    jit_ctx_ = std::make_unique<jit::Context>();
     ASSERT_NE(jit_ctx_, nullptr) << "Failed creating jit context";
   }
 
   void TearDown() override {
-    delete jit_ctx_;
-    jit_ctx_ = nullptr;
+    jit_ctx_.reset();
     RuntimeTest::TearDown();
   }
 
-  _PyJITContext* jit_ctx_;
+  std::unique_ptr<jit::Context> jit_ctx_;
 };
 
-TEST_F(PyJITContextTest, UnwatchableBuiltins) {
+TEST_F(JITContextTest, UnwatchableBuiltins) {
   // This is a C++ test rather than in test_cinderjit so we can guarantee a
   // fresh runtime state with a watchable builtins dict when the test begins.
   const char* py_src = R"(
@@ -45,7 +44,7 @@ foo = "hello"
 )";
 
   Ref<PyFunctionObject> func(compileAndGet(py_src, "func"));
-  ASSERT_EQ(_PyJITContext_CompileFunction(jit_ctx_, func), PYJIT_RESULT_OK);
+  ASSERT_EQ(jit_ctx_->compileFunc(func), PYJIT_RESULT_OK);
 
   auto empty_tuple = Ref<>::steal(PyTuple_New(0));
   auto result = Ref<>::steal(PyObject_Call(func, empty_tuple, nullptr));