optimization: with scope lookup

Cargo.lock

@@ -1013,12 +1013,6 @@ dependencies = [
  "windows-sys 0.60.2",
 ]
 
-[[package]]
-name = "fixedbitset"
-version = "0.5.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99"
-
 [[package]]
 name = "flate2"
 version = "1.1.5"
@@ -1957,7 +1951,6 @@ dependencies = [
  "nix-js-macros",
  "nix-nar",
  "num_enum",
- "petgraph",
  "regex",
  "reqwest",
  "rnix",
@@ -1984,7 +1977,6 @@ name = "nix-js-macros"
 version = "0.1.0"
 dependencies = [
  "convert_case 0.8.0",
- "proc-macro2",
  "quote",
  "syn",
 ]
@@ -2170,18 +2162,6 @@ version = "2.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
 
-[[package]]
-name = "petgraph"
-version = "0.8.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455"
-dependencies = [
- "fixedbitset",
- "hashbrown 0.15.5",
- "indexmap",
- "serde",
-]
-
 [[package]]
 name = "pin-project"
 version = "1.1.10"

flake.lock

@@ -8,11 +8,11 @@
         "rust-analyzer-src": "rust-analyzer-src"
       },
       "locked": {
-        "lastModified": 1770447430,
+        "lastModified": 1770966612,
-        "narHash": "sha256-smrRbWhvJF6BATB6pXbD8Cp04HRrVcYQkXqOhUF81nk=",
+        "narHash": "sha256-S6k14z/JsDwX6zZyLucDBTOe/9RsvxH9GTUxHn2o4vc=",
         "owner": "nix-community",
         "repo": "fenix",
-        "rev": "e1b28f6ca0d1722edceec1f2f3501558988d1aed",
+        "rev": "e90d48dcfaebac7ea7a5687888a2d0733be26343",
         "type": "github"
       },
       "original": {
@@ -37,11 +37,11 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1770197578,
+        "lastModified": 1770841267,
-        "narHash": "sha256-AYqlWrX09+HvGs8zM6ebZ1pwUqjkfpnv8mewYwAo+iM=",
+        "narHash": "sha256-9xejG0KoqsoKEGp2kVbXRlEYtFFcDTHjidiuX8hGO44=",
         "owner": "nixos",
         "repo": "nixpkgs",
-        "rev": "00c21e4c93d963c50d4c0c89bfa84ed6e0694df2",
+        "rev": "ec7c70d12ce2fc37cb92aff673dcdca89d187bae",
         "type": "github"
       },
       "original": {
@@ -61,11 +61,11 @@
     "rust-analyzer-src": {
       "flake": false,
       "locked": {
-        "lastModified": 1770290336,
+        "lastModified": 1770934477,
-        "narHash": "sha256-rJ79U68ZLjCSg1Qq+63aBXi//W7blaKiYq9NnfeTboA=",
+        "narHash": "sha256-GX0cINHhhzUbQHyDYN2Mc+ovb6Sx/4yrF95VVou9aW4=",
         "owner": "rust-lang",
         "repo": "rust-analyzer",
-        "rev": "d2a00da09293267e5be2efb216698762929d7140",
+        "rev": "931cd553be123b11db1435ac7ea5657e62e5e601",
         "type": "github"
       },
       "original": {

flake.nix

@@ -27,6 +27,7 @@
                 "rust-analyzer"
               ])
               cargo-outdated
+              cargo-machete
               lldb
               valgrind
               hyperfine

nix-js-macros/Cargo.toml

@@ -9,5 +9,4 @@ proc-macro = true
 [dependencies]
 convert_case = "0.8"
 quote = "1.0"
-proc-macro2 = "1.0"
 syn = { version = "2.0", features = ["full"] }

nix-js/Cargo.toml

@@ -23,7 +23,6 @@ thiserror = "2"
 miette = { version = "7.4", features = ["fancy"] }
 
 hashbrown = "0.16"
-petgraph = "0.8"
 string-interner = "0.19"
 
 rust-embed="8.11"

nix-js/runtime-ts/src/helpers.ts

@@ -380,3 +380,20 @@ export const ifFunc = (cond: NixValue, consq: NixValue, alter: NixValue) => {
 export const mkPos = (span: string): NixAttrs => {
   return Deno.core.ops.op_decode_span(span);
 };
+
+interface WithScope {
+  env: NixValue;
+  last: WithScope | null;
+}
+
+export const lookupWith = (name: string, withScope: WithScope | null): NixValue => {
+  let current = withScope;
+  while (current !== null) {
+    const attrs = forceAttrs(current.env);
+    if (name in attrs) {
+      return attrs[name];
+    }
+    current = current.last;
+  }
+  throw new Error(`undefined variable '${name}'`);
+};

nix-js/runtime-ts/src/index.ts

@@ -27,6 +27,7 @@ import {
   popContext,
   withContext,
   mkPos,
+  lookupWith,
 } from "./helpers";
 import { op } from "./operators";
 import { builtins, PRIMOP_METADATA } from "./builtins";
@@ -60,6 +61,7 @@ export const Nix = {
   hasAttr,
   select,
   selectWithDefault,
+  lookupWith,
   validateParams,
   resolvePath,
   coerceToString,

nix-js/src/codegen.rs

@@ -4,6 +4,71 @@ use std::path::Path;
 use crate::ir::*;
 use crate::value::Symbol;
 
+macro_rules! code {
+    ($buf:expr, $ctx:expr; $($item:expr)*) => {{
+        $(
+            ($item).compile($ctx, $buf);
+        )*
+    }};
+
+    ($buf:expr, $ctx:expr; $($item:expr)*) => {{
+        $(
+            ($item).compile($ctx, $buf);
+        )*
+    }};
+
+    ($buf:expr, $fmt:literal, $($arg:tt)*) => {
+        write!($buf, $fmt, $($arg)*).unwrap()
+    };
+
+    ($buf:expr, $fmt:literal) => {
+        write!($buf, $fmt).unwrap()
+    };
+}
+
+pub(crate) fn compile(expr: &Ir, ctx: &impl CodegenContext) -> String {
+    let mut buf = CodeBuffer::with_capacity(8192);
+
+    code!(&mut buf, ctx; "(()=>{");
+
+    if std::env::var("NIX_JS_DEBUG_THUNKS").is_ok() {
+        code!(&mut buf, ctx; "Nix.DEBUG_THUNKS.enabled=true;");
+    }
+
+    code!(&mut buf, ctx;
+        "Nix.builtins.storeDir="
+        quoted(ctx.get_store_dir())
+        ";const __currentDir="
+        quoted(&ctx.get_current_dir().display().to_string())
+        ";const __with=null;return "
+        expr
+        "})()");
+
+    buf.into_string()
+}
+
+pub(crate) fn compile_scoped(expr: &Ir, ctx: &impl CodegenContext) -> String {
+    let mut buf = CodeBuffer::with_capacity(8192);
+
+    code!(&mut buf, ctx; "((__scope)=>{");
+
+    if std::env::var("NIX_JS_DEBUG_THUNKS").is_ok() {
+        code!(&mut buf, ctx; "Nix.DEBUG_THUNKS.enabled=true;");
+    }
+
+    code!(&mut buf, ctx;
+        "Nix.builtins.storeDir="
+        quoted(ctx.get_store_dir())
+        ";const __currentDir="
+        quoted(&ctx.get_current_dir().display().to_string())
+        ";return "
+        expr
+        "})"
+    );
+
+    buf.into_string()
+}
+
 pub(crate) struct CodeBuffer {
     buf: String,
 }
@@ -83,71 +148,6 @@ fn joined<Ctx: CodegenContext, I: Iterator, F: Fn(&Ctx, &mut CodeBuffer, I::Item
     }
 }
 
-macro_rules! code {
-    ($buf:expr, $ctx:expr; $($item:expr)*) => {{
-        $(
-            ($item).compile($ctx, $buf);
-        )*
-    }};
-
-    ($buf:expr, $ctx:expr; $($item:expr)*) => {{
-        $(
-            ($item).compile($ctx, $buf);
-        )*
-    }};
-
-    ($buf:expr, $fmt:literal, $($arg:tt)*) => {
-        write!($buf, $fmt, $($arg)*).unwrap()
-    };
-
-    ($buf:expr, $fmt:literal) => {
-        write!($buf, $fmt).unwrap()
-    };
-}
-
-pub(crate) fn compile(expr: &Ir, ctx: &impl CodegenContext) -> String {
-    let mut buf = CodeBuffer::with_capacity(8192);
-
-    code!(&mut buf, ctx; "(()=>{");
-
-    if std::env::var("NIX_JS_DEBUG_THUNKS").is_ok() {
-        code!(&mut buf, ctx; "Nix.DEBUG_THUNKS.enabled=true;");
-    }
-
-    code!(&mut buf, ctx;
-        "Nix.builtins.storeDir="
-        quoted(ctx.get_store_dir())
-        ";const __currentDir="
-        quoted(&ctx.get_current_dir().display().to_string())
-        ";return "
-        expr
-        "})()");
-
-    buf.into_string()
-}
-
-pub(crate) fn compile_scoped(expr: &Ir, ctx: &impl CodegenContext) -> String {
-    let mut buf = CodeBuffer::with_capacity(8192);
-
-    code!(&mut buf, ctx; "((__scope)=>{");
-
-    if std::env::var("NIX_JS_DEBUG_THUNKS").is_ok() {
-        code!(&mut buf, ctx; "Nix.DEBUG_THUNKS.enabled=true;");
-    }
-
-    code!(&mut buf, ctx;
-        "Nix.builtins.storeDir="
-        quoted(ctx.get_store_dir())
-        ";const __currentDir="
-        quoted(&ctx.get_current_dir().display().to_string())
-        ";return "
-        expr
-        "})"
-    );
-
-    buf.into_string()
-}
-
 trait Compile<Ctx: CodegenContext> {
     fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer);
 }
@@ -318,6 +318,14 @@ impl<Ctx: CodegenContext> Compile<Ctx> for Ir {
                     "]"
                 );
             }
