refactor(codegen): less allocation

nix-js/src/codegen.rs

@@ -1,34 +1,158 @@
+use std::fmt::{self, Display, 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);
-
-    let mut debug_flags = Vec::new();
-    if std::env::var("NIX_JS_DEBUG_THUNKS").is_ok() {
-        debug_flags.push("Nix.DEBUG_THUNKS.enabled=true");
+/// High-performance code buffer for JavaScript generation
+pub(crate) struct CodeBuffer {
+    buf: String,
 }
-    let debug_prefix = if debug_flags.is_empty() {
-        String::new()
-    } else {
-        format!("{};", debug_flags.join(";"))
-    };
 
-    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
+impl fmt::Write for CodeBuffer {
+    #[inline]
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        self.buf.push_str(s);
+        Ok(())
+    }
+}
+
+impl CodeBuffer {
+    #[inline]
+    fn with_capacity(capacity: usize) -> Self {
+        Self {
+            buf: String::with_capacity(capacity),
+        }
+    }
+
+    #[inline]
+    pub fn push_str(&mut self, s: &str) {
+        self.buf.push_str(s);
+    }
+
+    #[inline]
+    pub fn push(&mut self, c: char) {
+        self.buf.push(c);
+    }
+
+    #[inline]
+    pub fn push_escaped_quote(&mut self, s: &str) {
+        self.buf.push('"');
+        for c in s.chars() {
+            match c {
+                '\\' => self.buf.push_str("\\\\"),
+                '"' => self.buf.push_str("\\\""),
+                '\n' => self.buf.push_str("\\n"),
+                '\r' => self.buf.push_str("\\r"),
+                '\t' => self.buf.push_str("\\t"),
+                _ => self.buf.push(c),
+            }
+        }
+        self.buf.push('"');
+    }
+
+    /// Write items joined by separator, using a closure to write each item
+    #[inline]
+    pub fn write_join<I, F>(&mut self, mut items: I, sep: &str, mut f: F)
+    where
+        I: Iterator,
+        F: FnMut(&mut Self, I::Item),
+    {
+        if let Some(first) = items.next() {
+            f(self, first);
+            for item in items {
+                self.buf.push_str(sep);
+                f(self, item);
+            }
+        }
+    }
+
+    #[inline]
+    fn into_string(self) -> String {
+        self.buf
+    }
+}
+
+/// Helper type for automatically escaping and quoting strings
+struct Quoted<'a>(&'a str);
+
+impl Display for Quoted<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_char('"')?;
+        for c in self.0.chars() {
+            match c {
+                '\\' => f.write_str("\\\\")?,
+                '"' => f.write_str("\\\"")?,
+                '\n' => f.write_str("\\n")?,
+                '\r' => f.write_str("\\r")?,
+                '\t' => f.write_str("\\t")?,
+                _ => f.write_char(c)?,
+            }
+        }
+        f.write_char('"')
+    }
+}
+
+/// Helper function to create a quoted string
+#[inline]
+#[allow(dead_code)]
+fn quoted(s: &str) -> Quoted<'_> {
+    Quoted(s)
+}
+
+/// Helper type for formatting spans
+struct Span<'a, Ctx> {
+    span: rnix::TextRange,
+    ctx: &'a Ctx,
+}
+
+impl<Ctx: CodegenContext> Display for Span<'_, Ctx> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "\"{}:{}:{}\"",
+            self.ctx.get_current_source_id(),
+            usize::from(self.span.start()),
+            usize::from(self.span.end())
         )
     }
