Repository: mitsuhiko/minijinja
Pull request #109 (2022-09-05 22:17:06)

Add finalizer support and custom escape types


Patch:
diff --git a/minijinja/src/environment.rs b/minijinja/src/environment.rs
index 71df36a..637a816 100644
--- a/minijinja/src/environment.rs
+++ b/minijinja/src/environment.rs
@@ -1,5 +1,5 @@
 use std::collections::BTreeMap;
-use std::fmt::{self, Write};
+use std::fmt;
 use std::sync::Arc;
 
 use serde::Serialize;
@@ -7,11 +7,12 @@ use serde::Serialize;
 use crate::compiler::Compiler;
 use crate::error::Error;
 use crate::instructions::Instructions;
+use crate::output::Output;
 use crate::parser::{parse, parse_expr};
 use crate::utils::{AutoEscape, BTreeMapKeysDebug, HtmlEscape};
 use crate::value::{ArgType, FunctionArgs, FunctionResult, Value, ValueKind};
 use crate::vm::Vm;
-use crate::{filters, functions, tests};
+use crate::{filters, functions, tests, ErrorKind};
 
 #[cfg(test)]
 use similar_asserts::assert_eq;
@@ -50,6 +51,19 @@ pub(crate) struct CompiledTemplate<'source> {
     blocks: BTreeMap<&'source str, Instructions<'source>>,
 }
 
+impl<'source> CompiledTemplate<'source> {
+    fn eval(
+        &self,
+        env: &Environment,
+        root: Value,
+        out: &mut Output,
+    ) -> Result<Option<Value>, Error> {
+        let vm = Vm::new(env);
+        let blocks = &self.blocks;
+        vm.eval(&self.instructions, root, blocks, out)
+    }
+}
+
 impl<'env> fmt::Debug for CompiledTemplate<'env> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut ds = f.debug_struct("CompiledTemplate");
@@ -139,15 +153,8 @@ impl<'env> Template<'env> {
 
     fn _render(&self, root: Value) -> Result<String, Error> {
         let mut output = String::new();
-        let vm = Vm::new(self.env);
-        let blocks = &self.compiled.blocks;
-        vm.eval(
-            &self.compiled.instructions,
-            root,
-            blocks,
-            self.initial_auto_escape,
-            &mut output,
-        )?;
+        let mut formatter = Output::string(&mut output, self.initial_auto_escape);
+        self.compiled.eval(self.env, root, &mut formatter)?;
         Ok(output)
     }
 
@@ -186,6 +193,9 @@ impl<'source> fmt::Debug for Source<'source> {
     }
 }
 
+type AutoEscapeFunc = dyn Fn(&str) -> AutoEscape + Sync + Send;
+type FormatterFunc = dyn Fn(&mut Output, &Value) -> Result<(), Error> + Sync + Send;
+
 /// An abstraction that holds the engine configuration.
 ///
 /// This object holds the central configuration state for templates.  It is also
@@ -214,7 +224,8 @@ pub struct Environment<'source> {
     filters: BTreeMap<&'source str, filters::BoxedFilter>,
     tests: BTreeMap<&'source str, tests::BoxedTest>,
     pub(crate) globals: BTreeMap<&'source str, Value>,
-    default_auto_escape: Arc<dyn Fn(&str) -> AutoEscape + Sync + Send>,
+    default_auto_escape: Arc<AutoEscapeFunc>,
+    formatter: Arc<FormatterFunc>,
     #[cfg(feature = "debug")]
     debug: bool,
 }
@@ -236,19 +247,6 @@ impl<'source> fmt::Debug for Environment<'source> {
     }
 }
 
-fn default_auto_escape(name: &str) -> AutoEscape {
-    match name.rsplit('.').next() {
-        Some("html") | Some("htm") | Some("xml") => AutoEscape::Html,
-        #[cfg(feature = "json")]
-        Some("json") | Some("js") | Some("yaml") | Some("yml") => AutoEscape::Json,
-        _ => AutoEscape::None,
-    }
-}
-
-fn no_auto_escape(_: &str) -> AutoEscape {
-    AutoEscape::None
-}
-
 /// A handle to a compiled expression.
 ///
 /// An expression is created via the
