Fix whitespace stripping after control blocks

src/template/lib.zig

@@ -8,20 +8,26 @@ pub fn main() !void {
             &.{ "5", "4", "3", "2", "1" },
             &.{ "5", "4", "3", "2", "1" },
         },
+        .qux = true,
+        .quxx = false,
     });
 }
 
 pub fn execute(writer: anytype, comptime template: []const u8, args: anytype) !void {
-    @setEvalBranchQuota(@intCast(u32, template.len * 6));
-    const tmpl = comptime parseTemplate(TokenIter{ .text = template }, .root);
+    @setEvalBranchQuota(@intCast(u32, template.len * 8));
+
+    const tokens = comptime parseTemplateTokens(ControlTokenIter{ .text = template });
+    const tmpl = comptime parseTemplate(tokens, 0, .root);
     try executeTemplate(writer, tmpl.item, args, .{});
 }
 
 fn executeTemplate(writer: anytype, comptime items: []const TemplateItem, args: anytype, captures: anytype) !void {
-    inline for (items) |it| switch (it) {
-        .text => |text| try writer.writeAll(text),
-        .statement => |stmt| try executeStatement(writer, stmt, args, captures),
-    };
+    inline for (items) |it| {
+        switch (it) {
+            .text => |text| try writer.writeAll(text),
+            .statement => |stmt| try executeStatement(writer, stmt, args, captures),
+        }
+    }
 }
 
 fn executeStatement(writer: anytype, comptime stmt: Statement, args: anytype, captures: anytype) !void {
@@ -30,19 +36,25 @@ fn executeStatement(writer: anytype, comptime stmt: Statement, args: anytype, ca
             const val = evaluateExpression(expr, args, captures);
             try print(writer, val);
         },
-        .for_loop => |loop| {
-            const iterable = evaluateExpression(loop.iterable, args, captures);
+        .@"for" => |loop| {
+            const iterable = evaluateExpression(loop.header.iterable, args, captures);
             const subtemplate = loop.subtemplate;
+            //std.log.debug("{any}", .{subtemplate});
             for (iterable) |v| {
                 try executeTemplate(
                     writer,
                     subtemplate,
                     args,
-                    addCapture(captures, loop.capture, v),
+                    addCapture(captures, loop.header.capture, v),
                 );
             }
         },
-        else => @compileError("TODO"),
+        .@"if" => |if_stmt| {
+            const condition = evaluateExpression(if_stmt.header.condition, args, captures);
+            const subtemplate = if_stmt.subtemplate;
+            if (condition) try executeTemplate(writer, subtemplate, args, captures);
+        },
+        //else => @compileError("TODO"),
     }
 }
 
@@ -112,120 +124,250 @@ fn addCapture(root: anytype, comptime name: []const u8, val: anytype) AddCapture
 
 const TemplateType = enum {
     root,
-    subtemplate,
+    for_block,
+    if_block,
 };
 
