Fix loop.last

2 files changed, 37 insertions, 36 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index 6c02a26..288ddae 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -467,15 +467,15 @@ impl<'a> Generator<'a> {
         let expr_code = self.visit_expr_root(iter);
 
         self.write_buf_writable(buf);
-        buf.write("for (_loop_index, _loop_last, ");
+        buf.write("for (");
         self.visit_target(buf, var);
         match iter {
             Expr::Range(_, _, _) => buf.writeln(&format!(
-                ") in ::askama::helpers::enumerate({}) {{",
+                ", _loop_item) in ::askama::helpers::TemplateLoop::new({}) {{",
                 expr_code
             )),
             _ => buf.writeln(&format!(
-                ") in ::askama::helpers::enumerate((&{}).into_iter()) {{",
+                ", _loop_item) in ::askama::helpers::TemplateLoop::new((&{}).into_iter()) {{",
                 expr_code
             )),
         };
@@ -913,16 +913,16 @@ impl<'a> Generator<'a> {
         if let Expr::Var(name) = *obj {
             if name == "loop" {
                 if attr == "index" {
-                    buf.write("(_loop_index + 1)");
+                    buf.write("(_loop_item.index + 1)");
                     return DisplayWrap::Unwrapped;
                 } else if attr == "index0" {
-                    buf.write("_loop_index");
+                    buf.write("_loop_item.index");
                     return DisplayWrap::Unwrapped;
                 } else if attr == "first" {
-                    buf.write("(_loop_index == 0)");
+                    buf.write("_loop_item.first");
                     return DisplayWrap::Unwrapped;
                 } else if attr == "last" {
-                    buf.write("_loop_last");
+                    buf.write("_loop_item.last");
                     return DisplayWrap::Unwrapped;
                 } else {
                     panic!("unknown loop variable");
diff --git a/askama_shared/src/helpers/mod.rs b/askama_shared/src/helpers/mod.rs
index 1a3fd7b..87316fd 100644
--- a/askama_shared/src/helpers/mod.rs
+++ b/askama_shared/src/helpers/mod.rs
@@ -1,48 +1,49 @@
+use std::iter::Enumerate;
 use std::iter::Peekable;
 
-pub struct Enumerate<I>
+pub struct TemplateLoop<I>
 where
     I: Iterator,
 {
-    iter: Peekable<I>,
-    count: usize,
+    iter: Peekable<Enumerate<I>>,
 }
 
-impl<I> Iterator for Enumerate<I>
+impl<I> TemplateLoop<I>
 where
     I: Iterator,
 {
-    type Item = (usize, bool, <I as Iterator>::Item);
-
     #[inline]
-    fn next(&mut self) -> Option<(usize, bool, <I as Iterator>::Item)> {
-        self.iter.next().map(|a| {
-            let last = self.iter.peek().is_none();
-            let ret = (self.count, last, a);
-            // Possible undefined overflow.
-            self.count += 1;
-            ret
-        })
-    }
-
-    #[inline]
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        self.iter.size_hint()
-    }
-
-    #[inline]
-    fn count(self) -> usize {
-        self.iter.count()
+    pub fn new(iter: I) -> Self {
+        TemplateLoop {
+            iter: iter.enumerate().peekable(),
+        }
     }
 }
 
-#[inline]
-pub fn enumerate<I>(iter: I) -> Enumerate<I>
+impl<I> Iterator for TemplateLoop<I>
 where
     I: Iterator,
 {
-    Enumerate {
-        iter: iter.peekable(),
-        count: 0,
+    type Item = (<I as Iterator>::Item, LoopItem);
+
+    #[inline]
+    fn next(&mut self) -> Option<(<I as Iterator>::Item, LoopItem)> {
+        self.iter.next().map(|(index, item)| {
+            (
+                item,
+                LoopItem {
+                    index,
+                    first: index == 0,
+                    last: self.iter.peek().is_none(),
+                },
+            )
+        })
     }
 }
+
+#[derive(Copy, Clone)]
+pub struct LoopItem {
+    pub index: usize,
+    pub first: bool,
+    pub last: bool,
+}