chore: comment

evaluator/nixjit_jit/src/lib.rs

@@ -1,3 +1,14 @@
+//! The Just-In-Time (JIT) compilation module for nixjit.
+//!
+//! This module provides functionality to compile Low-Level IR (LIR) expressions
+//! into optimized machine code using Cranelift. The JIT compiler translates
+//! Nix expressions into efficient native code for faster evaluation.
+//!
+//! The main components are:
+//! - `JITCompiler`: The core compiler that manages the compilation process
+//! - `JITContext`: A trait that provides the execution context for JIT-compiled code
+//! - `Context`: An internal compilation context used during code generation
+
 use std::marker::PhantomData;
 use std::ops::Deref;
 use std::rc::Rc;
@@ -18,15 +29,33 @@ mod helpers;
 pub use compile::JITCompile;
 use helpers::*;
 
-pub trait JITContext: EvalContext + Sized {
+/// A trait that provides the execution context for JIT-compiled code.
+///
+/// This trait extends `EvalContext` with additional methods needed
+/// for JIT compilation, such as stack and argument lookups, and
+/// managing `with` expression scopes.
+pub trait JITContext: EvalContext {
+    /// Looks up a value in the evaluation stack by offset.
     fn lookup_stack(&self, offset: usize) -> &Value;
+    /// Looks up a function argument by offset.
     fn lookup_arg(&self, offset: usize) -> &Value;
+    /// Enters a `with` expression scope with the given namespace.
     fn enter_with(&mut self, namespace: Rc<HashMap<String, Value>>);
+    /// Exits the current `with` expression scope.
     fn exit_with(&mut self);
 }
 
+/// Type alias for a JIT-compiled function.
+///
+/// This represents a function pointer to JIT-compiled code that takes
+/// a context pointer and a mutable value pointer as arguments.
 type F<Ctx> = unsafe extern "C" fn(*const Ctx, *mut Value);
 
+/// A JIT-compiled function.
+///
+/// This struct holds a function pointer to the compiled code and
+/// a set of strings that were used during compilation, which need
+/// to be kept alive for the function to work correctly.
 pub struct JITFunc<Ctx: JITContext> {
     func: F<Ctx>,
     strings: HashSet<String>,
@@ -39,10 +68,18 @@ impl<Ctx: JITContext> Deref for JITFunc<Ctx> {
     }
 }
 
+/// The internal compilation context used during code generation.
+///
+/// This context holds references to the compiler, the Cranelift function builder,
+/// and manages resources like stack slots and string literals during compilation.
 struct Context<'comp, 'ctx, Ctx: JITContext> {
+    /// Reference to the JIT compiler.
     pub compiler: &'comp mut JITCompiler<Ctx>,
+    /// The Cranelift function builder used to generate IR.
     pub builder: FunctionBuilder<'ctx>,
+    /// Stack slots available for reuse.
     free_slots: Vec<StackSlot>,
+    /// String literals used during compilation.
     strings: HashSet<String>,
 }
 
@@ -374,6 +411,7 @@ impl<'comp, 'ctx, Ctx: JITContext> Context<'comp, 'ctx, Ctx> {
     }
 }
 
+/// The main JIT compiler that manages the compilation process.
 pub struct JITCompiler<Ctx: JITContext> {
     ctx: codegen::Context,
     module: JITModule,