@@ -293,16 +291,12 @@ impl<'env, 'source> Expression<'env, 'source> {
     }
 
     fn _eval(&self, root: Value) -> Result<Value, Error> {
-        let mut output = String::new();
-        let vm = Vm::new(self.env);
-        let blocks = BTreeMap::new();
-        Ok(vm
+        Ok(Vm::new(self.env)
             .eval(
                 &self.instructions,
                 root,
-                &blocks,
-                AutoEscape::None,
-                &mut output,
+                &BTreeMap::new(),
+                &mut Output::null(),
             )?
             .unwrap())
     }
@@ -321,7 +315,8 @@ impl<'source> Environment<'source> {
             filters: filters::get_builtin_filters(),
             tests: tests::get_builtin_tests(),
             globals: functions::get_globals(),
-            default_auto_escape: Arc::new(default_auto_escape),
+            default_auto_escape: Arc::new(default_auto_escape_callback),
+            formatter: Arc::new(escape_formatter),
             #[cfg(feature = "debug")]
             debug: cfg!(debug_assertions),
         }
@@ -338,6 +333,7 @@ impl<'source> Environment<'source> {
             tests: Default::default(),
             globals: Default::default(),
             default_auto_escape: Arc::new(no_auto_escape),
+            formatter: Arc::new(escape_formatter),
             #[cfg(feature = "debug")]
             debug: cfg!(debug_assertions),
         }
@@ -428,16 +424,8 @@ impl<'source> Environment<'source> {
         let name = "<string>";
         let compiled = CompiledTemplate::from_name_and_source(name, source)?;
         let mut output = String::new();
-        let vm = Vm::new(self);
-        let blocks = &compiled.blocks;
-        let initial_auto_escape = self.get_initial_auto_escape(name);
-        vm.eval(
-            &compiled.instructions,
-            root,
-            blocks,
-            initial_auto_escape,
-            &mut output,
-        )?;
+        let mut formatter = Output::string(&mut output, self.get_initial_auto_escape(name));
+        compiled.eval(self, root, &mut formatter)?;
         Ok(output)
     }
 
@@ -446,22 +434,67 @@ impl<'source> Environment<'source> {
     /// This function is invoked when templates are loaded from the environment
     /// to determine the default auto escaping behavior.  The function is
     /// invoked with the name of the template and can make an initial auto
-    /// escaping decision based on that.  The default implementation is to
-    /// turn on escaping depending on the file extension:
+    /// escaping decision based on that.  The default implementation
+    /// ([`default_auto_escape_callback`]) turn on escaping depending on the file
+    /// extension.
     ///
-    /// * [`Html`](AutoEscape::Html): `.html`, `.htm`, `.xml`
-    #[cfg_attr(
-        feature = "json",
-        doc = r" * [`Json`](AutoEscape::Json): `.json`, `.js`, `.yml`"
-    )]
-    /// * [`None`](AutoEscape::None): _all others_
-    pub fn set_auto_escape_callback<F: Fn(&str) -> AutoEscape + 'static + Sync + Send>(
-        &mut self,
-        f: F,
-    ) {
+    /// ```
+    /// # use minijinja::{Environment, AutoEscape};
+    /// # let mut env = Environment::new();
+    /// env.set_auto_escape_callback(|name| {
+    ///     if matches!(name.rsplit('.').next().unwrap_or(""), "html" | "htm" | "aspx") {
+    ///         AutoEscape::Html
+    ///     } else {
+    ///         AutoEscape::None
+    ///     }
+    /// });
+    /// ```
+    pub fn set_auto_escape_callback<F>(&mut self, f: F)
+    where
+        F: Fn(&str) -> AutoEscape + 'static + Sync + Send,
+    {
         self.default_auto_escape = Arc::new(f);
     }
 
+    /// Sets a different formatter function.
+    ///
+    /// The formatter is invoked to format the given value into the provided
+    /// [`Output`].  The default implementation is the [`escape_formatter`].
+    ///
+    /// When implementing a custom formatter it depends on if auto escaping
+    /// should be supported or not.  If auto escaping should be supported then
+    /// it's easiest to just wrap the default [`escape_formatter`].  The
+    /// following example swaps out `None` values before rendering for
+    /// `Undefined` which renders as an empty string instead.
+    ///
+    /// The current value of the auto escape flag can be retrieved directly
+    /// from the output with [`Output::auto_escape`].
+    ///
+    /// ```
+    /// # use minijinja::Environment;
+    /// # let mut env = Environment::new();
+    /// use minijinja::escape_formatter;
+    /// use minijinja::value::Value;
+    ///
+    /// env.set_formatter(|out, value| {
+    ///     escape_formatter(
+    ///         out,
+    ///         if value.is_none() {
+    ///             &Value::UNDEFINED
+    ///         } else {
+    ///             value
+    ///         },
+    ///     )
+    ///});
+    /// # assert_eq!(env.render_str("{{ none }}", ()).unwrap(), "");
+    /// ```
+    pub fn set_formatter<F>(&mut self, f: F)
+    where
+        F: Fn(&mut Output, &Value) -> Result<(), Error> + 'static + Sync + Send,
+    {
+        self.formatter = Arc::new(f);
+    }
+
     /// Enable or disable the debug mode.
     ///
     /// When the debug mode is enabled the engine will dump out some of the
@@ -616,59 +649,87 @@ impl<'source> Environment<'source> {
         (self.default_auto_escape)(name)
     }
 
-    pub(crate) fn escape(
-        &self,
-        value: &Value,
-        autoescape: AutoEscape,
-        out: &mut String,
-    ) -> Result<(), Error> {
-        // safe values do not get escaped
-        if value.is_safe() {
-            write!(out, "{}", value).unwrap();
-            return Ok(());
-        }
-
-        // TODO: this should become pluggable
-        match autoescape {
-            AutoEscape::None => write!(out, "{}", value).unwrap(),
-            AutoEscape::Html => html_escape_value(out, value),
-            #[cfg(feature = "json")]
-            AutoEscape::Json => {
-                let value = serde_json::to_string(&value).map_err(|err| {
-                    Error::new(
-                        crate::ErrorKind::BadSerialization,
-                        "unable to format to JSON",
-                    )
-                    .with_source(err)
-                })?;
-                write!(out, "{}", value).unwrap()
-            }
-        }
-        Ok(())
-    }
-
-    /// Finalizes a value.
-    pub(crate) fn finalize(
-        &self,
-        value: &Value,
-        autoescape: AutoEscape,
-        out: &mut String,
-    ) -> Result<(), Error> {
-        self.escape(value, autoescape, out)
+    /// Formats a value into the final format.
+    ///
+    /// This step is called finalization in Jinja2 but since we are writing into
+    /// the output stream rather than converting values, it's renamed to format
+    /// here.
+    pub(crate) fn format(&self, value: &Value, out: &mut Output) -> Result<(), Error> {
+        (self.formatter)(out, value)
     }
 }
 
-pub fn html_escape_value(out: &mut String, value: &Value) {
+pub fn write_with_html_escaping(out: &mut Output, value: &Value) -> fmt::Result {
     if matches!(
         value.kind(),
         ValueKind::Undefined | ValueKind::None | ValueKind::Bool | ValueKind::Number
     ) {
-        write!(out, "{}", value).unwrap()
+        write!(out, "{}", value)
     } else if let Some(s) = value.as_str() {
-        write!(out, "{}", HtmlEscape(s)).unwrap()
+        write!(out, "{}", HtmlEscape(s))
     } else {
-        write!(out, "{}", HtmlEscape(&value.to_string())).unwrap()
+        write!(out, "{}", HtmlEscape(&value.to_string()))
+    }
+}
+
+fn no_auto_escape(_: &str) -> AutoEscape {
+    AutoEscape::None
+}
+
+/// The default logic for auto escaping based on file extension.
+///
+/// * [`Html`](AutoEscape::Html): `.html`, `.htm`, `.xml`
+#[cfg_attr(
+    feature = "json",
+    doc = r" * [`Json`](AutoEscape::Json): `.json`, `.js`, `.yml`"
+)]
+/// * [`None`](AutoEscape::None): _all others_
+pub fn default_auto_escape_callback(name: &str) -> AutoEscape {
+    match name.rsplit('.').next() {
+        Some("html") | Some("htm") | Some("xml") => AutoEscape::Html,
+        #[cfg(feature = "json")]
+        Some("json") | Some("js") | Some("yaml") | Some("yml") => AutoEscape::Json,
+        _ => AutoEscape::None,
+    }
+}
+
+/// The default formatter.
+///
+/// This formatter takes a value and directly writes it into the output format
+/// while honoring the requested auto escape format of the output.  If the
+/// value is already marked as safe, it's handled as if no auto escaping
+/// was requested.
+///
+/// * [`Html`](AutoEscape::Html): performs HTML escaping
+#[cfg_attr(
+    feature = "json",
+    doc = r" * [`Json`](AutoEscape::Json): serializes values to JSON"
+)]
+/// * [`None`](AutoEscape::None): no escaping
+/// * [`Custom(..)`](AutoEscape::Custom): results in an error
+pub fn escape_formatter(out: &mut Output, value: &Value) -> Result<(), Error> {
+    match (value.is_safe(), out.auto_escape()) {
+        // safe values do not get escaped
+        (true, _) | (_, AutoEscape::None) => write!(out, "{}", value)?,
+        (false, AutoEscape::Html) => write_with_html_escaping(out, value)?,
+        #[cfg(feature = "json")]
+        (false, AutoEscape::Json) => {
+            let value = serde_json::to_string(&value).map_err(|err| {
+                Error::new(ErrorKind::BadSerialization, "unable to format to JSON").with_source(err)
+            })?;
+            write!(out, "{}", value)?
+        }
+        (false, AutoEscape::Custom(name)) => {
+            return Err(Error::new(
+                ErrorKind::ImpossibleOperation,
+                format!(
+                    "Default formatter does not know how to format to custom format '{}'",
+                    name
+                ),
+            ));
+        }
     }
+    Ok(())
 }
 
 #[test]