+            Ir::WithExpr(x) => x.compile(ctx, buf),
+            &Ir::WithLookup(WithLookup { inner: name, .. }) => {
+                code!(buf, ctx;
+                    "Nix.lookupWith("
+                    ctx.get_sym(name)
+                    ",__with)"
+                );
+            }
         }
     }
 }
@@ -522,6 +530,19 @@ impl<Ctx: CodegenContext> Compile<Ctx> for TopLevel {
     }
 }
 
+impl<Ctx: CodegenContext> Compile<Ctx> for WithExpr {
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        let namespace = ctx.get_ir(self.namespace);
+        let body = ctx.get_ir(self.body);
+        let has_thunks = !self.thunks.is_empty();
+        if has_thunks {
+            code!(buf, ctx; "((__with)=>{" self.thunks "return " body "})({env:" namespace ",last:__with})");
+        } else {
+            code!(buf, ctx; "((__with)=>(" body "))({env:" namespace ",last:__with})");
+        }
+    }
+}
+
 impl<Ctx: CodegenContext> Compile<Ctx> for Select {
     fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
         if let Some(default) = self.default {

nix-js/src/context.rs

@@ -11,7 +11,7 @@ use crate::downgrade::*;
 use crate::error::{Error, Result, Source};
 use crate::ir::{
     Arg, ArgId, Bool, Builtin, ExprId, Ir, Null, ReplBinding, ScopedImportBinding, SymId, Thunk,
-    ToIr as _,
+    ToIr as _, WithLookup,
 };
 use crate::runtime::{Runtime, RuntimeContext};
 use crate::store::{DaemonStore, Store, StoreConfig};
@@ -342,7 +342,6 @@ enum Scope<'ctx> {
     ScopedImport(HashSet<SymId>),
     Let(HashMap<SymId, ExprId>),
     Param(SymId, ExprId),
-    With(ExprId),
 }
 
 struct ScopeGuard<'a, 'ctx> {
@@ -365,6 +364,7 @@ pub struct DowngradeCtx<'ctx> {
     ctx: &'ctx mut Ctx,
     irs: Vec<Ir>,
     scopes: Vec<Scope<'ctx>>,
+    with_scope_count: usize,
     arg_id: usize,
     thunk_scopes: Vec<Vec<(ExprId, ExprId)>>,
 }
@@ -381,6 +381,7 @@ impl<'ctx> DowngradeCtx<'ctx> {
                 .collect(),
             irs: vec![],
             arg_id: 0,
+            with_scope_count: 0,
             thunk_scopes: vec![Vec::new()],
             ctx,
         }
@@ -472,39 +473,18 @@ impl DowngradeContext for DowngradeCtx<'_> {
                         return Ok(expr);
                     }
                 }
