refactor(codegen): less allocation

nix-js/src/codegen.rs

@@ -1,34 +1,187 @@
+use std::fmt::{self, Write as _};
 use std::path::Path;
 
-use itertools::Itertools as _;
-
 use crate::ir::*;
 
-pub(crate) fn compile(expr: &Ir, ctx: &impl CodegenContext) -> String {
-    let code = expr.compile(ctx);
+pub(crate) struct CodeBuffer {
+    buf: String,
+}
 
-    let mut debug_flags = Vec::new();
-    if std::env::var("NIX_JS_DEBUG_THUNKS").is_ok() {
-        debug_flags.push("Nix.DEBUG_THUNKS.enabled=true");
+impl fmt::Write for CodeBuffer {
+    #[inline]
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        self.buf.push_str(s);
+        Ok(())
     }
-    let debug_prefix = if debug_flags.is_empty() {
-        String::new()
-    } else {
-        format!("{};", debug_flags.join(";"))
+}
+
+impl CodeBuffer {
+    #[inline]
+    fn with_capacity(capacity: usize) -> Self {
+        Self {
+            buf: String::with_capacity(capacity),
+        }
+    }
+
+    #[inline]
+    fn push_str(&mut self, s: &str) {
+        self.buf.push_str(s);
+    }
+
+    #[inline]
+    fn into_string(self) -> String {
+        self.buf
+    }
+}
+
+struct Quoted<'a>(&'a str);
+
+#[inline]
+fn quoted(s: &str) -> Quoted<'_> {
+    Quoted(s)
+}
+
+struct Escaped<'a>(&'a str);
+
+impl<Ctx: CodegenContext> Compile<Ctx> for Escaped<'_> {
+    fn compile(&self, _ctx: &Ctx, buf: &mut CodeBuffer) {
+        for c in self.0.chars() {
+            let _ = match c {
+                '\\' => buf.write_str("\\\\"),
+                '"' => buf.write_str("\\\""),
+                '\n' => buf.write_str("\\n"),
+                '\r' => buf.write_str("\\r"),
+                '\t' => buf.write_str("\\t"),
+                _ => buf.write_char(c),
+            };
+        }
+    }
+}
+
+#[inline]
+fn escaped(s: &str) -> Escaped<'_> {
+    Escaped(s)
+}
+
+struct Joined<I, F> {
+    items: I,
+    sep: &'static str,
+    write_fn: F,
+}
+
+#[inline]
+fn joined<Ctx: CodegenContext, I: Iterator, F: Fn(&Ctx, &mut CodeBuffer, I::Item)>(
+    items: I,
+    sep: &'static str,
+    write_fn: F,
+) -> Joined<I, F> {
+    Joined {
+        items,
+        sep,
+        write_fn,
+    }
+}
+
+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()
     };
 