diff --git a/minijinja/src/error.rs b/minijinja/src/error.rs
index 19eda52..37e1e03 100644
--- a/minijinja/src/error.rs
+++ b/minijinja/src/error.rs
@@ -83,6 +83,7 @@ pub enum ErrorKind {
     BadEscape,
     UndefinedError,
     BadSerialization,
+    WriteFailure,
 }
 
 impl ErrorKind {
@@ -101,6 +102,7 @@ impl ErrorKind {
             ErrorKind::BadEscape => "bad string escape",
             ErrorKind::UndefinedError => "variable or attribute undefined",
             ErrorKind::BadSerialization => "could not serialize to internal format",
+            ErrorKind::WriteFailure => "failed to write output",
         }
     }
 }
@@ -213,6 +215,12 @@ impl From<ErrorKind> for Error {
     }
 }
 
+impl From<fmt::Error> for Error {
+    fn from(_: fmt::Error) -> Self {
+        Error::new(ErrorKind::WriteFailure, "formatting failed")
+    }
+}
+
 impl serde::ser::Error for Error {
     fn custom<T>(msg: T) -> Self
     where
diff --git a/minijinja/src/filters.rs b/minijinja/src/filters.rs
index c58ed0e..c208388 100644
--- a/minijinja/src/filters.rs
+++ b/minijinja/src/filters.rs
@@ -58,7 +58,9 @@
 use std::collections::BTreeMap;
 use std::sync::Arc;
 
+use crate::environment::escape_formatter;
 use crate::error::Error;
+use crate::output::Output;
 use crate::value::{ArgType, FunctionArgs, FunctionResult, Value};
 use crate::vm::State;
 use crate::AutoEscape;
@@ -200,7 +202,6 @@ pub(crate) fn get_builtin_filters() -> BTreeMap<&'static str, BoxedFilter> {
 ///
 /// When a value is marked as safe, no further auto escaping will take place.
 pub fn safe(_state: &State, v: String) -> Value {
-    // TODO: this ideally understands which type of escaping is in use
     Value::from_safe_string(v)
 }
 
@@ -226,7 +227,8 @@ pub fn escape(state: &State, v: Value) -> Result<Value, Error> {
         other => other,
     };
     let mut out = String::new();
-    state.env().escape(&v, auto_escape, &mut out)?;
+    let mut formatter = Output::string(&mut out, auto_escape);
+    escape_formatter(&mut formatter, &v)?;
     Ok(Value::from_safe_string(out))
 }
 
