Add support for loop.index and loop.index0 variables

1 files changed, 30 insertions, 0 deletions

diff --git a/askama/src/generator.rs b/askama/src/generator.rs
index 895ec25..72c5ecd 100644
--- a/askama/src/generator.rs
+++ b/askama/src/generator.rs
@@ -35,6 +35,8 @@ struct Generator<'a> {
     locals: HashSet<String>,
     next_ws: Option<&'a str>,
     skip_ws: bool,
+    loop_vars: bool,
+    loop_index: u8,
 }
 
 impl<'a> Generator<'a> {
@@ -47,6 +49,8 @@ impl<'a> Generator<'a> {
             locals: HashSet::new(),
             next_ws: None,
             skip_ws: false,
+            loop_vars: false,
+            loop_index: 0,
         }
     }
 
@@ -120,6 +124,19 @@ impl<'a> Generator<'a> {
     }
 
     fn visit_attr(&mut self, obj: &Expr, attr: &str) {
+        if let Expr::Var(name) = *obj {
+            if name == "loop" {
+                self.write("_loop_indexes[_loop_cur]");
+                if attr == "index" {
+                    return;
+                } else if attr == "index0" {
+                    self.write(" - 1");
+                    return;
+                } else {
+                    panic!("unknown loop variable");
+                }
+            }
+        }
         self.visit_expr(obj);
         self.write(&format!(".{}", attr));
     }
@@ -239,6 +256,16 @@ impl<'a> Generator<'a> {
                   body: &'a [Node], ws2: &WS) {
 
         self.handle_ws(ws1);
+        if !self.loop_vars {
+            self.writeln("let mut _loop_indexes = Vec::new();");
+            self.writeln("let mut _loop_cur = 0;");
+            self.loop_vars = true;
+        }
+
+        self.writeln("_loop_indexes.push(0);");
+        let cur_index = self.loop_index;
+        self.loop_index += 1;
+        self.writeln(&format!("_loop_cur = {};", cur_index));
         self.write("for ");
         let targets = self.visit_target(var);
         for name in &targets {
@@ -249,9 +276,12 @@ impl<'a> Generator<'a> {
         self.visit_expr(iter);
         self.writeln(" {");
 
+        self.writeln("_loop_indexes[_loop_cur] += 1;");
         self.handle(body);
         self.handle_ws(ws2);
         self.writeln("}");
+        self.loop_index -= 1;
+        self.writeln("_loop_indexes.pop();");
         for name in &targets {
             self.locals.remove(name);
         }