+}
+
+/// Helper function to create a span formatter
+#[inline]
+#[allow(dead_code)]
+fn span<Ctx>(s: rnix::TextRange, ctx: &Ctx) -> Span<'_, Ctx> {
+    Span { span: s, ctx }
+}
+
+/// Ergonomic macro for code generation
+macro_rules! code {
+    ($buf:expr, $($arg:tt)*) => {
+        write!($buf, $($arg)*).unwrap()
+    };
+}
+
+pub(crate) fn compile(expr: &Ir, ctx: &impl CodegenContext) -> String {
+    let mut buf = CodeBuffer::with_capacity(8192);
+
+    buf.push_str("(()=>{");
+
+    if std::env::var("NIX_JS_DEBUG_THUNKS").is_ok() {
+        buf.push_str("Nix.DEBUG_THUNKS.enabled=true;");
+    }
+
+    buf.push_str("Nix.builtins.storeDir=");
+    buf.push_escaped_quote(ctx.get_store_dir());
+    buf.push_str(";const currentDir=");
+    buf.push_escaped_quote(&ctx.get_current_dir().display().to_string());
+    buf.push_str(";return ");
+    expr.compile(ctx, &mut buf);
+    buf.push_str("})()");
+
+    buf.into_string()
+}
 
 trait Compile<Ctx: CodegenContext> {
-    fn compile(&self, ctx: &Ctx) -> String;
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer);
 }
 
 pub(crate) trait CodegenContext {
@@ -40,50 +164,28 @@ 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
-    }
-}
-
-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(_) => {
+                buf.push_str("null");
+            }
+            Ir::Str(s) => {
+                buf.push_escaped_quote(&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)
+                buf.push_str("Nix.resolvePath(currentDir,");
+                ctx.get_ir(p.expr).compile(ctx, buf);
+                buf.push(')');
             }
             &Ir::If(If {
                 cond,
@@ -91,128 +193,195 @@ impl<Ctx: CodegenContext> Compile<Ctx> for Ir {
                 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
-                )
+                let cond_ir = ctx.get_ir(cond);
+                let cond_span = cond_ir.span();
+
+                code!(
+                    buf,
+                    "(Nix.withContext(\"while evaluating a branch condition\",{},()=>Nix.forceBool(",
+                    span(cond_span, ctx)
+                );
+                cond_ir.compile(ctx, buf);
+                buf.push_str(")))?(");
+                ctx.get_ir(consq).compile(ctx, buf);
+                buf.push_str("):(");
+                ctx.get_ir(alter).compile(ctx, buf);
+                buf.push(')');
             }
-            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::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::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(_) => {
+                buf.push_str("Nix.builtins");
             }
-            Ir::Builtins(_) => "Nix.builtins".to_string(),
             &Ir::Builtin(Builtin { inner: name, .. }) => {
-                format!("Nix.builtins[{}]", ctx.get_sym(name).escape_quote())
+                code!(buf, "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\",{},()=>({})),{},{},{})",
-                    assertion_span,
-                    assertion_code,
-                    expr,
-                    assertion_raw.escape_quote(),
-                    span
-                )
+                let assertion_ir = ctx.get_ir(assertion);
+                let assertion_span = assertion_ir.span();
+
+                code!(
+                    buf,
+                    "Nix.assert(Nix.withContext(\"while evaluating the condition of the assert statement\",{},()=>(",
+                    span(assertion_span, ctx)
+                );
+                assertion_ir.compile(ctx, buf);
+                buf.push_str(")),");
+                ctx.get_ir(expr).compile(ctx, buf);
+                code!(buf, ",{},{})", quoted(assertion_raw), span(assert_span, ctx));
             }
             Ir::CurPos(cur_pos) => {
-                let span_str = encode_span(cur_pos.span, ctx);
-                format!("Nix.mkPos({})", span_str)
+                code!(buf, "Nix.mkPos({})", span(cur_pos.span, ctx));
             }
         }
     }
 }
 
 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,
+                    "Nix.withContext(\"while evaluating the {} operator\",{},()=>({}",
+                    op_name,
+                    span(self.span, ctx),
+                    op_func
+                );
+                lhs.compile(ctx, buf);
+                buf.push(',');
+                rhs.compile(ctx, buf);
+                buf.push_str(")))");
+            }
+            And => {
+                code!(
+                    buf,
+                    "Nix.withContext(\"while evaluating the && operator\",{},()=>(Nix.forceBool(",
+                    span(self.span, ctx)
+                );
+                lhs.compile(ctx, buf);
+                buf.push_str(")&&Nix.forceBool(");
+                rhs.compile(ctx, buf);
+                buf.push_str(")))");
+            }
+            Or => {
+                code!(
+                    buf,
+                    "Nix.withContext(\"while evaluating the || operator\",{},()=>(Nix.forceBool(",
+                    span(self.span, ctx)
+                );
+                lhs.compile(ctx, buf);
+                buf.push_str(")||Nix.forceBool(");
+                rhs.compile(ctx, buf);
+                buf.push_str(")))");
+            }
+            Impl => {
+                code!(
+                    buf,
+                    "Nix.withContext(\"while evaluating the -> operator\",{},()=>((!Nix.forceBool(",
+                    span(self.span, ctx)
+                );
+                lhs.compile(ctx, buf);
+                buf.push_str(")||Nix.forceBool(");
+                rhs.compile(ctx, buf);
+                buf.push_str("))))");
+            }
+            PipeL => {
+                buf.push_str("Nix.call(");
+                rhs.compile(ctx, buf);
+                buf.push(',');
+                lhs.compile(ctx, buf);
+                buf.push(')');
+            }
+            PipeR => {
+                buf.push_str("Nix.call(");
+                lhs.compile(ctx, buf);
+                buf.push(',');
+                rhs.compile(ctx, buf);
+                buf.push(')');
+            }
         }
     }
 }
 
 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 => {
+                buf.push_str("Nix.op.sub(0n,");
+                ctx.get_ir(self.rhs).compile(ctx, buf);
+                buf.push(')');
+            }
+            Not => {
+                buf.push_str("Nix.op.bnot(");
+                ctx.get_ir(self.rhs).compile(ctx, buf);
+                buf.push(')');
+            }
         }
     }
 }
 
 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 +389,268 @@ 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 {
+                buf.push('{');
+                compile_thunks(&self.thunks, ctx, buf);
+                buf.push_str("return ");
+                ctx.get_ir(self.body).compile(ctx, buf);
+                buf.push('}');
             } else {
-                format!("Nix.mkFunction(arg{id}=>{body},{required},{optional},{ellipsis})")
+                buf.push('(');
+                ctx.get_ir(self.body).compile(ctx, buf);
+                buf.push(')');
             }