-    let cur_dir = ctx.get_current_dir().display().to_string().escape_quote();
-    format!(
-        "(()=>{{{}Nix.builtins.storeDir={};const currentDir={};return {}}})()",
-        debug_prefix,
-        ctx.get_store_dir().escape_quote(),
-        cur_dir,
-        code
-    )
+    ($buf:expr, $lit:literal) => {
+        $buf.push_str($lit)
+    };
+}
+
+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()
 }
 
 trait Compile<Ctx: CodegenContext> {
-    fn compile(&self, ctx: &Ctx) -> String;
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer);
+}
+
+impl<Ctx: CodegenContext> Compile<Ctx> for str {
+    fn compile(&self, _ctx: &Ctx, buf: &mut CodeBuffer) {
+        buf.push_str(self);
+    }
+}
+
+impl<Ctx: CodegenContext> Compile<Ctx> for usize {
+    fn compile(&self, _ctx: &Ctx, buf: &mut CodeBuffer) {
+        let _ = write!(buf, "{self}");
+    }
+}
+
+impl<Ctx: CodegenContext> Compile<Ctx> for Quoted<'_> {
+    fn compile(&self, _ctx: &Ctx, buf: &mut CodeBuffer) {
+        buf.push_escaped_quote(self.0);
+    }
+}
+
+// Implement Compile for Span
+// impl<Ctx: CodegenContext> Compile<Ctx> for Span<'_, Ctx> {
+//     fn compile(&self, _ctx: &Ctx, buf: &mut CodeBuffer) {
+//         write!(buf, "{}", self).unwrap();
+//     }
+// }
+
+// Implement Compile for Joined
+impl<Ctx: CodegenContext, I, F> Compile<Ctx> for Joined<I, F>
+where
+    I: IntoIterator + Clone,
+    F: Fn(&Ctx, &mut CodeBuffer, I::Item) + Clone,
+{
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        let mut iter = self.items.clone().into_iter();
+        if let Some(first) = iter.next() {
+            (self.write_fn)(ctx, buf, first);
+            for item in iter {
+                buf.push_str(self.sep);
+                (self.write_fn)(ctx, buf, item);
+            }
+        }
+    }
+}
+
+impl<Ctx: CodegenContext> Compile<Ctx> for rnix::TextRange {
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        code!(
+            buf,
+            "\"{}:{}:{}\"",
+            ctx.get_current_source_id(),
+            usize::from(self.start()),
+            usize::from(self.end())
+        );
+    }
 }
 
 pub(crate) trait CodegenContext {
@@ -40,179 +193,224 @@ pub(crate) trait CodegenContext {
     fn get_current_source(&self) -> crate::error::Source;
 }
 
-trait EscapeQuote {
-    fn escape_quote(&self) -> String;
-}
-
-impl EscapeQuote for str {
-    fn escape_quote(&self) -> String {
-        let mut escaped = String::with_capacity(self.len() + 2);
-        escaped.push('"');
-        for c in self.chars() {
-            match c {
-                '\\' => escaped.push_str("\\\\"),
-                '\"' => escaped.push_str("\\\""),
-                '\n' => escaped.push_str("\\n"),
-                '\r' => escaped.push_str("\\r"),
-                '\t' => escaped.push_str("\\t"),
-                _ => escaped.push(c),
-            }
-        }
-        escaped.push('"');
-        escaped
+impl<Ctx: CodegenContext> Compile<Ctx> for ExprId {
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        ctx.get_ir(*self).compile(ctx, buf);
     }
 }
 
-fn encode_span(span: rnix::TextRange, ctx: &impl CodegenContext) -> String {
-    format!(
-        "\"{}:{}:{}\"",
-        ctx.get_current_source_id(),
-        usize::from(span.start()),
-        usize::from(span.end())
-    )
-}
-
 impl<Ctx: CodegenContext> Compile<Ctx> for Ir {
-    fn compile(&self, ctx: &Ctx) -> String {
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
         match self {
-            Ir::Int(int) => format!("{}n", int.inner), // Generate BigInt literal
-            Ir::Float(float) => float.inner.to_string(),
-            Ir::Bool(bool) => bool.inner.to_string(),
-            Ir::Null(_) => "null".to_string(),
-            Ir::Str(s) => s.val.escape_quote(),
+            Ir::Int(int) => {
+                code!(buf, "{}n", int.inner);
+            }
+            Ir::Float(float) => {
+                code!(buf, "{}", float.inner);
+            }
+            Ir::Bool(bool) => {
+                code!(buf, "{}", bool.inner);
+            }
+            Ir::Null(_) => {
+                code!(buf, ctx; "null");
+            }
+            Ir::Str(s) => {
+                code!(buf, ctx; quoted(&s.val));
+            }
             Ir::Path(p) => {
-                // Path needs runtime resolution
-                let path_expr = ctx.get_ir(p.expr).compile(ctx);
-                format!("Nix.resolvePath(currentDir,{})", path_expr)
+                code!(buf, ctx; "Nix.resolvePath(currentDir," ctx.get_ir(p.expr) ")");
             }
-            &Ir::If(If {
-                cond,
-                consq,
-                alter,
-                span: _,
-            }) => {
-                let cond_code = ctx.get_ir(cond).compile(ctx);
-                let consq = ctx.get_ir(consq).compile(ctx);
-                let alter = ctx.get_ir(alter).compile(ctx);
-                let cond_span = encode_span(ctx.get_ir(cond).span(), ctx);
-                format!(
-                    "(Nix.withContext(\"while evaluating a branch condition\",{},()=>Nix.forceBool({})))?({}):({})",
-                    cond_span, cond_code, consq, alter
-                )
+            Ir::If(x) => x.compile(ctx, buf),
+            Ir::BinOp(x) => x.compile(ctx, buf),
+            Ir::UnOp(x) => x.compile(ctx, buf),
+            Ir::Func(x) => x.compile(ctx, buf),
+            Ir::AttrSet(x) => x.compile(ctx, buf),
+            Ir::List(x) => x.compile(ctx, buf),
+            Ir::Call(x) => x.compile(ctx, buf),
+            Ir::Arg(x) => {
+                code!(buf, "arg{}", x.inner.0);
             }
-            Ir::BinOp(x) => x.compile(ctx),
-            Ir::UnOp(x) => x.compile(ctx),
-            Ir::Func(x) => x.compile(ctx),
-            Ir::AttrSet(x) => x.compile(ctx),
-            Ir::List(x) => x.compile(ctx),
-            Ir::Call(x) => x.compile(ctx),
-            Ir::Arg(x) => format!("arg{}", x.inner.0),
-            Ir::TopLevel(x) => x.compile(ctx),
-            Ir::Select(x) => x.compile(ctx),
+            Ir::TopLevel(x) => x.compile(ctx, buf),
+            Ir::Select(x) => x.compile(ctx, buf),
             &Ir::Thunk(Thunk { inner: expr_id, .. }) => {
-                format!("expr{}", expr_id.0)
+                code!(buf, "expr{}", expr_id.0);
+            }
+            Ir::Builtins(_) => {
+                code!(buf, ctx; "Nix.builtins");
             }
-            Ir::Builtins(_) => "Nix.builtins".to_string(),
             &Ir::Builtin(Builtin { inner: name, .. }) => {
-                format!("Nix.builtins[{}]", ctx.get_sym(name).escape_quote())
+                code!(buf, ctx;
+                    "Nix.builtins[",
+                    quoted(ctx.get_sym(name)),
+                    "]"
+                );
             }
-            Ir::ConcatStrings(x) => x.compile(ctx),
-            Ir::HasAttr(x) => x.compile(ctx),
+            Ir::ConcatStrings(x) => x.compile(ctx, buf),
+            Ir::HasAttr(x) => x.compile(ctx, buf),
             &Ir::Assert(Assert {
                 assertion,
                 expr,
                 ref assertion_raw,
-                span,
+                span: assert_span,
             }) => {
-                let assertion_code = ctx.get_ir(assertion).compile(ctx);
-                let expr = ctx.get_ir(expr).compile(ctx);
-                let assertion_span = encode_span(ctx.get_ir(assertion).span(), ctx);
-                let span = encode_span(span, ctx);
-                format!(
-                    "Nix.assert(Nix.withContext(\"while evaluating the condition of the assert statement\",{},()=>({})),{},{},{})",
+                let assertion_ir = ctx.get_ir(assertion);
+                let assertion_span = assertion_ir.span();
+
+                code!(buf, ctx;
+                    "Nix.assert(Nix.withContext(\"while evaluating the condition of the assert statement\",",
                     assertion_span,
-                    assertion_code,
-                    expr,
-                    assertion_raw.escape_quote(),
-                    span
-                )
+                    ",()=>(",
+                    assertion_ir,
+                    ")),",
+                    ctx.get_ir(expr),
+                    ",",
+                    quoted(assertion_raw),
+                    ",",
+                    assert_span,
+                    ")"
+                );
             }
             Ir::CurPos(cur_pos) => {
-                let span_str = encode_span(cur_pos.span, ctx);
-                format!("Nix.mkPos({})", span_str)
+                code!(buf, ctx;
+                    "Nix.mkPos(",
+                    cur_pos.span,
+                    ")"
+                );
             }
         }
     }
 }
 
+impl<Ctx: CodegenContext> Compile<Ctx> for If {
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        let &If {
+            cond,
+            consq,
+            alter,
+            span: _,
+        } = self;
+        let cond_ir = ctx.get_ir(cond);
+        let cond_span = cond_ir.span();
+
+        code!(buf, ctx;
+            "(Nix.withContext(\"while evaluating a branch condition\"," cond_span ",()=>Nix.forceBool(" cond_ir ")))"
+            "?(" consq "):(" alter ")"
+        );
+    }
+}
+
 impl<Ctx: CodegenContext> Compile<Ctx> for BinOp {
-    fn compile(&self, ctx: &Ctx) -> String {
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
         use BinOpKind::*;
 
-        let lhs = ctx.get_ir(self.lhs).compile(ctx);
-        let rhs = ctx.get_ir(self.rhs).compile(ctx);
-
-        let with_ctx = |op_name: &str, op_call: String| {
-            let span = encode_span(self.span, ctx);
-            format!(
-                "Nix.withContext(\"while evaluating the {} operator\",{},()=>({}))",
-                op_name, span, op_call
-            )
-        };
+        let lhs = ctx.get_ir(self.lhs);
+        let rhs = ctx.get_ir(self.rhs);
 
         match self.kind {
-            Add => with_ctx("+", format!("Nix.op.add({},{})", lhs, rhs)),
-            Sub => with_ctx("-", format!("Nix.op.sub({},{})", lhs, rhs)),
-            Mul => with_ctx("*", format!("Nix.op.mul({},{})", lhs, rhs)),
-            Div => with_ctx("/", format!("Nix.op.div({},{})", lhs, rhs)),
-            Eq => with_ctx("==", format!("Nix.op.eq({},{})", lhs, rhs)),
-            Neq => with_ctx("!=", format!("Nix.op.neq({},{})", lhs, rhs)),
-            Lt => with_ctx("<", format!("Nix.op.lt({},{})", lhs, rhs)),
-            Gt => with_ctx(">", format!("Nix.op.gt({},{})", lhs, rhs)),
-            Leq => with_ctx("<=", format!("Nix.op.lte({},{})", lhs, rhs)),
-            Geq => with_ctx(">=", format!("Nix.op.gte({},{})", lhs, rhs)),
-            // Short-circuit operators: use JavaScript native && and ||
-            And => with_ctx(
-                "&&",
-                format!("Nix.forceBool({})&&Nix.forceBool({})", lhs, rhs),
-            ),
-            Or => with_ctx(
-                "||",
-                format!("Nix.forceBool({})||Nix.forceBool({})", lhs, rhs),
-            ),
-            Impl => with_ctx(
-                "->",
-                format!("(!Nix.forceBool({})||Nix.forceBool({}))", lhs, rhs),
-            ),
-            Con => with_ctx("++", format!("Nix.op.concat({},{})", lhs, rhs)),
-            Upd => with_ctx("//", format!("Nix.op.update({},{})", lhs, rhs)),
-            PipeL => format!("Nix.call({},{})", rhs, lhs),
-            PipeR => format!("Nix.call({},{})", lhs, rhs),
+            Add | Sub | Mul | Div | Eq | Neq | Lt | Gt | Leq | Geq | Con | Upd => {
+                let op_name = match self.kind {
+                    Add => "+",
+                    Sub => "-",
+                    Mul => "*",
+                    Div => "/",
+                    Eq => "==",
+                    Neq => "!=",
+                    Lt => "<",
+                    Gt => ">",
+                    Leq => "<=",
+                    Geq => ">=",
+                    Con => "++",
+                    Upd => "//",
+                    _ => unreachable!(),
+                };
+                let op_func = match self.kind {
+                    Add => "Nix.op.add",
+                    Sub => "Nix.op.sub",
+                    Mul => "Nix.op.mul",
+                    Div => "Nix.op.div",
+                    Eq => "Nix.op.eq",
+                    Neq => "Nix.op.neq",
+                    Lt => "Nix.op.lt",
+                    Gt => "Nix.op.gt",
+                    Leq => "Nix.op.lte",
+                    Geq => "Nix.op.gte",
+                    Con => "Nix.op.concat",
+                    Upd => "Nix.op.update",
+                    _ => unreachable!(),
+                };
+
+                code!(
+                    buf, ctx;
+                    "Nix.withContext(\"while evaluating the " op_name " operator\"," self.span ",()=>(" op_func "(" lhs "," rhs ")))"
+                );
+            }
+            And => {
+                code!(
+                    buf, ctx;
+                    "Nix.withContext(\"while evaluating the && operator\"," self.span ",()=>(Nix.forceBool(" lhs ")&&Nix.forceBool(" rhs ")))"
+                );
+            }
+            Or => {
+                code!(
+                    buf, ctx;
+                    "Nix.withContext(\"while evaluating the || operator\"," self.span ",()=>(Nix.forceBool(" lhs ")||Nix.forceBool(" rhs ")))"
+                );
+            }
+            Impl => {
+                code!(
+                    buf, ctx;
+                    "Nix.withContext(\"while evaluating the || operator\"," self.span ",()=>(Nix.forceBool(" lhs ")||Nix.forceBool(" rhs ")))"
+                );
+            }
+            PipeL => {
+                code!(buf, ctx;
+                    "Nix.call(",
+                    rhs,
+                    ",",
+                    lhs,
+                    ")"
+                );
+            }
+            PipeR => {
+                code!(buf, ctx;
+                    "Nix.call(",
+                    lhs,
+                    ",",
+                    rhs,
+                    ")"
+                );
+            }
         }
     }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for UnOp {
-    fn compile(&self, ctx: &Ctx) -> String {
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
         use UnOpKind::*;
-        let rhs = ctx.get_ir(self.rhs).compile(ctx);
         match self.kind {
-            Neg => format!("Nix.op.sub(0n,{rhs})"),
-            Not => format!("Nix.op.bnot({rhs})"),
+            Neg => {
+                code!(buf, ctx;
+                    "Nix.op.sub(0n,",
+                    ctx.get_ir(self.rhs),
+                    ")"
+                );
+            }
+            Not => {
+                code!(buf, ctx;
+                    "Nix.op.bnot(",
+                    ctx.get_ir(self.rhs),
+                    ")"
+                );
+            }
         }
     }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for Func {
-    fn compile(&self, ctx: &Ctx) -> String {
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
         let id = ctx.get_ir(self.arg).as_ref().unwrap_arg().inner.0;
-        let thunk_defs = compile_thunks(&self.thunks, ctx);
-        let body_code = ctx.get_ir(self.body).compile(ctx);
-        let body = if thunk_defs.is_empty() {
-            body_code
-        } else {
-            format!("{{{}return {}}}", thunk_defs, body_code)
-        };
+
+        let has_thunks = !self.thunks.is_empty();
 
         if let Some(Param {
             required,
@@ -220,199 +418,230 @@ impl<Ctx: CodegenContext> Compile<Ctx> for Func {
             ellipsis,
         }) = &self.param
         {
-            let mut required = required.iter().map(|&sym| ctx.get_sym(sym).escape_quote());
-            let required = format!("[{}]", required.join(","));
-            let mut optional = optional.iter().map(|&sym| ctx.get_sym(sym).escape_quote());
-            let optional = format!("[{}]", optional.join(","));
-            if thunk_defs.is_empty() {
-                format!("Nix.mkFunction(arg{id}=>({body}),{required},{optional},{ellipsis})")
+            code!(buf, "Nix.mkFunction(arg{}=>", id);
+            if has_thunks {
+                code!(buf, ctx; "{", &self.thunks, "return ", self.body, "}");
             } else {
-                format!("Nix.mkFunction(arg{id}=>{body},{required},{optional},{ellipsis})")
+                code!(buf, ctx; "(", self.body, ")");
             }
+            code!(buf, ctx;
+                ",["
+                joined(required.iter(), ",", |ctx: &Ctx, buf, &sym| {
+                    code!(buf, ctx; quoted(ctx.get_sym(sym)));
+                })
+                "],["
+                joined(optional.iter(), ",", |ctx: &Ctx, buf, &sym| {
+                    code!(buf, ctx; quoted(ctx.get_sym(sym)));
+                })
+                "],"
+            );
+            code!(buf, "{})", ellipsis);
         } else {
-            if thunk_defs.is_empty() {
-                format!("arg{id}=>({body})")
+            code!(buf, "arg{}=>", id);
+            if has_thunks {
+                code!(buf, ctx; "{", &self.thunks, "return ", self.body, "}");
             } else {
-                format!("arg{id}=>{body}")
+                code!(buf, ctx; "(", self.body, ")");
             }
         }
     }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for Call {
-    fn compile(&self, ctx: &Ctx) -> String {
-        let func = ctx.get_ir(self.func).compile(ctx);
-        let arg = ctx.get_ir(self.arg).compile(ctx);
-        let span_str = encode_span(self.span, ctx);
-        format!("Nix.call({func},{arg},{span_str})")
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        code!(buf, ctx;
+            "Nix.call(",
+            ctx.get_ir(self.func),
+            ",",
+            ctx.get_ir(self.arg),
+            ",",
+            self.span,
+            ")"
+        );
     }
 }
 
-fn compile_thunks<Ctx: CodegenContext>(thunks: &[(ExprId, ExprId)], ctx: &Ctx) -> String {
-    if thunks.is_empty() {
-        return String::new();
-    }
-    thunks
-        .iter()
-        .map(|&(slot, inner)| {
-            let inner_code = ctx.get_ir(inner).compile(ctx);
-            let inner_span = ctx.get_ir(inner).span();
-            format!(
-                "let expr{}=Nix.createThunk(()=>({}),\"expr{} {}:{}:{}\")",
-                slot.0,
-                inner_code,
+impl<Ctx: CodegenContext> Compile<Ctx> for [(ExprId, ExprId)] {
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        if self.is_empty() {
+            return;
+        }
+
+        for &(slot, inner) in self {
+            let inner_ir = ctx.get_ir(inner);
+            let inner_span = inner_ir.span();
+
+            code!(buf, "let expr{}=Nix.createThunk(()=>(", slot.0);
+            inner_ir.compile(ctx, buf);
+            code!(
+                buf,
+                "),\"expr{} {}:{}:{}\");",
                 slot.0,
                 ctx.get_current_source().get_name(),
                 usize::from(inner_span.start()),
                 usize::from(inner_span.end())
-            )
-        })
-        .join(";")
-        + ";"
+            );
+        }
+    }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for TopLevel {
-    fn compile(&self, ctx: &Ctx) -> String {
-        let thunk_defs = compile_thunks(&self.thunks, ctx);
-        let body = ctx.get_ir(self.body).compile(ctx);
-        if thunk_defs.is_empty() {
-            body
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        if self.thunks.is_empty() {
+            ctx.get_ir(self.body).compile(ctx, buf);
         } else {
-            format!("(()=>{{{}return {}}})()", thunk_defs, body)
+            code!(buf, "(()=>{");
+            code!(buf, ctx; &self.thunks);
+            code!(buf, "return ");
+            ctx.get_ir(self.body).compile(ctx, buf);
+            code!(buf, "})()");
         }
     }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for Select {
-    fn compile(&self, ctx: &Ctx) -> String {
-        let lhs = ctx.get_ir(self.expr).compile(ctx);
-        let attrpath = self
-            .attrpath
-            .iter()
-            .map(|attr| match attr {
-                Attr::Str(sym, _) => ctx.get_sym(*sym).escape_quote(),
-                Attr::Dynamic(expr_id, _) => ctx.get_ir(*expr_id).compile(ctx),
-            })
-            .join(",");
-        let span_str = encode_span(self.span, ctx);
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
         if let Some(default) = self.default {
-            format!(
-                "Nix.selectWithDefault({lhs},[{attrpath}],{},{span_str})",
-                ctx.get_ir(default).compile(ctx)
-            )
+            code!(buf, ctx;
+                "Nix.selectWithDefault("
+                ctx.get_ir(self.expr)
+                ",["
+                joined(self.attrpath.iter(), ",", |ctx: &Ctx, buf, attr| {
+                    match attr {
+                        Attr::Str(sym, _) => code!(buf, ctx; quoted(ctx.get_sym(*sym))),
+                        Attr::Dynamic(expr_id, _) => code!(buf, ctx; ctx.get_ir(*expr_id)),
+                    }
+                })
+                "],"
+                ctx.get_ir(default)
+                ","
+                self.span
+                ")"
+            );
         } else {
-            format!("Nix.select({lhs},[{attrpath}],{span_str})")
+            code!(buf, ctx;
+                "Nix.select("
+                ctx.get_ir(self.expr)
+                ",["
+                joined(self.attrpath.iter(), ",", |ctx: &Ctx, buf, attr| {
+                    match attr {
+                        Attr::Str(sym, _) => code!(buf, ctx; quoted(ctx.get_sym(*sym))),
+                        Attr::Dynamic(expr_id, _) => code!(buf, ctx; ctx.get_ir(*expr_id)),
+                    }
+                })
+                "],"
+                self.span
+                ")"
+            );
         }
     }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for AttrSet {
-    fn compile(&self, ctx: &Ctx) -> String {
-        let mut attrs = Vec::new();
-        let mut attr_positions = Vec::new();
-
-        for (&sym, &(expr, attr_span)) in &self.stcs {
-            let key = ctx.get_sym(sym);
-            let value_code = ctx.get_ir(expr).compile(ctx);
-
-            let value_span = encode_span(ctx.get_ir(expr).span(), ctx);
-            let value = format!(
-                "Nix.withContext(\"while evaluating the attribute '{}'\",{},()=>({}))",
-                key, value_span, value_code
-            );
-            attrs.push(format!("{}:{}", key.escape_quote(), value));
-
-            let attr_pos_str = encode_span(attr_span, ctx);
-            attr_positions.push(format!("{}:{}", key.escape_quote(), attr_pos_str));
-        }
-
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
         if !self.dyns.is_empty() {
-            let (keys, vals, dyn_spans) = self
-                .dyns
-                .iter()
-                .map(|(key, val, attr_span)| {
-                    let key = ctx.get_ir(*key).compile(ctx);
-                    let val_expr = ctx.get_ir(*val);
-                    let val = val_expr.compile(ctx);
-                    let span = val_expr.span();
-                    let span = encode_span(span, ctx);
-                    let val = format!(
-                        "Nix.withContext(\"while evaluating a dynamic attribute\",{},()=>({}))",
-                        span, val
+            code!(buf, ctx;
+                "Nix.mkAttrsWithPos({"
+                joined(self.stcs.iter(), ",", |ctx: &Ctx, buf, (&sym, &(expr, _))| {
+                    let key = ctx.get_sym(sym);
+                    let val = ctx.get_ir(expr);
+
+                    code!(
+                        buf, ctx;
+                        quoted(key) ":Nix.withContext(\"while evaluating the attribute '" escaped(key) "'\"," val.span() ",()=>(" val "))"
                     );
-                    let dyn_span_str = encode_span(*attr_span, ctx);
-                    (key, val, dyn_span_str)
                 })
-                .multiunzip::<(Vec<_>, Vec<_>, Vec<_>)>();
-            format!(
-                "Nix.mkAttrsWithPos({{{}}},{{{}}},{{dynKeys:[{}],dynVals:[{}],dynSpans:[{}]}})",
-                attrs.join(","),
-                attr_positions.join(","),
-                keys.join(","),
-                vals.join(","),
-                dyn_spans.join(",")
-            )
-        } else if !attr_positions.is_empty() {
-            format!(
-                "Nix.mkAttrsWithPos({{{}}},{{{}}})",
-                attrs.join(","),
-                attr_positions.join(",")
-            )
+                "},{"
+                joined(self.stcs.iter(), ",", |ctx: &Ctx, buf, (&sym, &(_, span))| {
+                    code!(buf, ctx; quoted(ctx.get_sym(sym)) ":" span);
+                })
+                "},{dynKeys:["
+                joined(self.dyns.iter(), ",", |ctx: &Ctx, buf, (key, _, _)| {
+                    code!(buf, ctx; ctx.get_ir(*key));
+                })
+                "],dynVals:["
+                joined(self.dyns.iter(), ",", |ctx: &Ctx, buf, (_, val, _)| {
+                    let val = ctx.get_ir(*val);
+                    code!(
+                        buf, ctx;
+                        "Nix.withContext(\"while evaluating a dynamic attribute\"," val.span() ",()=>(" val "))"
+                    );
+                })
+                "],dynSpans:["
+                joined(self.dyns.iter(), ",", |ctx: &Ctx, buf, (_, _, attr_span)| {
+                    code!(buf, ctx; attr_span);
+                })
+                "]})"
+            );
+        } else if !self.stcs.is_empty() {
+            code!(buf, ctx;
+                "Nix.mkAttrsWithPos({"
+                joined(self.stcs.iter(), ",", |ctx: &Ctx, buf, (&sym, &(expr, _))| {
+                    let key = ctx.get_sym(sym);
+                    let val = ctx.get_ir(expr);
+
+                    code!(
+                        buf, ctx;
+                        quoted(key) ":Nix.withContext(\"while evaluating the attribute '" escaped(key) "'\"," val.span() ",()=>(" val "))"
+                    );
+                })
+                "},{"
+                joined(self.stcs.iter(), ",", |ctx: &Ctx, buf, (&sym, &(_, span))| {
+                    code!(buf, ctx; quoted(ctx.get_sym(sym)) ":" span);
+                })
+                "})"
+            );
         } else {
-            format!("{{{}}}", attrs.join(","))
+            code!(buf, ctx; "{}");
         }
     }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for List {
-    fn compile(&self, ctx: &Ctx) -> String {
-        let list = self
-            .items
-            .iter()
-            .enumerate()
-            .map(|(idx, item)| {
-                let item_code = ctx.get_ir(*item).compile(ctx);
-                let item_span = encode_span(ctx.get_ir(*item).span(), ctx);
-                format!(
-                    "Nix.withContext(\"while evaluating list element {}\",{},()=>({}))",
-                    idx, item_span, item_code
-                )
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        code!(buf, ctx;
+            "["
+            joined(self.items.iter().enumerate(), ",", |ctx: &Ctx, buf, (idx, item)| {
+                let item = ctx.get_ir(*item);
+                code!(
+                    buf, ctx;
+                    "Nix.withContext(\"while evaluating list element " idx "\"," item.span() ",()=>(" item "))"
+                );
             })
-            .join(",");
-        format!("[{list}]")
+            "]"
+        );
     }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for ConcatStrings {
-    fn compile(&self, ctx: &Ctx) -> String {
-        let parts: Vec<String> = self
-            .parts
-            .iter()
-            .map(|part| {
-                let part_code = ctx.get_ir(*part).compile(ctx);
-                let part_span = encode_span(ctx.get_ir(*part).span(), ctx);
-                format!(
-                    "Nix.withContext(\"while evaluating a path segment\",{},()=>({}))",
-                    part_span, part_code
-                )
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        code!(buf, ctx;
+            "Nix.concatStringsWithContext(["
+            joined(self.parts.iter(), ",", |ctx: &Ctx, buf, part| {
+                let part = ctx.get_ir(*part);
+                code!(
+                    buf, ctx;
+                    "Nix.withContext(\"while evaluating a path segment\"," part.span() ",()=>(" part "))"
+                );
             })
-            .collect();
-
-        format!("Nix.concatStringsWithContext([{}])", parts.join(","))
+            "])"
+        );
     }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for HasAttr {
-    fn compile(&self, ctx: &Ctx) -> String {
-        let lhs = ctx.get_ir(self.lhs).compile(ctx);
-        let attrpath = self
-            .rhs
-            .iter()
-            .map(|attr| match attr {
-                Attr::Str(sym, _) => ctx.get_sym(*sym).escape_quote(),
-                Attr::Dynamic(expr_id, _) => ctx.get_ir(*expr_id).compile(ctx),
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        code!(buf, ctx;
+            "Nix.hasAttr("
+            ctx.get_ir(self.lhs)
+            ",["
+            joined(self.rhs.iter(), ",", |ctx: &Ctx, buf, attr| {
+                match attr {
+                    Attr::Str(sym, _) => code!(buf, ctx; quoted(ctx.get_sym(*sym))),
+                    Attr::Dynamic(expr_id, _) => code!(buf, ctx; ctx.get_ir(*expr_id)),
+                }
             })
-            .join(",");
-        format!("Nix.hasAttr({lhs},[{attrpath}])")
+            "])"
+        );
     }
 }

nix-js/src/ir/utils.rs

@@ -4,8 +4,8 @@
 use hashbrown::hash_map::Entry;
 use hashbrown::{HashMap, HashSet};
 use itertools::Itertools as _;
-use rnix::ast;
 use rnix::TextRange;
+use rnix::ast;
 use rowan::ast::AstNode;
 
 use crate::error::{Error, Result};
@@ -388,11 +388,11 @@ where
                             ));
                         }
                     }
-                } else if attrs_vec.len() > 1 {
-                    if let Some(ast::Attr::Ident(ident)) = attrs_vec.first() {
-                        let sym = ctx.new_sym(ident.to_string());
-                        binding_syms.insert(sym);
-                    }
+                } else if attrs_vec.len() > 1
+                    && let Some(ast::Attr::Ident(ident)) = attrs_vec.first()
+                {
+                    let sym = ctx.new_sym(ident.to_string());
+                    binding_syms.insert(sym);
                 }
             }
         }