Consolidate feature set management into a with_feature_set method

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Changed
 
 - Upgrade Solana crates to 2.2.0 ([#138](https://github.com/LiteSVM/litesvm/pull/138)).
+- Consolidate feature set management into a `with_feature_set` method and remove the `feature_set` param from `with_builtins` and `with_precompiles` ([#141](https://github.com/LiteSVM/litesvm/pull/141)).
 
 ## [0.5.0] - 2025-01-23
 

crates/litesvm/src/lib.rs

@@ -373,10 +373,11 @@ impl LiteSVM {
     /// Creates the basic test environment.
     pub fn new() -> Self {
         LiteSVM::default()
-            .with_builtins(None)
+            .with_feature_set(FeatureSet::all_enabled())
+            .with_builtins()
             .with_lamports(1_000_000u64.wrapping_mul(LAMPORTS_PER_SOL))
             .with_sysvars()
-            .with_precompiles(None)
+            .with_precompiles()
             .with_spl_programs()
             .with_sigverify(true)
             .with_blockhash_check(true)
@@ -446,28 +447,38 @@ impl LiteSVM {
         self
     }
 
+    /// Set the FeatureSet used by the VM instance.
+    pub fn with_feature_set(mut self, feature_set: FeatureSet) -> Self {
+        self.set_feature_set(feature_set);
+        self
+    }
+
     #[cfg_attr(feature = "nodejs-internal", qualifiers(pub))]
-    fn set_builtins(&mut self, feature_set: Option<FeatureSet>) {
-        let mut feature_set = feature_set.unwrap_or(FeatureSet::all_enabled());
+    fn set_feature_set(&mut self, feature_set: FeatureSet) {
+        self.feature_set = feature_set;
+    }
 
+    #[cfg_attr(feature = "nodejs-internal", qualifiers(pub))]
+    fn set_builtins(&mut self) {
         BUILTINS.iter().for_each(|builtint| {
-            let loaded_program =
-                ProgramCacheEntry::new_builtin(0, builtint.name.len(), builtint.entrypoint);
-            self.accounts
-                .programs_cache
-                .replenish(builtint.program_id, Arc::new(loaded_program));
-            self.accounts.add_builtin_account(
-                builtint.program_id,
-                crate::utils::create_loadable_account_for_test(builtint.name),
-            );
-
-            if let Some(feature_id) = builtint.feature_id {
-                feature_set.activate(&feature_id, 0);
+            if builtint
+                .feature_id
+                .map_or(true, |x| self.feature_set.is_active(&x))
+            {
+                let loaded_program =
+                    ProgramCacheEntry::new_builtin(0, builtint.name.len(), builtint.entrypoint);
+                self.accounts
+                    .programs_cache
+                    .replenish(builtint.program_id, Arc::new(loaded_program));
+                self.accounts.add_builtin_account(
+                    builtint.program_id,
+                    crate::utils::create_loadable_account_for_test(builtint.name),
+                );
             }
         });
 
         let program_runtime_v1 = create_program_runtime_environment_v1(
-            &feature_set,
+            &self.feature_set,
             &ComputeBudget::default(),
             false,
             true,
@@ -479,12 +490,13 @@ impl LiteSVM {
 
         self.accounts.programs_cache.environments.program_runtime_v1 = Arc::new(program_runtime_v1);
         self.accounts.programs_cache.environments.program_runtime_v2 = Arc::new(program_runtime_v2);
-        self.feature_set = feature_set;
     }
 
     /// Changes the default builtins.
-    pub fn with_builtins(mut self, feature_set: Option<FeatureSet>) -> Self {
-        self.set_builtins(feature_set);
+    // 
+    // Use `with_feature_set` beforehand to change change what builtins are added.
+    pub fn with_builtins(mut self) -> Self {
+        self.set_builtins();
         self
     }
 
@@ -536,13 +548,15 @@ impl LiteSVM {
     }
 
     #[cfg_attr(feature = "nodejs-internal", qualifiers(pub))]
-    fn set_precompiles(&mut self, feature_set: Option<FeatureSet>) {
-        let feature_set = feature_set.unwrap_or_else(FeatureSet::all_enabled);
-        load_precompiles(self, feature_set);
+    fn set_precompiles(&mut self) {
+        load_precompiles(self);
     }
 
-    pub fn with_precompiles(mut self, feature_set: Option<FeatureSet>) -> Self {
-        self.set_precompiles(feature_set);
+    /// Adds the standard precompiles to the VM.
+    // 
+    // Use `with_feature_set` beforehand to change change what precompiles are added.
+    pub fn with_precompiles(mut self) -> Self {
+        self.set_precompiles();
         self
     }
 

crates/litesvm/src/precompiles.rs

@@ -1,13 +1,12 @@
 use {
     solana_account::{AccountSharedData, WritableAccount},
-    solana_feature_set::FeatureSet,
     solana_precompiles::get_precompiles,
     solana_sdk_ids::native_loader,
 };
 
 use crate::LiteSVM;
 
-pub(crate) fn load_precompiles(svm: &mut LiteSVM, feature_set: FeatureSet) {
+pub(crate) fn load_precompiles(svm: &mut LiteSVM) {
     let mut account = AccountSharedData::default();
     account.set_owner(native_loader::id());
     account.set_lamports(1);
@@ -16,7 +15,7 @@ pub(crate) fn load_precompiles(svm: &mut LiteSVM, feature_set: FeatureSet) {
     for precompile in get_precompiles() {
         if precompile
             .feature
-            .map_or(true, |feature_id| feature_set.is_active(&feature_id))
+            .map_or(true, |feature_id| svm.feature_set.is_active(&feature_id))
         {
             svm.set_account(precompile.program_id, account.clone().into())
                 .unwrap();

crates/node-litesvm/litesvm/index.ts

@@ -145,13 +145,21 @@ export class LiteSVM {
 	}
 
 	/**
-	 * Adds the standard builtin programs.
-	 * @param featureSet if provided, decides what builtins to add based on what
-	 *  features are active
+	 * Set the FeatureSet used by the VM instance.
+	 * @param featureSet The FeatureSet to use.
 	 * @returns The modified LiteSVM instance
 	 */
-	withBuiltins(featureSet?: FeatureSet): LiteSVM {
-		this.inner.setBuiltins(featureSet);
+	withFeatureSet(featureSet: FeatureSet): LiteSVM {
+		this.inner.setFeatureSet(featureSet);
+		return this;
+	}
+
+	/**
+	 * Adds the standard builtin programs. Use `withFeatureSet` beforehand to change change what builtins are added.
+	 * @returns The modified LiteSVM instance
+	 */
+	withBuiltins(): LiteSVM {
+		this.inner.setBuiltins();
 		return this;
 	}
 
@@ -196,13 +204,11 @@ export class LiteSVM {
 	}
 
 	/**
-	 * Adds the standard precompiles.
-	 * @param featureSet if provided, decides what precompiles to add based on what
-	 *  features are active
+	 * Adds the standard precompiles. Use `withFeatureSet` beforehand to change change what builtins are added.
 	 * @returns The modified LiteSVM instance
 	 */
-	withPrecompiles(featureSet?: FeatureSet): LiteSVM {
-		this.inner.setPrecompiles(featureSet);
+	withPrecompiles(): LiteSVM {
+		this.inner.setPrecompiles();
 		return this;
 	}
 

crates/node-litesvm/litesvm/internal.d.ts

@@ -525,7 +525,9 @@ export declare class LiteSvm {
   /** Includes the default sysvars */
   setSysvars(): void
   /** Changes the default builtins */
-  setBuiltins(featureSet?: FeatureSet | undefined | null): void
+  setFeatureSet(featureSet: FeatureSet): void
+  /** Changes the default builtins */
+  setBuiltins(): void
   /** Changes the initial lamports in LiteSVM's airdrop account */
   setLamports(lamports: bigint): void
   /** Includes the standard SPL programs */
@@ -536,7 +538,7 @@ export declare class LiteSvm {
    */
   setTransactionHistory(capacity: bigint): void
   setLogBytesLimit(limit?: bigint | undefined | null): void
-  setPrecompiles(featureSet?: FeatureSet | undefined | null): void
+  setPrecompiles(): void
   /** Returns minimum balance required to make an account with specified data length rent exempt. */
   minimumBalanceForRentExemption(dataLen: bigint): bigint
   /** Returns all information associated with the account of the provided pubkey. */

crates/node-litesvm/src/lib.rs

@@ -125,8 +125,14 @@ impl LiteSvm {
 
     #[napi]
     /// Changes the default builtins
-    pub fn set_builtins(&mut self, feature_set: Option<&FeatureSet>) {
-        self.0.set_builtins(feature_set.map(|x| x.0.clone()));
+    pub fn set_feature_set(&mut self, feature_set: &FeatureSet) {
+        self.0.set_feature_set(feature_set.0.clone());
+    }
+
+    #[napi]
+    /// Changes the default builtins
+    pub fn set_builtins(&mut self) {
+        self.0.set_builtins();
     }
 
     #[napi]
@@ -160,8 +166,8 @@ impl LiteSvm {
     }
 
     #[napi]
-    pub fn set_precompiles(&mut self, feature_set: Option<&FeatureSet>) {
-        self.0.set_precompiles(feature_set.map(|x| x.0.clone()));
+    pub fn set_precompiles(&mut self) {
+        self.0.set_precompiles();
     }
 
     #[napi]