+            buf.push_str(",[");
+            buf.write_join(required.iter(), ",", |b, &sym| {
+                b.push_escaped_quote(ctx.get_sym(sym));
+            });
+            buf.push_str("],[");
+            buf.write_join(optional.iter(), ",", |b, &sym| {
+                b.push_escaped_quote(ctx.get_sym(sym));
+            });
+            code!(buf, "],{})", ellipsis);
         } else {
-            if thunk_defs.is_empty() {
-                format!("arg{id}=>({body})")
+            code!(buf, "arg{}=>", id);
+            if has_thunks {
+                buf.push('{');
+                compile_thunks(&self.thunks, ctx, buf);
+                buf.push_str("return ");
+                ctx.get_ir(self.body).compile(ctx, buf);
+                buf.push('}');
             } else {
-                format!("arg{id}=>{body}")
+                buf.push('(');
+                ctx.get_ir(self.body).compile(ctx, buf);
+                buf.push(')');
             }
         }
     }
 }
 
 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) {
+        buf.push_str("Nix.call(");
+        ctx.get_ir(self.func).compile(ctx, buf);
+        buf.push(',');
+        ctx.get_ir(self.arg).compile(ctx, buf);
+        code!(buf, ",{})", span(self.span, ctx));
     }
 }
 
-fn compile_thunks<Ctx: CodegenContext>(thunks: &[(ExprId, ExprId)], ctx: &Ctx) -> String {
+fn compile_thunks<Ctx: CodegenContext>(
+    thunks: &[(ExprId, ExprId)],
+    ctx: &Ctx,
+    buf: &mut CodeBuffer,
+) {
     if thunks.is_empty() {
-        return String::new();
+        return;
     }
-    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,
+
+    for &(slot, inner) in thunks {
+        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)
+            buf.push_str("(()=>{");
+            compile_thunks(&self.thunks, ctx, buf);
+            buf.push_str("return ");
+            ctx.get_ir(self.body).compile(ctx, buf);
+            buf.push_str("})()");
         }
     }
 }
 
 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)