-                &Scope::With(_) => (),
             }
         }
 
-        let namespaces: Vec<ExprId> = self
+        if self.with_scope_count > 0 {
-            .scopes
+            Ok(self.new_expr(WithLookup { inner: sym, span }.to_ir()))
-            .iter()
+        } else {
-            .filter_map(|scope| {
+            Err(Error::downgrade_error(
-                if let &Scope::With(namespace) = scope {
-                    Some(namespace)
-                } else {
-                    None
-                }
-            })
-            .collect();
-        let mut result = None;
-        for namespace in namespaces {
-            use crate::ir::{Attr, Select};
-            let select = Select {
-                expr: namespace,
-                attrpath: vec![Attr::Str(sym, rnix::TextRange::default())],
-                default: result, // Link to outer With or None
-                span,
-            };
-            result = Some(self.new_expr(select.to_ir()));
-        }
-        result.ok_or_else(|| {
-            Error::downgrade_error(
                 format!("'{}' not found", self.get_sym(sym)),
                 self.get_current_source(),
                 span,
-            )
+            ))
-        })
+        }
     }
 
     fn replace_ir(&mut self, id: ExprId, expr: Ir) {
@@ -558,14 +538,14 @@ impl DowngradeContext for DowngradeCtx<'_> {
         f(guard.as_ctx())
     }
 
