Improved implicit borrowing and changed to resolve Askama variables at compile time

1 files changed, 127 insertions, 49 deletions

diff --git a/askama_shared/src/generator.rs b/askama_shared/src/generator.rs
index 0053b38..ab964e9 100644
--- a/askama_shared/src/generator.rs
+++ b/askama_shared/src/generator.rs
@@ -10,7 +10,7 @@ use proc_macro2::Span;
 
 use quote::{quote, ToTokens};
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::path::PathBuf;
 use std::{cmp, hash, mem, str};
 
@@ -20,7 +20,7 @@ pub fn generate<S: std::hash::BuildHasher>(
     heritage: &Option<Heritage>,
     integrations: Integrations,
 ) -> Result<String, CompileError> {
-    Generator::new(input, contexts, heritage, integrations, SetChain::new())
+    Generator::new(input, contexts, heritage, integrations, MapChain::new())
         .build(&contexts[&input.path])
 }
 
@@ -34,7 +34,7 @@ struct Generator<'a, S: std::hash::BuildHasher> {
     // What integrations need to be generated
     integrations: Integrations,
     // Variables accessible directly from the current scope (not redirected to context)
-    locals: SetChain<'a, &'a str>,
+    locals: MapChain<'a, &'a str, LocalMeta>,
     // Suffix whitespace from the previous literal. Will be flushed to the
     // output buffer unless suppressed by whitespace suppression on the next
     // non-literal.
@@ -56,7 +56,7 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
         contexts: &'n HashMap<&'n PathBuf, Context<'n>, S>,
         heritage: &'n Option<Heritage>,
         integrations: Integrations,
-        locals: SetChain<'n, &'n str>,
+        locals: MapChain<'n, &'n str, LocalMeta>,
     ) -> Generator<'n, S> {
         Generator {
             input,
@@ -73,7 +73,7 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
     }
 
     fn child(&mut self) -> Generator<'_, S> {
