Merge pull request #661 from cmrschwarz/avoid_context_clone_v2

sunng87 · web-flow · commit 080b5e0b23e2 · 2024-07-15T22:00:09.000+08:00
Avoid cloning the Render Context
diff --git a/src/partial.rs b/src/partial.rs
@@ -1,4 +1,3 @@
-use std::borrow::Cow;
 use std::collections::HashMap;
 
 use serde_json::value::Value as Json;
@@ -15,22 +14,26 @@ use crate::RenderErrorReason;
 
 pub(crate) const PARTIAL_BLOCK: &str = "@partial-block";
 
-fn find_partial<'reg: 'rc, 'rc: 'a, 'a>(
-    rc: &'a RenderContext<'reg, 'rc>,
+fn find_partial<'reg: 'rc, 'rc>(
+    rc: &RenderContext<'reg, 'rc>,
     r: &'reg Registry<'reg>,
     d: &Decorator<'rc>,
     name: &str,
-) -> Result<Option<Cow<'a, Template>>, RenderError> {
+) -> Result<Option<&'rc Template>, RenderError> {
     if let Some(partial) = rc.get_partial(name) {
-        return Ok(Some(Cow::Borrowed(partial)));
+        return Ok(Some(partial));
     }
 
-    if let Some(tpl) = r.get_or_load_template_optional(name) {
-        return tpl.map(Option::Some);
+    if let Some(t) = rc.get_dev_mode_template(name) {
+        return Ok(Some(t));
+    }
+
+    if let Some(t) = r.get_template(name) {
+        return Ok(Some(t));
     }
 
     if let Some(tpl) = d.template() {
-        return Ok(Some(Cow::Borrowed(tpl)));
+        return Ok(Some(tpl));
     }
 
     Ok(None)
@@ -49,114 +52,103 @@ pub fn expand_partial<'reg: 'rc, 'rc>(
     }
 
     let tname = d.name();
+
+    let current_template_before = rc.get_current_template_name();
+    let indent_before = rc.get_indent_string().cloned();
+
     if rc.is_current_template(tname) {
         return Err(RenderErrorReason::CannotIncludeSelf.into());
     }
 
     let partial = find_partial(rc, r, d, tname)?;
 