-    fn with_with_scope<F, R>(&mut self, namespace: ExprId, f: F) -> R
+    fn with_with_scope<F, R>(&mut self, f: F) -> R
     where
         F: FnOnce(&mut Self) -> R,
     {
-        let namespace = self.maybe_thunk(namespace);
+        self.with_scope_count += 1;
-        self.scopes.push(Scope::With(namespace));
+        let ret = f(self);
-        let mut guard = ScopeGuard { ctx: self };
+        self.with_scope_count -= 1;
-        f(guard.as_ctx())
+        ret
     }
 
     fn with_thunk_scope<F, R>(&mut self, f: F) -> (R, Vec<(ExprId, ExprId)>)

nix-js/src/downgrade.rs

@@ -35,7 +35,7 @@ pub trait DowngradeContext {
     fn with_let_scope<F, R>(&mut self, bindings: HashMap<SymId, ExprId>, f: F) -> R
     where
         F: FnOnce(&mut Self) -> R;
-    fn with_with_scope<F, R>(&mut self, namespace: ExprId, f: F) -> R
+    fn with_with_scope<F, R>(&mut self, f: F) -> R
     where
         F: FnOnce(&mut Self) -> R;
     fn with_thunk_scope<F, R>(&mut self, f: F) -> (R, Vec<(ExprId, ExprId)>)
@@ -402,13 +402,23 @@ impl<Ctx: DowngradeContext> Downgrade<Ctx> for ast::LetIn {
 /// Downgrades a `with` expression.
 impl<Ctx: DowngradeContext> Downgrade<Ctx> for ast::With {
     fn downgrade(self, ctx: &mut Ctx) -> Result<ExprId> {
-        // with namespace; expr
+        let span = self.syntax().text_range();
         let namespace = self.namespace().unwrap().downgrade(ctx)?;
+        let namespace = ctx.maybe_thunk(namespace);
 
-        // Downgrade body in With scope
+        let (body, thunks) = ctx
-        let expr = ctx.with_with_scope(namespace, |ctx| self.body().unwrap().downgrade(ctx))?;
+            .with_thunk_scope(|ctx| ctx.with_with_scope(|ctx| self.body().unwrap().downgrade(ctx)));
+        let body = body?;
 
-        Ok(expr)
+        Ok(ctx.new_expr(
+            WithExpr {
+                namespace,
+                body,
+                thunks,
+                span,
+            }
+            .to_ir(),
+        ))
     }
 }
 

nix-js/src/fetcher.rs

@@ -350,43 +350,16 @@ pub fn op_fetch_hg(
 }
 
 fn normalize_hash(hash: &str) -> String {
+    use base64::prelude::*;
     if hash.starts_with("sha256-")
         && let Some(b64) = hash.strip_prefix("sha256-")
-        && let Ok(bytes) = base64_decode(b64)
+        && let Ok(bytes) = BASE64_STANDARD.decode(b64)
     {
         return hex::encode(bytes);
     }
     hash.to_string()
 }
 
-fn base64_decode(input: &str) -> Result<Vec<u8>, String> {
-    const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
-
-    let input = input.trim_end_matches('=');
-    let mut output = Vec::with_capacity(input.len() * 3 / 4);
-
-    let mut buffer = 0u32;
-    let mut bits = 0;
-
-    for c in input.bytes() {
-        let value = ALPHABET
-            .iter()
-            .position(|&x| x == c)
-            .ok_or_else(|| format!("Invalid base64 character: {}", c as char))?;
-
-        buffer = (buffer << 6) | (value as u32);
-        bits += 6;
-
-        if bits >= 8 {
-            bits -= 8;
-            output.push((buffer >> bits) as u8);
-            buffer &= (1 << bits) - 1;
-        }
-    }
-
-    Ok(output)
-}
-
 pub fn register_ops<Ctx: RuntimeContext>() -> Vec<deno_core::OpDecl> {
     vec![
         op_fetch_url::<Ctx>(),

nix-js/src/ir.rs

@@ -34,6 +34,8 @@ ir! {
     CurPos,
     ReplBinding(SymId),
     ScopedImportBinding(SymId),
+    WithExpr { pub namespace: ExprId, pub body: ExprId, pub thunks: Vec<(ExprId, ExprId)> },
+    WithLookup(SymId),
 }
 
 #[repr(transparent)]