Extend for loops with custom index support (#119)

* Extend for loops with custom index support

* Swap index and element to match Swift's enumerated()

Sources/LeafKit/LeafSerialize/LeafSerializer.swift

@@ -141,7 +141,7 @@ internal struct LeafSerializer {
 
             innerContext["isFirst"] = .bool(idx == array.startIndex)
             innerContext["isLast"] = .bool(idx == array.index(before: array.endIndex))
-            innerContext["index"] = .int(idx)
+            innerContext[loop.index] = .int(idx)
             innerContext[loop.item] = item
 
             var serializer = LeafSerializer(

Sources/LeafKit/LeafSyntax/LeafSyntax.swift

@@ -533,6 +533,8 @@ extension Syntax {
         public let item: String
         /// the key to use to access the array
         public let array: String
+        /// the key to use when accessing the current index
+        public let index: String
 
         /// the body of the looop
         public internal(set) var body: [Syntax]
@@ -542,30 +544,50 @@ extension Syntax {
 
         /// initialize a new loop
         public init(_ params: [ParameterDeclaration], body: [Syntax]) throws {
-            guard
-                params.count == 1,
-                case .expression(let list) = params[0],
-                list.count == 3,
-                case .parameter(let left) = list[0],
-                case .variable(let item) = left,
-                case .parameter(let `in`) = list[1],
-                case .keyword(let k) = `in`,
-                k == .in,
-                case .parameter(let right) = list[2],
-                case .variable(let array) = right
-                else { throw "for loops expect single expression, 'name in names'" }
-            self.item = item
-            self.array = array
+            if params.count == 1 {
+                guard
+                    case .expression(let list) = params[0],
+                    list.count == 3,
+                    case .parameter(let left) = list[0],
+                    case .variable(let item) = left,
+                    case .parameter(let `in`) = list[1],
+                    case .keyword(let k) = `in`,
+                    k == .in,
+                    case .parameter(let right) = list[2],
+                    case .variable(let array) = right
+                    else { throw "for loops expect one of the following expressions: 'name in names' or 'name, nameIndex in names'" }
+                self.item = item
+                self.array = array
+                self.index = "index"
+            } else {
+                guard
+                    params.count == 2,
+                    case .parameter(.variable(let index)) = params[0],
+                    case .expression(let list) = params[1],
+                    list.count == 3,
+                    case .parameter(let left) = list[0],
+                    case .variable(let item) = left,
+                    case .parameter(let `in`) = list[1],
+                    case .keyword(let k) = `in`,
+                    k == .in,
+                    case .parameter(let right) = list[2],
+                    case .variable(let array) = right
+                else { throw "for loops expect one of the following expressions: 'name in names' or 'name, nameIndex in names'" }
+                self.item = item
+                self.array = array
+                self.index = index
+            }
 
             guard !body.isEmpty else { throw "for loops require a body" }
             self.body = body
             self.externalsSet = body.externals()
             self.importSet = body.imports()
         }
 
-        internal init(item: String, array: String, body: [Syntax]) {
+        internal init(item: String, array: String, index: String, body: [Syntax]) {
             self.item = item
             self.array = array
+            self.index = index
             self.body = body
             self.externalsSet = body.externals()
             self.importSet = body.imports()
@@ -581,12 +603,12 @@ extension Syntax {
 
         func inlineRefs(_ externals: [String: LeafAST], _ imports: [String : Syntax.Export]) -> [Syntax] {
             guard !externalsSet.isEmpty || !importSet.isEmpty else { return [.loop(self)] }
-            return [.loop(.init(item: item, array: array, body: body.inlineRefs(externals, imports)))]
+            return [.loop(.init(item: item, array: array, index: index, body: body.inlineRefs(externals, imports)))]
         }
 
         func print(depth: Int) -> String {
             var print = indent(depth)
-            print += "for(" + item + " in " + array + "):\n"
+            print += "for(" + (index == "index" ? "" : "\(index), ") + item + " in " + array + "):\n"
             print += body.map { $0.print(depth: depth + 1) } .joined(separator: "\n")
             return print
         }

Tests/LeafKitTests/LeafTests.swift

@@ -328,6 +328,34 @@ final class LeafTests: XCTestCase {
         try XCTAssertEqual(render(template, ["arrays": data]), expected)
     }
 
+    func testNestedLoopCustomIndices() throws {
+        let template = """
+        #for(i, array in arrays):#for(j, element in array):
+        (#(i), #(j)): #(element)#endfor#endfor
+        """
+
+        let expected = """
+
+        (0, 0): zero
+        (0, 1): one
+        (0, 2): two
+        (1, 0): a
+        (1, 1): b
+        (1, 2): c
+        (2, 0): red fish
+        (2, 1): blue fish
+        (2, 2): green fish
+        """
+
+        let data = LeafData.array([
+            LeafData.array(["zero", "one", "two"]),
+            LeafData.array(["a", "b", "c"]),
+            LeafData.array(["red fish", "blue fish", "green fish"])
+        ])
+
+        try XCTAssertEqual(render(template, ["arrays": data]), expected)
+    }
+
     // It would be nice if a pre-render phase could catch things like calling
     // tags that would normally ALWAYS throw in serializing (eg, calling index
     // when not in a loop) so that warnings can be provided and AST can be minimized.

Tests/LeafKitTests/TestHelpers.swift

@@ -196,6 +196,27 @@ final class PrintTests: XCTestCase {
         XCTAssertEqual(output, expectation)
     }
 
+    func testLoopCustomIndex() throws {
+        let template = """
+        #for(i, name in names):
+            #(i): hello, #(name).
+        #endfor
+        """
+        let expectation = """
+        for(i, name in names):
+          raw("\\n    ")
+          expression[variable(i)]
+          raw(": hello, ")
+          expression[variable(name)]
+          raw(".\\n")
+        """
+
+        let v = try parse(template).first!
+        guard case .loop(let test) = v else { throw "nope" }
+        let output = test.print(depth: 0)
+        XCTAssertEqual(output, expectation)
+    }
+
     func testConditional() throws {
         let template = """
         #if(foo):