-            )
+            buf.push_str("Nix.selectWithDefault(");
+            ctx.get_ir(self.expr).compile(ctx, buf);
+            buf.push_str(",[");
+            buf.write_join(self.attrpath.iter(), ",", |b, attr| match attr {
+                Attr::Str(sym, _) => code!(b, "{}", quoted(ctx.get_sym(*sym))),
+                Attr::Dynamic(expr_id, _) => ctx.get_ir(*expr_id).compile(ctx, b),
+            });
+            buf.push_str("],");
+            ctx.get_ir(default).compile(ctx, buf);
+            code!(buf, ",{})", span(self.span, ctx));
         } else {
-            format!("Nix.select({lhs},[{attrpath}],{span_str})")
+            buf.push_str("Nix.select(");
+            ctx.get_ir(self.expr).compile(ctx, buf);
+            buf.push_str(",[");
+            buf.write_join(self.attrpath.iter(), ",", |b, attr| match attr {
+                Attr::Str(sym, _) => code!(b, "{}", quoted(ctx.get_sym(*sym))),
+                Attr::Dynamic(expr_id, _) => ctx.get_ir(*expr_id).compile(ctx, b),
+            });
+            code!(buf, "],{})", span(self.span, ctx));
         }
     }
 }
 
 impl<Ctx: CodegenContext> Compile<Ctx> for AttrSet {
-    fn compile(&self, ctx: &Ctx) -> String {
-        let mut attrs = Vec::new();
-        let mut attr_positions = Vec::new();
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        if !self.dyns.is_empty() {
+            buf.push_str("Nix.mkAttrsWithPos({");
+
+            let mut first = true;
+            for (&sym, &(expr, _)) in &self.stcs {
+                if !first {
+                    buf.push(',');
+                }
+                first = false;
 
-        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_ir = ctx.get_ir(expr);
 
-            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
+                code!(
+                    buf,
+                    "{}:Nix.withContext(\"while evaluating the attribute '{}'\",{},()=>(",
+                    quoted(key),
+                    key,
+                    span(value_ir.span(), ctx)
                 );
-            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));
+                value_ir.compile(ctx, buf);
+                buf.push_str("))");
             }
 
-        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
+            buf.push_str("},{");
+
+            let mut first = true;
+            for (&sym, &(_, attr_span)) in &self.stcs {
+                if !first {
+                    buf.push(',');
+                }
+                first = false;
+
+                code!(buf, "{}:{}", quoted(ctx.get_sym(sym)), span(attr_span, ctx));
+            }
+
+            buf.push_str("},{dynKeys:[");
+            buf.write_join(self.dyns.iter(), ",", |b, (key, _, _)| {
+                ctx.get_ir(*key).compile(ctx, b);
+            });
+
+            buf.push_str("],dynVals:[");
+            buf.write_join(self.dyns.iter(), ",", |b, (_, val, _)| {
+                let val_ir = ctx.get_ir(*val);
+                code!(
+                    b,
+                    "Nix.withContext(\"while evaluating a dynamic attribute\",{},()=>(",
+                    span(val_ir.span(), ctx)
                 );
-                    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(",")
-            )
+                val_ir.compile(ctx, b);
+                b.push_str("))");
+            });
+
+            buf.push_str("],dynSpans:[");
+            buf.write_join(self.dyns.iter(), ",", |b, (_, _, attr_span)| {
+                code!(b, "{}", span(*attr_span, ctx));
+            });
+
+            buf.push_str("]})");
+        } else if !self.stcs.is_empty() {
+            buf.push_str("Nix.mkAttrsWithPos({");
+
+            let mut first = true;
+            for (&sym, &(expr, _)) in &self.stcs {
+                if !first {
+                    buf.push(',');
+                }
+                first = false;
+
+                let key = ctx.get_sym(sym);
+                let value_ir = ctx.get_ir(expr);
+
+                code!(
+                    buf,
+                    "{}:Nix.withContext(\"while evaluating the attribute '{}'\",{},()=>(",
+                    quoted(key),
+                    key,
+                    span(value_ir.span(), ctx)
+                );
+                value_ir.compile(ctx, buf);
+                buf.push_str("))");
+            }
+
+            buf.push_str("},{");
+
+            let mut first = true;
+            for (&sym, &(_, attr_span)) in &self.stcs {
+                if !first {
+                    buf.push(',');
+                }
+                first = false;
+
+                code!(buf, "{}:{}", quoted(ctx.get_sym(sym)), span(attr_span, ctx));
+            }
+
+            buf.push_str("})");
         } else {
-            format!("{{{}}}", attrs.join(","))
+            buf.push_str("{}");
         }
     }
 }
 
 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
-                )
-            })
-            .join(",");
-        format!("[{list}]")
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        buf.push('[');
+        buf.write_join(self.items.iter().enumerate(), ",", |b, (idx, item)| {
+            let item_ir = ctx.get_ir(*item);
+            code!(
+                b,
+                "Nix.withContext(\"while evaluating list element {}\",{},()=>(",
+                idx,
+                span(item_ir.span(), ctx)
+            );
+            item_ir.compile(ctx, b);
+            b.push_str("))");
+        });
+        buf.push(']');
     }
 }
 
 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
-                )
-            })
-            .collect();
-
-        format!("Nix.concatStringsWithContext([{}])", parts.join(","))
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        buf.push_str("Nix.concatStringsWithContext([");
+        buf.write_join(self.parts.iter(), ",", |b, part| {
+            let part_ir = ctx.get_ir(*part);
+            code!(
+                b,
+                "Nix.withContext(\"while evaluating a path segment\",{},()=>(",
+                span(part_ir.span(), ctx)
+            );
+            part_ir.compile(ctx, b);
+            b.push_str("))");
+        });
+        buf.push_str("])");
     }
 }
 
 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),
-            })
-            .join(",");
-        format!("Nix.hasAttr({lhs},[{attrpath}])")
+    fn compile(&self, ctx: &Ctx, buf: &mut CodeBuffer) {
+        buf.push_str("Nix.hasAttr(");
+        ctx.get_ir(self.lhs).compile(ctx, buf);
+        buf.push_str(",[");
+        buf.write_join(self.rhs.iter(), ",", |b, attr| match attr {
+            Attr::Str(sym, _) => code!(b, "{}", quoted(ctx.get_sym(*sym))),
+            Attr::Dynamic(expr_id, _) => ctx.get_ir(*expr_id).compile(ctx, b),
+        });
+        buf.push_str("])");
     }
 }