-    if let Some(t) = partial {
-        // clone to avoid lifetime issue
-        // FIXME refactor this to avoid
-        let mut local_rc = rc.clone();
-
-        // if tname == PARTIAL_BLOCK
-        let is_partial_block = tname == PARTIAL_BLOCK;
-
-        // add partial block depth there are consecutive partial
-        // blocks in the stack.
-        if is_partial_block {
-            local_rc.inc_partial_block_depth();
-        } else {
-            // depth cannot be lower than 0, which is guaranted in the
-            // `dec_partial_block_depth` method
-            local_rc.dec_partial_block_depth();
-        }
+    let Some(partial) = partial else {
+        return Err(RenderErrorReason::PartialNotFound(tname.to_owned()).into());
+    };
+
+    let is_partial_block = tname == PARTIAL_BLOCK;
+
+    // add partial block depth there are consecutive partial
+    // blocks in the stack.
+    if is_partial_block {
+        rc.inc_partial_block_depth();
+    } else {
+        // depth cannot be lower than 0, which is guaranted in the
+        // `dec_partial_block_depth` method
+        rc.dec_partial_block_depth();
+    }
+
+    let mut block_created = false;
+
+    // create context if param given
+    if let Some(base_path) = d.param(0).and_then(|p| p.context_path()) {
+        // path given, update base_path
+        let mut block_inner = BlockContext::new();
+        *block_inner.base_path_mut() = base_path.to_vec();
 
-        let mut block_created = false;
+        block_created = true;
+        rc.push_block(block_inner);
+    }
 
-        // create context if param given
-        if let Some(base_path) = d.param(0).and_then(|p| p.context_path()) {
-            // path given, update base_path
-            let mut block_inner = BlockContext::new();
-            base_path.clone_into(block_inner.base_path_mut());
+    if !d.hash().is_empty() {
+        // hash given, update base_value
+        let hash_ctx = d
+            .hash()
+            .iter()
+            .map(|(k, v)| (*k, v.value()))
+            .collect::<HashMap<&str, &Json>>();
+
+        // create block if we didn't (no param provided for partial expression)
+        if !block_created {
+            let block_inner = if let Some(block) = rc.block() {
+                // reuse current block information, including base_path and
+                // base_value if any
+                block.clone()
+            } else {
+                BlockContext::new()
+            };
 
-            // because block is moved here, we need another bool variable to track
-            // its status for later cleanup
             block_created = true;
-            // clear blocks to prevent block params from parent
-            // template to be leaked into partials
-            // see `test_partial_context_issue_495` for the case.
-            local_rc.clear_blocks();
-            local_rc.push_block(block_inner);
+            rc.push_block(block_inner);
         }
 
-        if !d.hash().is_empty() {
-            // hash given, update base_value
-            let hash_ctx = d
-                .hash()
-                .iter()
-                .map(|(k, v)| (*k, v.value()))
-                .collect::<HashMap<&str, &Json>>();
-
-            // create block if we didn't (no param provided for partial expression)
-            if !block_created {
-                let block_inner = if let Some(block) = local_rc.block() {
-                    // reuse current block information, including base_path and
-                    // base_value if any
-                    block.clone()
-                } else {
-                    BlockContext::new()
-                };
-
-                local_rc.clear_blocks();
-                local_rc.push_block(block_inner);
-            }
-
-            // evaluate context within current block, this includes block
-            // context provided by partial expression parameter
-            let merged_context = merge_json(
-                local_rc.evaluate2(ctx, &Path::current())?.as_json(),
-                &hash_ctx,
-            );
-
-            // update the base value, there must be a block for this so it's
-            // also safe to unwrap.
-            if let Some(block) = local_rc.block_mut() {
-                block.set_base_value(merged_context);
-            }
-        }
+        // evaluate context within current block, this includes block
+        // context provided by partial expression parameter
+        let merged_context = merge_json(rc.evaluate2(ctx, &Path::current())?.as_json(), &hash_ctx);
 
-        // @partial-block
-        if let Some(pb) = d.template() {
-            local_rc.push_partial_block(pb);
+        // update the base value, there must be a block for this so it's
+        // also safe to unwrap.
+        if let Some(block) = rc.block_mut() {
+            block.set_base_value(merged_context);
         }
+    }
 
-        // indent
-        local_rc.set_indent_string(d.indent().cloned());
-
-        let result = t.render(r, ctx, &mut local_rc, out);
+    // @partial-block
+    if let Some(pb) = d.template() {
+        rc.push_partial_block(pb);
+    }
 
-        // cleanup
+    // indent
+    rc.set_indent_string(d.indent().cloned());
 
-        let trailing_newline = local_rc.get_trailine_newline();
+    let result = partial.render(r, ctx, rc, out);
 
-        if block_created {
-            local_rc.pop_block();
-        }
+    // cleanup
+    let trailing_newline = rc.get_trailine_newline();
 
-        if d.template().is_some() {
-            local_rc.pop_partial_block();
-        }
+    if d.template().is_some() {
+        rc.pop_partial_block();
+    }
 
-        drop(local_rc);
+    if block_created {
+        rc.pop_block();
+    }
 
-        rc.set_trailing_newline(trailing_newline);
+    rc.set_trailing_newline(trailing_newline);
+    rc.set_current_template_name(current_template_before);
+    rc.set_indent_string(indent_before);
 
-        result
-    } else {
-        Err(RenderErrorReason::PartialNotFound(tname.to_owned()).into())
-    }
+    result
 }
 
 #[cfg(test)]
diff --git a/src/registry.rs b/src/registry.rs
@@ -1,5 +1,5 @@
 use std::borrow::Cow;
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use std::convert::AsRef;
 use std::fmt::{self, Debug, Formatter};
 use std::io::{Error as IoError, Write};
@@ -671,6 +671,52 @@ impl<'reg> Registry<'reg> {
         self.template_sources.clear();
     }
 
+    fn gather_dev_mode_templates(
+        &'reg self,
+        prebound: Option<(&str, Cow<'reg, Template>)>,
+    ) -> Result<BTreeMap<String, Cow<'reg, Template>>, RenderError> {
+        let prebound_name = prebound.as_ref().map(|(name, _)| *name);
+        let mut res = BTreeMap::new();
+        for name in self.template_sources.keys() {
+            if Some(&**name) == prebound_name {
+                continue;
+            }
+            res.insert(name.clone(), self.get_or_load_template(name)?);
+        }
+        if let Some((name, prebound)) = prebound {
+            res.insert(name.to_owned(), prebound);
+        }
+        Ok(res)
+    }
+
+    fn render_resolved_template_to_output(
+        &self,
+        name: Option<&str>,
+        template: Cow<'_, Template>,
+        ctx: &Context,
+        output: &mut impl Output,
+    ) -> Result<(), RenderError> {
+        if !self.dev_mode {
+            let mut render_context = RenderContext::new(template.name.as_ref());
+            return template.render(self, ctx, &mut render_context, output);
+        }
+
+        let dev_mode_templates;
+        let template = if let Some(name) = name {
+            dev_mode_templates = self.gather_dev_mode_templates(Some((name, template)))?;
+            &dev_mode_templates[name]
+        } else {
+            dev_mode_templates = self.gather_dev_mode_templates(None)?;
+            &template
+        };
+
+        let mut render_context = RenderContext::new(template.name.as_ref());
+
+        render_context.set_dev_mode_templates(Some(&dev_mode_templates));
+
+        template.render(self, ctx, &mut render_context, output)
+    }
+
     #[inline]
     fn render_to_output<O>(
         &self,
@@ -681,10 +727,12 @@ impl<'reg> Registry<'reg> {
     where
         O: Output,
     {
-        self.get_or_load_template(name).and_then(|t| {
-            let mut render_context = RenderContext::new(t.name.as_ref());
-            t.render(self, ctx, &mut render_context, output)
-        })
+        self.render_resolved_template_to_output(
+            Some(name),
+            self.get_or_load_template(name)?,
+            ctx,
+            output,
+        )
     }
 
     /// Render a registered template with some data into a string
@@ -762,10 +810,7 @@ impl<'reg> Registry<'reg> {
         .map_err(RenderError::from)?;
 
         let mut out = StringOutput::new();
-        {
-            let mut render_context = RenderContext::new(None);
-            tpl.render(self, ctx, &mut render_context, &mut out)?;
-        }
+        self.render_resolved_template_to_output(None, Cow::Owned(tpl), ctx, &mut out)?;
 
         out.into_string().map_err(RenderError::from)
     }
@@ -789,9 +834,9 @@ impl<'reg> Registry<'reg> {
             },
         )
         .map_err(RenderError::from)?;
-        let mut render_context = RenderContext::new(None);
         let mut out = WriteOutput::new(writer);
-        tpl.render(self, ctx, &mut render_context, &mut out)
+
+        self.render_resolved_template_to_output(None, Cow::Owned(tpl), ctx, &mut out)
     }
 
     /// Render a template string using current registry without registering it
diff --git a/src/render.rs b/src/render.rs