-        let locals = SetChain::with_parent(&self.locals);
+        let locals = MapChain::with_parent(&self.locals);
         Self::new(
             self.input,
             self.contexts,
@@ -597,7 +597,7 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
                         buf.write("(");
                         for (i, param) in params.iter().enumerate() {
                             if let MatchParameter::Name(p) = *param {
-                                self.locals.insert(p);
+                                self.locals.insert_with_default(p);
                             }
                             if i > 0 {
                                 buf.write(", ");
@@ -611,9 +611,9 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
                     buf.write("{");
                     for (i, param) in params.iter().enumerate() {
                         if let Some(MatchParameter::Name(p)) = param.1 {
-                            self.locals.insert(p);
+                            self.locals.insert_with_default(p);
                         } else {
-                            self.locals.insert(param.0);
+                            self.locals.insert_with_default(param.0);
                         }
 
                         if i > 0 {
@@ -744,21 +744,49 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
         let mut names = Buffer::new(0);
         let mut values = Buffer::new(0);
         for (i, arg) in def.args.iter().enumerate() {
-            if i > 0 {
-                names.write(", ");
-                values.write(", ");
+            let expr = args.get(i).ok_or_else(|| {
+                CompileError::String(format!("macro '{}' takes more than {} arguments", name, i))
+            })?;
+
+            match expr {
+                // If `expr` is already a form of variable then
+                // don't reintroduce a new variable. This is
+                // to avoid moving non-copyable values.
+                Expr::Var(name) => {
+                    let var = self.locals.resolve_or_self(name);
+                    self.locals.insert(arg, LocalMeta::with_ref(var));
+                }
+                Expr::Attr(obj, attr) => {
+                    let mut attr_buf = Buffer::new(0);
+                    self.visit_attr(&mut attr_buf, obj, attr)?;
+
+                    let var = self.locals.resolve(&attr_buf.buf).unwrap_or(attr_buf.buf);
+                    self.locals.insert(arg, LocalMeta::with_ref(var));
+                }
+                // Everything else still needs to become variables,
+                // to avoid having the same logic be executed
+                // multiple times, e.g. in the case of macro
+                // parameters being used multiple times.
+                _ => {
+                    if i > 0 {
+                        names.write(", ");
+                        values.write(", ");
+                    }
+                    names.write(arg);
+
+                    values.write("(");
+                    values.write(&self.visit_expr_root(expr)?);
+                    values.write(")");
+                    self.locals.insert_with_default(arg);
+                }
             }
-            names.write(arg);
+        }
 
-            values.write("&(");
-            values.write(&self.visit_expr_root(args.get(i).ok_or_else(|| {
-                CompileError::String(format!("macro '{}' takes more than {} arguments", name, i))
-            })?)?);
-            values.write(")");
-            self.locals.insert(arg);
+        debug_assert_eq!(names.buf.is_empty(), values.buf.is_empty());
+        if !names.buf.is_empty() {
+            buf.writeln(&format!("let ({}) = ({});", names.buf, values.buf))?;
         }
 
-        buf.writeln(&format!("let ({}) = ({});", names.buf, values.buf))?;
         let mut size_hint = self.handle(own_ctx, &def.nodes, buf, AstLevel::Nested)?;
 
         self.flush_ws(def.ws2);
@@ -819,13 +847,13 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
         buf.write("let ");
         match *var {
             Target::Name(name) => {
-                self.locals.insert(name);
+                self.locals.insert_with_default(name);
                 buf.write(name);
             }
             Target::Tuple(ref targets) => {
                 buf.write("(");
                 for name in targets {
-                    self.locals.insert(name);
+                    self.locals.insert_with_default(name);
                     buf.write(name);
                     buf.write(",");
                 }
@@ -848,16 +876,16 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
 
         match *var {
             Target::Name(name) => {
-                if !self.locals.contains(name) {
+                if !self.locals.contains(&name) {
                     buf.write("let ");
-                    self.locals.insert(name);
+                    self.locals.insert_with_default(name);
                 }
                 buf.write(name);
             }
             Target::Tuple(ref targets) => {
                 buf.write("let (");
                 for name in targets {
-                    self.locals.insert(name);
+                    self.locals.insert_with_default(name);
                     buf.write(name);
                     buf.write(",");
                 }
@@ -1396,12 +1424,12 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
     }
 
     fn visit_var(&mut self, buf: &mut Buffer, s: &str) -> DisplayWrap {
-        if self.locals.contains(s) || s == "self" {
-            buf.write(s);
-        } else {
-            buf.write("self.");
+        if s == "self" {
             buf.write(s);
+            return DisplayWrap::Unwrapped;
         }
+
+        buf.write(&self.locals.resolve_or_self(&s));
         DisplayWrap::Unwrapped
     }
 
@@ -1412,7 +1440,7 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
         args: &[Expr],
     ) -> Result<DisplayWrap, CompileError> {
         buf.write("(");
-        if self.locals.contains(s) || s == "self" {
+        if self.locals.contains(&s) || s == "self" {
             buf.write(s);
         } else {
             buf.write("self.");
@@ -1447,13 +1475,13 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
     fn visit_target(&mut self, buf: &mut Buffer, target: &'a Target) {
         match *target {
             Target::Name(name) => {
-                self.locals.insert(name);
+                self.locals.insert_with_default(name);
                 buf.write(name);
             }
             Target::Tuple(ref targets) => {
                 buf.write("(");
                 for name in targets {
-                    self.locals.insert(name);
+                    self.locals.insert_with_default(name);
                     buf.write(name);
                     buf.write(",");
                 }
@@ -1548,51 +1576,87 @@ impl Buffer {
     }
 }
 
+#[derive(Default)]
+struct LocalMeta {
+    refs: Option<String>,
+}
+
+impl LocalMeta {
+    fn with_ref(refs: String) -> Self {
+        Self { refs: Some(refs) }
+    }
+}
+
+// type SetChain<'a, T> = MapChain<'a, T, ()>;
+
 #[derive(Debug)]
-struct SetChain<'a, T: 'a>
+struct MapChain<'a, K: 'a, V: 'a>
 where
-    T: cmp::Eq + hash::Hash,
+    K: cmp::Eq + hash::Hash,
 {
-    parent: Option<&'a SetChain<'a, T>>,
-    scopes: Vec<HashSet<T>>,
+    parent: Option<&'a MapChain<'a, K, V>>,
+    scopes: Vec<HashMap<K, V>>,
 }
 
-impl<'a, T: 'a> SetChain<'a, T>
+impl<'a, K: 'a, V: 'a> MapChain<'a, K, V>
 where
-    T: cmp::Eq + hash::Hash,
+    K: cmp::Eq + hash::Hash,
 {
-    fn new() -> SetChain<'a, T> {
-        SetChain {
+    fn new() -> MapChain<'a, K, V> {
+        MapChain {
             parent: None,
-            scopes: vec![HashSet::new()],
+            scopes: vec![HashMap::new()],
         }
     }
 
-    fn with_parent<'p>(parent: &'p SetChain<T>) -> SetChain<'p, T> {
-        SetChain {
+    fn with_parent<'p>(parent: &'p MapChain<K, V>) -> MapChain<'p, K, V> {
+        MapChain {
             parent: Some(parent),
-            scopes: vec![HashSet::new()],
+            scopes: vec![HashMap::new()],
         }
     }
 
-    fn contains(&self, val: T) -> bool {
-        self.scopes.iter().rev().any(|set| set.contains(&val))
+    fn contains(&self, key: &K) -> bool {
+        self.scopes.iter().rev().any(|set| set.contains_key(&key))
             || match self.parent {
-                Some(set) => set.contains(val),
+                Some(set) => set.contains(key),
                 None => false,
             }
     }
 
+    /// Iterates the scopes in reverse and returns `Some(LocalMeta)`
+    /// from the first scope where `key` exists.
+    fn get(&self, key: &K) -> Option<&V> {
+        let scopes = self.scopes.iter().rev();
+        scopes
+            .filter_map(|set| set.get(&key))
+            .next()
+            .or_else(|| self.parent.and_then(|set| set.get(key)))
+    }
+
     fn is_current_empty(&self) -> bool {
         self.scopes.last().unwrap().is_empty()
     }
 
-    fn insert(&mut self, val: T) {
-        self.scopes.last_mut().unwrap().insert(val);
+    fn insert(&mut self, key: K, val: V) {
+        self.scopes.last_mut().unwrap().insert(key, val);
+
+        // Note that if `insert` returns `Some` then it implies
+        // an identifier is reused. For e.g. `{% macro f(a, a) %}`
+        // and `{% let (a, a) = ... %}` then this results in a
+        // generated template, which when compiled fails with the
+        // compile error "identifier `a` used more than once".
+    }
+
+    fn insert_with_default(&mut self, key: K)
+    where
+        V: Default,
+    {
+        self.insert(key, V::default());
     }
 
     fn push(&mut self) {
-        self.scopes.push(HashSet::new());
+        self.scopes.push(HashMap::new());
     }
 
     fn pop(&mut self) {
@@ -1601,6 +1665,20 @@ where
     }
 }
 
+impl MapChain<'_, &str, LocalMeta> {
+    fn resolve(&self, name: &str) -> Option<String> {
+        self.get(&name).map(|meta| match &meta.refs {
+            Some(expr) => expr.clone(),
+            None => name.to_string(),
+        })
+    }
+
+    fn resolve_or_self(&self, name: &str) -> String {
+        self.resolve(name)
+            .unwrap_or_else(|| format!("self.{}", name))
+    }
+}
+
 fn median(sizes: &mut [usize]) -> usize {
     sizes.sort_unstable();
     if sizes.len() % 2 == 1 {