-fn parseTemplate(comptime tokens: TokenIter, comptime template_type: TemplateType) ParseResult([]const TemplateItem) {
+fn parseTemplate(
+    comptime tokens: []const TemplateToken,
+    comptime start: usize,
+    comptime template_type: TemplateType,
+) ParseResult(usize, []const TemplateItem) {
     comptime {
-        var iter = tokens;
-        var items: []const TemplateItem = &.{};
+        var i: usize = start;
         var current_text: []const u8 = "";
+        var items: []const TemplateItem = &.{};
 
-        parse_loop: while (iter.next()) |token| {
-            switch (token) {
-                .whitespace, .text => |text| current_text = current_text ++ text,
-                .open_bracket => {
-                    const next = iter.peek() orelse @compileError("Unexpected end of template");
-                    if (next == .open_bracket) {
-                        current_text = current_text ++ "{";
-                        _ = iter.next();
-                    } else {
-                        if (current_text.len != 0) {
-                            items = items ++ [_]TemplateItem{.{ .text = current_text }};
-                            current_text = "";
+        while (i < tokens.len) : (i += 1) {
+            switch (tokens[i]) {
+                .text => |text| current_text = current_text ++ text,
+                .whitespace => |wsp| {
+                    if (i != tokens.len - 1 and tokens[i + 1] == .control_block)
+                        if (tokens[i + 1].control_block.strip_before)
+                            continue;
+                    current_text = current_text ++ wsp;
+                },
+                .control_block => |cb| {
+                    if (current_text.len != 0) {
+                        items = items ++ [_]TemplateItem{.{ .text = current_text }};
+                        current_text = "";
+                    }
+
+                    switch (cb.block) {
+                        .expression => |expr| items = items ++ [_]TemplateItem{.{ .statement = .{ .expression = expr } }},
+                        .if_header => |header| {
+                            if (i != tokens.len - 1 and tokens[i + 1] == .whitespace and cb.strip_after) i += 1;
+                            const subtemplate = parseTemplate(tokens, i + 1, .if_block);
+                            items = items ++ [_]TemplateItem{.{
+                                .statement = .{
+                                    .@"if" = .{
+                                        .subtemplate = subtemplate.item,
+                                        .header = header,
+                                    },
+                                },
+                            }};
+                            i = subtemplate.new_iter;
+                        },
+                        .for_header => |header| {
+                            if (i != tokens.len - 1 and tokens[i + 1] == .whitespace and cb.strip_after) i += 1;
+                            const subtemplate = parseTemplate(tokens, i + 1, .for_block);
+                            items = items ++ [_]TemplateItem{.{
+                                .statement = .{
+                                    .@"for" = .{
+                                        .subtemplate = subtemplate.item,
+                                        .header = header,
+                                    },
+                                },
+                            }};
+                            i = subtemplate.new_iter;
+                        },
+                        .end_for => if (template_type == .for_block)
+                            break
+                        else
+                            @compileError("Unexpected /for tag"),
+                        .end_if => if (template_type == .if_block)
+                            break
+                        else
+                            @compileError("Unexpected /if tag"),
+                    }
+
+                    if (i != tokens.len - 1 and tokens[i] == .control_block) {
+                        if (tokens[i].control_block.strip_after and tokens[i + 1] == .whitespace) {
+                            i += 1;
                         }
-                        const result = parseExpressionOrStatement(iter, true);
-                        iter = result.new_iter;
-                        if (result.item == .end_for) {
-                            if (template_type == .subtemplate) break :parse_loop else @compileError("Unexpected end statement");
-                        }
-                        items = items ++ [_]TemplateItem{.{ .statement = result.item }};
                     }
                 },
-                .close_bracket => {
-                    const next = iter.next() orelse @compileError("Unexpected end of template");
-                    if (next == .close_bracket) current_text = current_text ++ "}" else @compileError("Unpaired close bracket, did you mean \"}}\"?");
-                },
-                .period => current_text = current_text ++ ".",
-                .pound => current_text = current_text ++ "#",
-                .pipe => current_text = current_text ++ "|",
-                .dollar => current_text = current_text ++ "$",
             }
-        }
+        } else if (template_type != .root) @compileError("End tag not found");
 
-        if (current_text.len != 0) {
-            items = items ++ [_]TemplateItem{.{ .text = current_text }};
-        }
+        if (current_text.len != 0) items = items ++ [_]TemplateItem{.{ .text = current_text }};
 
         return .{
-            .new_iter = iter,
+            .new_iter = i,
             .item = items,
         };
     }
 }
 
-fn parseExpressionOrStatement(
-    comptime tokens: TokenIter,
-    comptime as_statement: bool,
-) ParseResult(if (as_statement) Statement else Expression) {
+const TemplateToken = union(enum) {
+    text: []const u8,
+    whitespace: []const u8,
+    control_block: ControlBlock,
+};
+
+fn parseTemplateTokens(comptime tokens: ControlTokenIter) []const TemplateToken {
     comptime {
         var iter = tokens;
-        var stmt: Statement = while (iter.next()) |token| switch (token) {
+        var items: []const TemplateToken = &.{};
+
+        while (iter.next()) |token| switch (token) {
+            .whitespace => |wsp| items = items ++ [_]TemplateToken{.{ .whitespace = wsp }},
+            .text => |text| items = items ++ [_]TemplateToken{.{ .text = text }},
+            .open_bracket => {
+                const next = iter.next() orelse @compileError("Unexpected end of template");
+                if (next == .open_bracket) {
+                    items = items ++ [_]TemplateToken{.{ .text = "{" }};
+                } else {
+                    iter.putBack(next);
+                    const result = parseControlBlock(iter);
+                    iter = result.new_iter;
+                    items = items ++ [_]TemplateToken{.{ .control_block = result.item }};
+                }
+            },
+            .close_bracket => {
+                const next = iter.next() orelse @compileError("Unexpected end of template");
+                if (next == .close_bracket)
+                    items = items ++ [_]TemplateToken{.{ .text = "}" }}
+                else
+                    @compileError("Unpaired close bracket, did you mean \"}}\"?");
+            },
+            .period => items = items ++ [_]TemplateToken{.{ .text = "." }},
+            .pound => items = items ++ [_]TemplateToken{.{ .text = "#" }},
+            .pipe => items = items ++ [_]TemplateToken{.{ .text = "|" }},
+            .dollar => items = items ++ [_]TemplateToken{.{ .text = "$" }},
+            .slash => items = items ++ [_]TemplateToken{.{ .text = "/" }},
+            .equals => items = items ++ [_]TemplateToken{.{ .text = "=" }},
+        };
+
+        return items;
+    }
+}
+
+fn parseExpression(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, Expression) {
+    comptime {
+        var iter = tokens;
+
+        var expr: Expression = while (iter.next()) |token| switch (token) {
+            .whitespace => {},
+            .period => {
+                const names = parseDeref(iter);
+                iter = names.new_iter;
+                break .{ .arg_deref = names.item };
+            },
+            .dollar => {
+                const names = parseDeref(iter);
+                iter = names.new_iter;
+                break .{ .capture_deref = names.item };
+            },
+            else => @compileError("TODO"),
+        };
+
+        return .{
+            .new_iter = iter,
+            .item = expr,
+        };
+    }
+}
+
+fn parseControlBlock(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, ControlBlock) {
+    comptime {
+        var iter = tokens;
+
+        const strip_before = if (iter.next()) |first| blk: {
+            if (first == .equals) {
+                break :blk true;
+            }
+
+            iter.putBack(first);
+            break :blk false;
+        } else @compileError("Unexpected end of template");
+
+        var stmt: ControlBlock.Data = while (iter.next()) |token| switch (token) {
+            .equals => @compileError("Unexpected '='"),
             .whitespace => {},
             .pound => {
-                if (!as_statement) @compileError("Unexpected Token");
                 const next = iter.next() orelse @compileError("Unexpected end of template");
                 if (next != .text) @compileError("Expected keyword following '#' character");
                 const text = next.text;
                 const keyword = std.meta.stringToEnum(Keyword, text) orelse @compileError("Unknown keyword: " ++ text);
 
                 switch (keyword) {
-                    .end_for => break .{ .end_for = {} },
                     .@"for" => {
-                        const result = parseForLoop(iter);
-                        // statemnt already finished so just return
-                        return .{
-                            .new_iter = result.new_iter,
-                            .item = .{ .for_loop = result.item },
-                        };
+                        const result = parseForHeader(iter);
+                        iter = result.new_iter;
+                        break .{ .for_header = result.item };
+                    },
+                    .@"if" => {
+                        const result = parseIfHeader(iter);
+                        iter = result.new_iter;
+                        break .{ .if_header = result.item };
                     },
 
                     //else => @compileError("TODO"),
                 }
             },
-            .period => {
-                const names = parseDeref(iter);
-                iter = names.new_iter;
-                break .{ .expression = .{ .arg_deref = names.item } };
+            .slash => {
+                const next = iter.next() orelse @compileError("Unexpected end of template");
+                if (next != .text) @compileError("Expected keyword following '/' character");
+                const text = next.text;
+                const keyword = std.meta.stringToEnum(EndKeyword, text) orelse @compileError("Unknown keyword: " ++ text);
+
+                switch (keyword) {
+                    .@"for" => break .{ .end_for = {} },
+                    .@"if" => break .{ .end_if = {} },
+                }
             },
-            .dollar => {
-                const names = parseDeref(iter);
-                iter = names.new_iter;
-                break .{ .expression = .{ .capture_deref = names.item } };
+            .period, .dollar => {
+                iter.putBack(token);
+                const expr = parseExpression(iter);
+                iter = expr.new_iter;
+                break .{ .expression = expr.item };
             },
-            else => if (as_statement) @compileError("TODO") else break,
+            else => @compileError("TODO"),
         };
 
-        if (as_statement) {
-            // search for end of statement
-            while (iter.next()) |token| switch (token) {
-                .whitespace => {},
-                .close_bracket => return .{
-                    .new_iter = iter,
-                    .item = stmt,
+        // search for end of statement
+        var strip_after: bool = false;
+        while (iter.next()) |token| switch (token) {
+            .whitespace => {},
+            .equals => {
+                if (iter.peek()) |t| {
+                    if (t == .close_bracket) {
+                        strip_after = true;
+                        continue;
+                    }
+                }
+                @compileError("Unexpected '='");
+            },
+            .close_bracket => return .{
+                .new_iter = iter,
+                .item = .{
+                    .block = stmt,
+                    .strip_before = strip_before,
+                    .strip_after = strip_after,
                 },
-                else => {
-                    @compileLog(iter.row);
-                    @compileError("TODO" ++ @tagName(token));
-                },
-            };
+            },
+            else => {
+                @compileLog(iter.row);
+                @compileError("TODO" ++ @tagName(token));
+            },
+        };
 
-            @compileError("Unexpected end of template");
-        } else return .{ .new_iter = iter, .item = stmt.expression };
+        @compileError("Unexpected end of template");
     }
 }
 
-fn skipWhitespace(comptime tokens: TokenIter) TokenIter {
+fn skipWhitespace(comptime tokens: ControlTokenIter) ControlTokenIter {
     comptime {
         var iter = tokens;
         while (iter.peek()) |token| switch (token) {
@@ -237,7 +379,7 @@ fn skipWhitespace(comptime tokens: TokenIter) TokenIter {
     }
 }
 
-fn endStatement(comptime tokens: TokenIter) TokenIter {
+fn endControlBlock(comptime tokens: ControlTokenIter) ControlTokenIter {
     comptime {
         var iter = skipWhitespace(tokens);
 
@@ -247,9 +389,9 @@ fn endStatement(comptime tokens: TokenIter) TokenIter {
     }
 }
 
-fn parseForLoop(comptime tokens: TokenIter) ParseResult(ForLoop) {
+fn parseForHeader(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, ForHeader) {
     comptime {
-        const iterable = parseExpressionOrStatement(tokens, false);
+        const iterable = parseExpression(tokens);
         var iter = iterable.new_iter;
 
         iter = skipWhitespace(iter);
@@ -270,15 +412,32 @@ fn parseForLoop(comptime tokens: TokenIter) ParseResult(ForLoop) {
             const token = iter.next() orelse @compileError("Unexpected end of template");
             if (token != .pipe) @compileError("Unexpected token");
         }
-        iter = endStatement(iter);
 
-        const subtemplate = parseTemplate(iter, .subtemplate);
-
-        return .{ .new_iter = subtemplate.new_iter, .item = .{ .iterable = iterable.item, .subtemplate = subtemplate.item, .capture = capture } };
+        return .{
+            .new_iter = iter,
+            .item = .{
+                .iterable = iterable.item,
+                .capture = capture,
+            },
+        };
     }
 }
 
-fn parseDeref(comptime tokens: TokenIter) ParseResult([]const []const u8) {
+fn parseIfHeader(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, IfHeader) {
+    comptime {
+        const condition = parseExpression(tokens);
+        var iter = condition.new_iter;
+
+        return .{
+            .new_iter = iter,
+            .item = .{
+                .condition = condition.item,
+            },
+        };
+    }
+}
+
+fn parseDeref(comptime tokens: ControlTokenIter) ParseResult(ControlTokenIter, []const []const u8) {
     comptime {
         var iter = tokens;
         var fields: []const []const u8 = &.{};
@@ -305,9 +464,9 @@ fn parseDeref(comptime tokens: TokenIter) ParseResult([]const []const u8) {
     }
 }
 
-fn ParseResult(comptime T: type) type {
+fn ParseResult(comptime It: type, comptime T: type) type {
     return struct {
-        new_iter: TokenIter,
+        new_iter: It,
         item: T,
     };
 }
@@ -322,24 +481,55 @@ const Expression = union(enum) {
     capture_deref: []const []const u8,
 };
 
-const ForLoop = struct {
+const For = struct {
     subtemplate: []const TemplateItem,
+    header: ForHeader,
+};
+
+const ForHeader = struct {
     iterable: Expression,
     capture: []const u8,
 };
 
+const If = struct {
+    subtemplate: []const TemplateItem,
+    header: IfHeader,
+};
+
+const IfHeader = struct {
+    condition: Expression,
+};
+
 const Statement = union(enum) {
     expression: Expression,
-    for_loop: ForLoop,
-    end_for: void,
+    @"for": For,
+    @"if": If,
+};
+
+const ControlBlock = struct {
+    const Data = union(enum) {
+        expression: Expression,
+        for_header: ForHeader,
+        end_for: void,
+        if_header: IfHeader,
+        end_if: void,
+    };
+    block: Data,
+    strip_before: bool,
+    strip_after: bool,
 };
 
 const Keyword = enum {
     @"for",
-    end_for,
+    @"if",
 };
 
-const Token = union(enum) {
+const EndKeyword = enum {
+    @"for",
+    @"if",
+};
+
+const ControlToken = union(enum) {
     text: []const u8,
     open_bracket: void,
     close_bracket: void,
@@ -348,16 +538,18 @@ const Token = union(enum) {
     pound: void,
     pipe: void,
     dollar: void,
+    slash: void,
+    equals: void,
 };
 
-const TokenIter = struct {
+const ControlTokenIter = struct {
     start: usize = 0,
     text: []const u8,
-    peeked_token: ?Token = null,
+    peeked_token: ?ControlToken = null,
 
     row: usize = 0,
 
-    fn next(self: *TokenIter) ?Token {
+    fn next(self: *ControlTokenIter) ?ControlToken {
         if (self.peeked_token) |token| {
             self.peeked_token = null;
             return token;
@@ -375,6 +567,8 @@ const TokenIter = struct {
             '#' => return .{ .pound = {} },
             '|' => return .{ .pipe = {} },
             '$' => return .{ .dollar = {} },
+            '/' => return .{ .slash = {} },
+            '=' => return .{ .equals = {} },
             ' ', '\t', '\n', '\r' => {
                 var idx: usize = 0;
                 while (idx < remaining.len and std.mem.indexOfScalar(u8, " \t\n\r", remaining[idx]) != null) : (idx += 1) {}
@@ -386,7 +580,7 @@ const TokenIter = struct {
             },
             else => {
                 var idx: usize = 0;
-                while (idx < remaining.len and std.mem.indexOfScalar(u8, "{}.#|$ \t\n\r", remaining[idx]) == null) : (idx += 1) {}
+                while (idx < remaining.len and std.mem.indexOfScalar(u8, "{}.#|$/= \t\n\r", remaining[idx]) == null) : (idx += 1) {}
 
                 self.start += idx - 1;
                 return .{ .text = remaining[0..idx] };
@@ -394,9 +588,14 @@ const TokenIter = struct {
         }
     }
 
-    fn peek(self: *TokenIter) ?Token {
+    fn peek(self: *ControlTokenIter) ?ControlToken {
         const token = self.next();
         self.peeked_token = token;
         return token;
     }
+
+    fn putBack(self: *ControlTokenIter, token: ControlToken) void {
+        std.debug.assert(self.peeked_token == null);
+        self.peeked_token = token;
+    }
 };

src/template/test.tmp.html

@@ -1,7 +1,7 @@
 <!DOCTYPE html>
 <html>
   <head>
-    <title>{.community.name}</title>
+    <title>   {= .community.name =}    </title>
   </head>
 
   <body>
@@ -9,8 +9,13 @@
     <h2> {{ REAL BRACKETS }} </h2>
 
     <section>
-      {#for .baz |$f|}{#for $f |$b|}{$b}:{#end_for}
-      {#end_for}
+      {#for .baz |$f| =}
+      {= #for $f |$b| =}
+      {$b}:
+      {= /for =}
+      {= /for}
+      {#if .qux}qux!{/if=}
+      {#if .quxx}quxx!{/if}
     </section>
   </body>
 </html>