@@ -796,9 +798,9 @@ mod builtins {
         let state = State {
             env: &env,
             ctx: crate::vm::Context::default(),
-            auto_escape: crate::AutoEscape::None,
             current_block: None,
             name: "<unknown>",
+            out: &mut Output::null(),
         };
         let bx = BoxedFilter::new(test);
         assert_eq!(
@@ -819,12 +821,13 @@ mod builtins {
         }
 
         let env = crate::Environment::new();
+        let mut out = Output::null();
         let state = State {
             env: &env,
             ctx: crate::vm::Context::default(),
-            auto_escape: crate::AutoEscape::None,
             current_block: None,
             name: "<unknown>",
+            out: &mut out,
         };
         let bx = BoxedFilter::new(add);
         assert_eq!(
diff --git a/minijinja/src/lib.rs b/minijinja/src/lib.rs
index 152ac27..7b8c3b0 100644
--- a/minijinja/src/lib.rs
+++ b/minijinja/src/lib.rs
@@ -134,6 +134,7 @@ mod error;
 mod instructions;
 mod lexer;
 mod macros;
+mod output;
 mod parser;
 mod tokens;
 mod utils;
@@ -148,8 +149,11 @@ pub mod value;
 #[cfg(feature = "source")]
 mod source;
 
-pub use self::environment::{Environment, Expression, Template};
+pub use self::environment::{
+    default_auto_escape_callback, escape_formatter, Environment, Expression, Template,
+};
 pub use self::error::{Error, ErrorKind};
+pub use self::output::Output;
 pub use self::utils::{AutoEscape, HtmlEscape};
 
 #[cfg(feature = "source")]
@@ -176,4 +180,9 @@ pub mod machinery {
     pub use crate::parser::parse;
     pub use crate::tokens::{Span, Token};
     pub use crate::vm::Vm;
+    use crate::{AutoEscape, Output};
+
+    pub fn make_string_formatter(s: &mut String, auto_escape: AutoEscape) -> Output<'_> {
+        Output::string(s, auto_escape)
+    }
 }
diff --git a/minijinja/src/output.rs b/minijinja/src/output.rs
new file mode 100644
index 0000000..be3b997
--- /dev/null
+++ b/minijinja/src/output.rs
@@ -0,0 +1,109 @@
+use std::fmt::{self, Arguments};
+
+use crate::utils::AutoEscape;
+use crate::value::Value;
+
+/// An abstraction over [`Write`](std::fmt::Write) for the rendering.
+///
+/// This is a utility type used in the engine which can be written into like one
+/// can write into an [`std::fmt::Write`] value.  It's primarily used internally
+/// in the engine but it's also passed to the custom formatter function.
+///
+/// Additionally it keeps track of the requested auto escape format of the
+/// template so that the formatter can customize it's behavior based on the auto
+/// escaping request.  The output itself however does not have functionality to
+/// escape by itself.
+pub struct Output<'a> {
+    w: &'a mut (dyn fmt::Write + 'a),
+    capture_stack: Vec<String>,
+    pub(crate) auto_escape: AutoEscape,
+}
+
+pub struct NullWriter;
+
+impl<'a> Output<'a> {
+    pub(crate) fn string(buf: &'a mut String, auto_escape: AutoEscape) -> Self {
+        Self {
+            w: buf,
+            capture_stack: Vec::new(),
+            auto_escape,
+        }
+    }
+
+    pub(crate) fn null() -> Self {
+        static mut NULL_WRITER: NullWriter = NullWriter;
+        Self {
+            // SAFETY: this is safe as the null writer is a ZST
+            w: unsafe { &mut NULL_WRITER },
+            capture_stack: Vec::new(),
+            auto_escape: AutoEscape::None,
+        }
+    }
+
+    pub(crate) fn begin_capture(&mut self) {
+        self.capture_stack.push(String::new());
+    }
+
+    pub(crate) fn end_capture(&mut self) -> Value {
+        let captured = self.capture_stack.pop().unwrap();
+        if !matches!(self.auto_escape, AutoEscape::None) {
+            Value::from_safe_string(captured)
+        } else {
+            Value::from(captured)
+        }
+    }
+
+    fn target(&mut self) -> &mut dyn fmt::Write {
+        match self.capture_stack.last_mut() {
+            Some(stream) => stream as _,
+            None => self.w,
+        }
+    }
+
+    /// Returns the current auto escape setting.
+    #[inline]
+    pub fn auto_escape(&self) -> AutoEscape {
+        self.auto_escape
+    }
+
+    /// Writes some data to the underlying buffer contained within this output.
+    #[inline]
+    pub fn write_str(&mut self, s: &str) -> fmt::Result {
+        self.target().write_str(s)
+    }
+
+    /// Writes some formatted information into this instance.
+    #[inline]
+    pub fn write_fmt(&mut self, a: Arguments<'_>) -> fmt::Result {
+        self.target().write_fmt(a)
+    }
+}
+
+impl fmt::Write for Output<'_> {
+    #[inline]
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        fmt::Write::write_str(self.target(), s)
+    }
+
+    #[inline]
+    fn write_char(&mut self, c: char) -> fmt::Result {
+        fmt::Write::write_char(self.target(), c)
+    }
+
+    #[inline]
+    fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
+        fmt::Write::write_fmt(self.target(), args)
+    }
+}
+
+impl fmt::Write for NullWriter {
+    #[inline]
+    fn write_str(&mut self, _s: &str) -> fmt::Result {
+        Ok(())
+    }
+
+    #[inline]
+    fn write_char(&mut self, _c: char) -> fmt::Result {
+        Ok(())
+    }
+}
diff --git a/minijinja/src/utils.rs b/minijinja/src/utils.rs
index 1807086..20c3810 100644
--- a/minijinja/src/utils.rs
+++ b/minijinja/src/utils.rs
@@ -41,6 +41,12 @@ pub enum AutoEscape {
     #[cfg(feature = "json")]
     #[cfg_attr(docsrs, doc(cfg(feature = "json")))]
     Json,
+    /// A custom auto escape format.
+    ///
+    /// The default formatter does not know how to deal with a custom escaping
+    /// format and would error.  The use of these requires a custom formatter.
+    /// See [`set_formatter`](crate::Environment::set_formatter).
+    Custom(&'static str),
 }
 
 /// Helper to HTML escape a string.
diff --git a/minijinja/src/vm.rs b/minijinja/src/vm.rs
index 5303511..5e407c6 100644
--- a/minijinja/src/vm.rs
+++ b/minijinja/src/vm.rs
@@ -1,6 +1,5 @@
 use std::collections::{BTreeMap, HashSet};
-use std::fmt::{self, Write};
-use std::ops::{Deref, DerefMut};
+use std::fmt;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
 
@@ -10,6 +9,7 @@ use crate::instructions::{
     Instruction, Instructions, LOOP_FLAG_RECURSIVE, LOOP_FLAG_WITH_LOOP_VAR,
 };
 use crate::key::Key;
+use crate::output::Output;
 use crate::value::{self, ops, Object, Value, ValueIterator, ValueRepr};
 use crate::AutoEscape;
 
@@ -362,27 +362,27 @@ fn process_err(mut err: Error, pc: usize, _state: &State, instructions: &Instruc
 /// only provides access to the template environment but also the context
 /// variables of the engine, the current auto escaping behavior as well as the
 /// auto escape flag.
-pub struct State<'vm, 'env> {
+pub struct State<'vm, 'env, 'out, 'buf> {
     pub(crate) env: &'env Environment<'env>,
     pub(crate) ctx: Context<'env, 'vm>,
     pub(crate) name: &'env str,
     pub(crate) current_block: Option<&'env str>,
-    pub(crate) auto_escape: AutoEscape,
+    pub(crate) out: &'out mut Output<'buf>,
 }
 
-impl<'vm, 'env> fmt::Debug for State<'vm, 'env> {
+impl<'vm, 'env, 'out, 'buf> fmt::Debug for State<'vm, 'env, 'out, 'buf> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut ds = f.debug_struct("State");
         ds.field("name", &self.name);
         ds.field("current_block", &self.current_block);
-        ds.field("auto_escape", &self.auto_escape);
+        ds.field("auto_escape", &self.out.auto_escape);
         ds.field("ctx", &self.ctx);
         ds.field("env", &self.env);
         ds.finish()
     }
 }
 
-impl<'vm, 'env> State<'vm, 'env> {
+impl<'vm, 'env, 'out, 'buf> State<'vm, 'env, 'out, 'buf> {
     /// Returns a reference to the current environment.
     pub fn env(&self) -> &Environment<'env> {
         self.env
@@ -395,7 +395,7 @@ impl<'vm, 'env> State<'vm, 'env> {
 
     /// Returns the current value of the auto escape flag.
     pub fn auto_escape(&self) -> AutoEscape {
-        self.auto_escape
+        self.out.auto_escape
     }
 
     /// Returns the name of the innermost block.
@@ -455,47 +455,6 @@ impl<'vm, 'env> State<'vm, 'env> {
     }
 }
 
-struct Output<'a> {
-    output: &'a mut String,
-    capture_stack: Vec<String>,
-}
-
-impl<'a> Output<'a> {
-    pub fn new(output: &'a mut String) -> Output<'a> {
-        Output {
-            output,
-            capture_stack: Vec::new(),
-        }
-    }
-
-    pub fn begin_capture(&mut self) {
-        self.capture_stack.push(String::new());
-    }
-
-    pub fn end_capture(&mut self, auto_escape: AutoEscape) -> Value {
-        let captured = self.capture_stack.pop().unwrap();
-        if !matches!(auto_escape, AutoEscape::None) {
-            Value::from_safe_string(captured)
-        } else {
-            Value::from(captured)
-        }
-    }
-}
-
-impl<'a> Deref for Output<'a> {
-    type Target = String;
-
-    fn deref(&self) -> &Self::Target {
-        self.capture_stack.last().unwrap_or(self.output)
-    }
-}
-
-impl<'a> DerefMut for Output<'a> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        self.capture_stack.last_mut().unwrap_or(self.output)
-    }
-}
-
 /// Helps to evaluate something.
 #[cfg_attr(feature = "internal_debug", derive(Debug))]
 pub struct Vm<'env> {
@@ -514,10 +473,8 @@ impl<'env> Vm<'env> {
         instructions: &Instructions<'env>,
         root: Value,
         blocks: &BTreeMap<&'env str, Instructions<'env>>,
-        initial_auto_escape: AutoEscape,
-        output: &mut String,
+        out: &mut Output,
     ) -> Result<Option<Value>, Error> {
-        let mut output = Output::new(output);
         let mut ctx = Context::default();
         ctx.push_frame(Frame::new(FrameBase::Value(root)));
         let mut referenced_blocks = BTreeMap::new();
@@ -527,24 +484,23 @@ impl<'env> Vm<'env> {
         let mut state = State {
             env: self.env,
             ctx,
-            auto_escape: initial_auto_escape,
             current_block: None,
             name: instructions.name(),
+            out,
         };
         value::with_value_optimization(|| {
-            self.eval_state(&mut state, instructions, referenced_blocks, &mut output)
+            self.eval_state(&mut state, instructions, referenced_blocks)
         })
     }
 
     /// This is the actual evaluation loop that works with a specific context.
     fn eval_state(
         &self,
-        state: &mut State<'_, 'env>,
+        state: &mut State<'_, 'env, '_, '_>,
         mut instructions: &Instructions<'env>,
         mut blocks: BTreeMap<&'env str, Vec<&'_ Instructions<'env>>>,
-        output: &mut Output,
     ) -> Result<Option<Value>, Error> {
-        let initial_auto_escape = state.auto_escape;
+        let initial_auto_escape = state.out.auto_escape;
         let mut stack = Stack { values: Vec::new() };
         let mut auto_escape_stack = vec![];
         let mut block_stack = vec![];
@@ -584,14 +540,7 @@ impl<'env> Vm<'env> {
 
         macro_rules! sub_eval {
             ($instructions:expr) => {{
-                self.sub_eval(
-                    state,
-                    $instructions,
-                    blocks.clone(),
-                    state.current_block,
-                    state.auto_escape,
-                    output,
-                )?
+                self.sub_eval(state, $instructions, blocks.clone(), state.current_block)?
             }};
         }
 
@@ -611,11 +560,11 @@ impl<'env> Vm<'env> {
                     layers.remove(0);
                     let instructions = layers.first().unwrap();
                     if $capture {
-                        output.begin_capture();
+                        state.out.begin_capture();
                     }
                     sub_eval!(instructions);
                     if $capture {
-                        stack.push(output.end_capture(state.auto_escape));
+                        stack.push(state.out.end_capture());
                     }
                 } else {
                     panic!("attempted to super unreferenced block");
@@ -633,7 +582,7 @@ impl<'env> Vm<'env> {
                         // to.
                         next_loop_recursion_jump = Some((pc + 1, $capture));
                         if $capture {
-                            output.begin_capture();
+                            state.out.begin_capture();
                         }
                         pc = recurse_jump_target;
                         continue;
@@ -655,10 +604,12 @@ impl<'env> Vm<'env> {
         while let Some(instr) = instructions.get(pc) {
             match instr {
                 Instruction::EmitRaw(val) => {
-                    write!(output, "{}", val).unwrap();
+                    // this only produces a format error, no need to attach
+                    // location information.
+                    write!(state.out, "{}", val)?;
                 }
                 Instruction::Emit => {
-                    try_ctx!(self.env.finalize(&stack.pop(), state.auto_escape, output));
+                    try_ctx!(self.env.format(&stack.pop(), state.out));
                 }
                 Instruction::StoreLocal(name) => {
                     state.ctx.store(name, stack.pop());
@@ -770,7 +721,7 @@ impl<'env> Vm<'env> {
                         {
                             pc = target;
                             if end_capture {
-                                stack.push(output.end_capture(state.auto_escape));
+                                stack.push(state.out.end_capture());
                             }
                             continue;
                         }
@@ -884,12 +835,12 @@ impl<'env> Vm<'env> {
                 }
                 Instruction::Include(ignore_missing) => {
                     let name = stack.pop();
-                    try_ctx!(self.perform_include(name, state, output, ignore_missing));
+                    try_ctx!(self.perform_include(name, state, *ignore_missing));
                 }
                 Instruction::PushAutoEscape => {
                     let value = stack.pop();
-                    auto_escape_stack.push(state.auto_escape);
-                    state.auto_escape = match (value.as_str(), value == Value::from(true)) {
+                    auto_escape_stack.push(state.out.auto_escape);
+                    state.out.auto_escape = match (value.as_str(), value == Value::from(true)) {
                         (Some("html"), _) => AutoEscape::Html,
                         #[cfg(feature = "json")]
                         (Some("json"), _) => AutoEscape::Json,
@@ -910,13 +861,13 @@ impl<'env> Vm<'env> {
                     };
                 }
                 Instruction::PopAutoEscape => {
-                    state.auto_escape = auto_escape_stack.pop().unwrap();
+                    state.out.auto_escape = auto_escape_stack.pop().unwrap();
                 }
                 Instruction::BeginCapture => {
-                    output.begin_capture();
+                    state.out.begin_capture();
                 }
                 Instruction::EndCapture => {
-                    stack.push(output.end_capture(state.auto_escape));
+                    stack.push(state.out.end_capture());
                 }
                 Instruction::ApplyFilter(name) => {
                     let top = stack.pop();
@@ -997,9 +948,8 @@ impl<'env> Vm<'env> {
     fn perform_include(
         &self,
         name: Value,
-        state: &mut State<'_, 'env>,
-        output: &mut Output,
-        ignore_missing: &bool,
+        state: &mut State<'_, 'env, '_, '_>,
+        ignore_missing: bool,
     ) -> Result<(), Error> {
         let choices = if let ValueRepr::Seq(ref choices) = name.0 {
             &choices[..]
@@ -1030,18 +980,13 @@ impl<'env> Vm<'env> {
             for (&name, instr) in tmpl.blocks().iter() {
                 referenced_blocks.insert(name, vec![instr]);
             }
-            self.sub_eval(
-                state,
-                instructions,
-                referenced_blocks,
-                None,
-                tmpl.initial_auto_escape(),
-                output,
-            )?;
-            templates_tried.clear();
-            break;
+            let original_escape = state.out.auto_escape;
+            state.out.auto_escape = tmpl.initial_auto_escape();
+            self.sub_eval(state, instructions, referenced_blocks, None)?;
+            state.out.auto_escape = original_escape;
+            return Ok(());
         }
-        if !templates_tried.is_empty() && !*ignore_missing {
+        if !templates_tried.is_empty() && !ignore_missing {
             Err(if templates_tried.len() == 1 {
                 Error::new(
                     ErrorKind::TemplateNotFound,
@@ -1066,23 +1011,21 @@ impl<'env> Vm<'env> {
 
     fn sub_eval(
         &self,
-        state: &mut State<'_, 'env>,
+        state: &mut State<'_, 'env, '_, '_>,
         instructions: &Instructions<'env>,
         blocks: BTreeMap<&'env str, Vec<&'_ Instructions<'env>>>,
         current_block: Option<&str>,
-        auto_escape: AutoEscape,
-        output: &mut Output,
     ) -> Result<(), Error> {
         let mut sub_context = Context::default();
         sub_context.push_frame(Frame::new(FrameBase::Context(&state.ctx)));
         let mut sub_state = State {
             env: self.env,
             ctx: sub_context,
-            auto_escape,
             current_block,
             name: instructions.name(),
+            out: state.out,
         };
-        self.eval_state(&mut sub_state, instructions, blocks, output)?;
+        self.eval_state(&mut sub_state, instructions, blocks)?;
         Ok(())
     }
 }


Test patch:
diff --git a/minijinja/src/tests.rs b/minijinja/src/tests.rs
index cd8fcee..13fbcfd 100644
--- a/minijinja/src/tests.rs
+++ b/minijinja/src/tests.rs
@@ -277,9 +277,9 @@ mod builtins {
         let state = State {
             env: &env,
             ctx: crate::vm::Context::default(),
-            auto_escape: crate::AutoEscape::None,
             current_block: None,
             name: "<unknown>",
+            out: &mut crate::Output::null(),
         };
         let bx = BoxedTest::new(test);
         assert!(bx
diff --git a/minijinja/tests/test_vm.rs b/minijinja/tests/test_vm.rs
index 821df7a..71a2253 100644
--- a/minijinja/tests/test_vm.rs
+++ b/minijinja/tests/test_vm.rs
@@ -1,7 +1,7 @@
 #![cfg(feature = "unstable_machinery")]
 use std::collections::BTreeMap;
 
-use minijinja::machinery::{Compiler, Instruction, Instructions, Vm};
+use minijinja::machinery::{make_string_formatter, Compiler, Instruction, Instructions, Vm};
 use minijinja::value::Value;
 use minijinja::{AutoEscape, Environment, Error};
 
@@ -16,7 +16,8 @@ pub fn simple_eval<S: serde::Serialize>(
     let empty_blocks = BTreeMap::new();
     let vm = Vm::new(&env);
     let root = Value::from_serializable(&ctx);
-    vm.eval(instructions, root, &empty_blocks, AutoEscape::None, output)
+    let mut formatter = make_string_formatter(output, AutoEscape::None);
+    vm.eval(instructions, root, &empty_blocks, &mut formatter)
 }
 
 #[test]