From 42698bb2195c6b06672dda1d19000004b08335a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marcel=20M=C3=BCller?= <neikos@neikos.email>
Date: Wed, 18 Mar 2026 09:59:36 +0100
Subject: [PATCH] Work on error messages
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Marcel Müller <neikos@neikos.email>
---
 benches/asting.rs                             |   4 +-
 src/compiler/mod.rs                           |  20 +-
 src/errors.rs                                 |  67 +++--
 src/eval/mod.rs                               |  19 +-
 src/lib.rs                                    |   4 +-
 src/parser/mod.rs                             | 238 ++++++++++--------
 ...__parser__tests__check_invalid_action.snap |  10 +-
 src/winnow_ext.rs                             |  51 ++++
 tests/checks.rs                               |   6 +-
 tests/errors/invalid_controls.1-parsed.snap   |  21 +-
 tests/errors/invalid_controls.2-ast.snap      |  37 ++-
 tests/errors/invalid_controls.nomo            |   1 +
 tests/file_tests.rs                           |  15 +-
 13 files changed, 333 insertions(+), 160 deletions(-)
 create mode 100644 src/winnow_ext.rs

diff --git a/benches/asting.rs b/benches/asting.rs
index 74e404f..5cf1b43 100644
--- a/benches/asting.rs
+++ b/benches/asting.rs
@@ -22,7 +22,7 @@ fn asting_benchmark(c: &mut Criterion) {
         parsing.bench_with_input(BenchmarkId::from_parameter(size), &input, |b, input| {
             b.iter(|| {
                 let tokens = nomo::lexer::parse(input.clone()).unwrap();
-                let _ast = nomo::parser::parse(tokens.tokens()).unwrap();
+                let _ast = nomo::parser::parse(input.clone(), tokens.tokens()).unwrap();
             });
         });
     }
@@ -46,7 +46,7 @@ fn asting_nested(c: &mut Criterion) {
         parsing.bench_with_input(BenchmarkId::from_parameter(size), &input, |b, input| {
             b.iter(|| {
                 let tokens = nomo::lexer::parse(input.clone()).unwrap();
-                let _ast = nomo::parser::parse(tokens.tokens()).unwrap();
+                let _ast = nomo::parser::parse(input.clone(), tokens.tokens()).unwrap();
             });
         });
     }
diff --git a/src/compiler/mod.rs b/src/compiler/mod.rs
index 59a7791..8a18816 100644
--- a/src/compiler/mod.rs
+++ b/src/compiler/mod.rs
@@ -579,14 +579,15 @@ fn emit_expr_load(
 #[cfg(test)]
 mod tests {
     use crate::compiler::emit_machine;
+    use crate::input::NomoInput;
 
     #[test]
     fn check_simple_variable_interpolation() {
-        let input = "Hello {{= world }}";
+        let input = NomoInput::from("Hello {{= world }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = crate::parser::parse(parsed.tokens()).unwrap();
+        let ast = crate::parser::parse(input, parsed.tokens()).unwrap();
 
         let emit = emit_machine(ast);
 
@@ -619,11 +620,12 @@ mod tests {
 
     #[test]
     fn check_if_else_if() {
-        let input = "{{ if foo }} foo {{ else if bar }} bar {{ else }} foobar {{ end }}";
+        let input =
+            NomoInput::from("{{ if foo }} foo {{ else if bar }} bar {{ else }} foobar {{ end }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = crate::parser::parse(parsed.tokens()).unwrap();
+        let ast = crate::parser::parse(input, parsed.tokens()).unwrap();
 
         let emit = emit_machine(ast);
 
@@ -632,11 +634,11 @@ mod tests {
 
     #[test]
     fn check_function_call() {
-        let input = "{{ if foo(23) }} bar {{ else }} foobar {{ end }}";
+        let input = NomoInput::from("{{ if foo(23) }} bar {{ else }} foobar {{ end }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = crate::parser::parse(parsed.tokens()).unwrap();
+        let ast = crate::parser::parse(input, parsed.tokens()).unwrap();
 
         let emit = emit_machine(ast);
 
diff --git a/src/errors.rs b/src/errors.rs
index 9d843ad..11df09a 100644
--- a/src/errors.rs
+++ b/src/errors.rs
@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use annotate_snippets::AnnotationKind;
 use annotate_snippets::Level;
+use annotate_snippets::Patch;
 use annotate_snippets::Renderer;
 use annotate_snippets::Snippet;
 use thiserror::Error;
@@ -15,25 +16,38 @@ use crate::parser::AstError;
 #[derive(Debug, Error)]
 pub struct AstFailure {
     errors: Vec<AstError>,
+    input: NomoInput,
 }
 
 impl std::fmt::Display for AstFailure {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.write_str("TODO")
+        self.render(f)
     }
 }
 
 impl AstFailure {
-    pub(crate) fn from_errors(errors: Vec<AstError>) -> AstFailure {
-        AstFailure { errors }
+    pub(crate) fn from_errors(errors: Vec<AstError>, source: NomoInput) -> AstFailure {
+        AstFailure {
+            errors,
+            input: source,
+        }
     }
 
     /// Create a CLI printable report
-    pub fn to_report(&self, source: &str) -> String {
-        let reports = self
-            .errors
-            .iter()
-            .map(|error| {
+    pub fn to_report(&self) -> String {
+        self.to_string()
+    }
+
+    /// Render this failure to the given formatter
+    ///
+    /// Note, you can also use [`to_string`](ToString::to_string), as this type also implements
+    /// [`Display`](std::fmt::Display).
+    pub fn render(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        let renderer = annotate_snippets::Renderer::styled()
+            .decor_style(annotate_snippets::renderer::DecorStyle::Unicode);
+
+        for error in &self.errors {
+            let mut snippets = vec![
                 annotate_snippets::Level::ERROR
                     .primary_title(
                         error
@@ -41,26 +55,39 @@ impl AstFailure {
                             .as_deref()
                             .unwrap_or("An error occurred while producing an Ast"),
                     )
-                    .element(annotate_snippets::Snippet::source(source).annotation(
-                        annotate_snippets::AnnotationKind::Primary.span(
-                            constrain_without_whitespace(
-                                source,
-                                error.span.clone().map(|s| s.range).unwrap_or_else(|| 0..0),
+                    .element(
+                        annotate_snippets::Snippet::source(self.input.as_str()).annotation(
+                            annotate_snippets::AnnotationKind::Primary.span(
+                                constrain_without_whitespace(
+                                    self.input.as_ref(),
+                                    error.span.clone().map(|s| s.range).unwrap_or_else(|| 0..0),
+                                ),
                             ),
                         ),
-                    ))
+                    )
                     .elements(
                         error
                             .help
                             .as_ref()
                             .map(|help| annotate_snippets::Level::HELP.message(help)),
-                    )
-            })
-            .collect::<Vec<_>>();
+                    ),
+            ];
+            if let Some((range, help)) = &error.replacement {
+                snippets.push(
+                    annotate_snippets::Level::NOTE
+                        .secondary_title("Try adding it")
+                        .element(
+                            Snippet::source(self.input.as_str())
+                                .patch(Patch::new(range.range.clone(), help.clone())),
+                        ),
+                );
+            }
 
-        let renderer = annotate_snippets::Renderer::styled()
-            .decor_style(annotate_snippets::renderer::DecorStyle::Unicode);
-        renderer.render(&reports)
+            writeln!(f, "{}", renderer.render(&snippets))?;
+            writeln!(f)?;
+        }
+
+        Ok(())
     }
 }
 
diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index 283432a..7a0957b 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -297,14 +297,15 @@ mod tests {
     use crate::eval::execute;
     use crate::functions::FunctionMap;
     use crate::functions::NomoFunctionError;
+    use crate::input::NomoInput;
 
     #[test]
     fn check_simple_variable_interpolation() {
-        let input = "Hello {{= world }}";
+        let input = NomoInput::from("Hello {{= world }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = crate::parser::parse(parsed.tokens()).unwrap();
+        let ast = crate::parser::parse(input, parsed.tokens()).unwrap();
 
         let emit = crate::compiler::emit_machine(ast);
 
@@ -321,11 +322,11 @@ mod tests {
 
     #[test]
     fn check_method_call() {
-        let input = "Hello {{= foo(world) }}";
+        let input = NomoInput::from("Hello {{= foo(world) }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = crate::parser::parse(parsed.tokens()).unwrap();
+        let ast = crate::parser::parse(input, parsed.tokens()).unwrap();
 
         let emit = crate::compiler::emit_machine(ast);
 
@@ -348,11 +349,11 @@ mod tests {
 
     #[test]
     fn check_conditional_access() {
-        let input = "Hello {{= unknown? }}";
+        let input = NomoInput::from("Hello {{= unknown? }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = crate::parser::parse(parsed.tokens()).unwrap();
+        let ast = crate::parser::parse(input, parsed.tokens()).unwrap();
 
         let emit = crate::compiler::emit_machine(ast);
 
diff --git a/src/lib.rs b/src/lib.rs
index b4b047e..abddcf5 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -137,6 +137,8 @@ unstable_pub!(
     mod parser
 );
 
+mod winnow_ext;
+
 /// Errors in this library
 pub mod errors;
 
@@ -203,7 +205,7 @@ impl Nomo {
     ) -> Result<(), NomoError> {
         let source = value.into();
         let parse = lexer::parse(source.clone())?;
-        let ast = parser::parse(parse.tokens())?;
+        let ast = parser::parse(source.clone(), parse.tokens())?;
 
         let instructions = compiler::emit_machine(ast);
 
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 5763f45..decda96 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -20,6 +20,7 @@ use winnow::error::AddContext;
 use winnow::error::FromRecoverableError;
 use winnow::error::ModalError;
 use winnow::error::ParserError;
+use winnow::stream::Location;
 use winnow::stream::Offset;
 use winnow::stream::Recoverable;
 use winnow::stream::Stream;
@@ -28,6 +29,7 @@ use winnow::token::any;
 
 use crate::SourceSpan;
 use crate::errors::AstFailure;
+use crate::input::NomoInput;
 use crate::lexer::TemplateToken;
 use crate::lexer::TokenKind;
 use crate::lexer::TokenOperator;
@@ -50,6 +52,7 @@ pub struct AstError {
     pub(crate) message: Option<String>,
     pub(crate) help: Option<String>,
     pub(crate) span: Option<crate::SourceSpan>,
+    pub(crate) replacement: Option<(SourceSpan, String)>,
 
     is_fatal: bool,
 }
@@ -60,6 +63,7 @@ impl AstError {
             message: None,
             help: None,
             span: None,
+            replacement: None,
 
             is_fatal: false,
         }
@@ -74,6 +78,11 @@ impl AstError {
         self.help = Some(help.to_string());
         self
     }
+
+    fn replacement(mut self, span: SourceSpan, replacement: &str) -> Self {
+        self.replacement = Some((span, replacement.to_string()));
+        self
+    }
 }
 
 impl ModalError for AstError {
@@ -132,6 +141,7 @@ impl AddContext<Input<'_>, AstError> for AstError {
     ) -> Self {
         self.message = context.message.or(self.message);
         self.help = context.help.or(self.help);
+        self.replacement = context.replacement.or(self.replacement);
         self
     }
 }
@@ -162,7 +172,7 @@ impl<'input> Parser<Input<'input>, TemplateToken, AstError> for TokenKind {
     }
 }
 
-pub fn parse(input: &[TemplateToken]) -> Result<TemplateAst<'_>, AstFailure> {
+pub fn parse(source: NomoInput, input: &[TemplateToken]) -> Result<TemplateAst<'_>, AstFailure> {
     let (_remaining, val, errors) = parse_asts.recoverable_parse(TokenSlice::new(input));
 
     if errors.is_empty()
@@ -170,7 +180,7 @@ pub fn parse(input: &[TemplateToken]) -> Result<TemplateAst<'_>, AstFailure> {
     {
         Ok(TemplateAst { root: val })
     } else {
-        Err(AstFailure::from_errors(errors))
+        Err(AstFailure::from_errors(errors, source))
     }
 }
 
@@ -309,26 +319,42 @@ fn parse_action<'input>(input: &mut Input<'input>) -> Result<TemplateAstExpr<'in
         alt((
             parse_conditional_chain,
             parse_for_chain,
-            (parse_block(
-                cut_err(not(repeat_till(
-                    0..,
-                    parse_expression,
-                    peek((ws, TokenKind::RightDelim)),
-                )
-                .map(|((), _)| ())))
-                .context(
-                    AstError::ctx()
-                        .msg("Standlone action block")
-                        .help("If you want to output this expression, add a '=' to the block"),
-                )
+            parse_block(parse_unknown_action)
                 .take()
                 .map(TemplateAstExpr::Invalid),
-            )),
         )),
     )
     .parse_next(input)
 }
 
+fn parse_unknown_action<'input>(
+    input: &mut Input<'input>,
+) -> Result<TemplateAstExpr<'input>, AstError> {
+    let expression = peek(parse_expression).parse_next(input);
+
+    let is_expr = match expression {
+        Ok(_) => true,
+        Err(err) if err.is_backtrack() => true,
+        _ => false,
+    };
+
+    if is_expr {
+        return Err(AstError::ctx()
+            .msg("Standlone action block")
+            .help("If you want to output this expression, add a '=' to the block")
+            .cut());
+    }
+
+    let keywords = peek(parse_keyword).parse_next(input);
+
+    let is_keyword = keywords.is_ok();
+
+    if is_keyword {
+        return Err(AstError::ctx().msg("Unexpected keyword").cut());
+    }
+
+    Err(AstError::ctx().cut())
+}
 fn parse_for_chain<'input>(input: &mut Input<'input>) -> Result<TemplateAstExpr<'input>, AstError> {
     trace("for_loop", |input: &mut Input<'input>| {
         let for_block = parse_for_loop.map(Box::new).parse_next(input)?;
@@ -368,26 +394,34 @@ fn parse_for_chain<'input>(input: &mut Input<'input>) -> Result<TemplateAstExpr<
 fn parse_for_loop<'input>(input: &mut Input<'input>) -> Result<TemplateAstExpr<'input>, AstError> {
     trace(
         "for_loop_inner",
-        parse_block(
-            (
-                ws,
-                TokenKind::For,
-                ws,
-                cut_err(TokenKind::Ident.context(AstError::ctx().msg("Expected identifier here"))),
-                ws,
-                cut_err(
-                    TokenKind::In.context(AstError::ctx().msg("Missing `in` in `for _ in <expr>`")),
+        parse_block(|input: &mut Input<'input>| {
+            let _ignored = surrounded(ws, TokenKind::For).parse_next(input)?;
+
+            let value_ident =
+                cut_err(TokenKind::Ident.context(AstError::ctx().msg("Missing ident here")))
+                    .parse_next(input)?;
+
+            let _ignored = cut_err(
+                preceded(ws, TokenKind::In).context(
+                    AstError::ctx()
+                        .msg("Missing `in` in `for .. in ` loop")
+                        .replacement(
+                            SourceSpan {
+                                range: input.current_token_start()..(input.current_token_start()),
+                            },
+                            " in",
+                        ),
                 ),
-                ws,
-                parse_expression.map(Box::new),
             )
-                .map(|(_, _for, _, value_ident, _, _in, _, value_expression)| {
-                    TemplateAstExpr::For {
-                        value_ident,
-                        value_expression,
-                    }
-                }),
-        ),
+            .parse_next(input)?;
+
+            let value_expression = parse_expression.map(Box::new).parse_next(input)?;
+
+            Ok(TemplateAstExpr::For {
+                value_ident,
+                value_expression,
+            })
+        }),
     )
     .parse_next(input)
 }
@@ -614,22 +648,25 @@ fn parse_operand<'input>(input: &mut Input<'input>) -> Result<TemplateAstExpr<'i
     .parse_next(input)
 }
 
-fn parse_keywords_fail<'input>(
-    input: &mut Input<'input>,
-) -> Result<TemplateAstExpr<'input>, AstError> {
-    let value = alt((
+fn parse_keyword<'input>(input: &mut Input<'input>) -> Result<TemplateToken, AstError> {
+    alt((
         TokenKind::ConditionalIf,
+        TokenKind::ConditionalElse,
         TokenKind::For,
         TokenKind::End,
         TokenKind::In,
     ))
-    .take()
-    .map(TemplateAstExpr::Invalid)
-    .parse_next(input)?;
+    .parse_next(input)
+}
 
-    cut_err(fail::<_, (), _>.context(AstError::ctx().msg("Found literal, expected expression")))
-        .value(value)
-        .parse_next(input)
+fn parse_keywords_fail<'input>(
+    input: &mut Input<'input>,
+) -> Result<TemplateAstExpr<'input>, AstError> {
+    let _value = parse_keyword.parse_next(input)?;
+
+    Err(AstError::ctx()
+        .msg("Found literal, expected expression")
+        .cut())
 }
 
 fn parse_literal<'input>(input: &mut Input<'input>) -> Result<TemplateAstExpr<'input>, AstError> {
@@ -735,6 +772,7 @@ mod tests {
     use winnow::combinator::fail;
     use winnow::stream::TokenSlice;
 
+    use crate::input::NomoInput;
     use crate::lexer::TokenKind;
     use crate::parser::AstError;
     use crate::parser::AstFailure;
@@ -744,25 +782,22 @@ mod tests {
     use crate::parser::parse_block;
     use crate::parser::parse_end;
 
-    fn panic_pretty<'a>(
-        input: &'_ str,
-        tokens: Result<TemplateAst<'a>, AstFailure>,
-    ) -> TemplateAst<'a> {
+    fn panic_pretty<'a>(tokens: Result<TemplateAst<'a>, AstFailure>) -> TemplateAst<'a> {
         match tokens {
             Ok(ast) => ast,
             Err(failure) => {
-                panic!("{}", failure.to_report(input));
+                panic!("{}", failure);
             }
         }
     }
 
     #[test]
     fn check_only_content() {
-        let input = "Hello World";
+        let input = NomoInput::from("Hello World");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = parse(parsed.tokens()).unwrap();
+        let ast = parse(input, parsed.tokens()).unwrap();
 
         insta::assert_debug_snapshot!(ast, @r#"
         TemplateAst {
@@ -777,11 +812,11 @@ mod tests {
 
     #[test]
     fn check_simple_variable_interpolation() {
-        let input = "Hello {{= world }}";
+        let input = NomoInput::from("Hello {{= world }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = parse(parsed.tokens()).unwrap();
+        let ast = parse(input, parsed.tokens()).unwrap();
 
         insta::assert_debug_snapshot!(ast, @r#"
         TemplateAst {
@@ -805,11 +840,11 @@ mod tests {
 
     #[test]
     fn check_simple_if() {
-        let input = "{{ if foo }} Hiii {{ end }}";
+        let input = NomoInput::from("{{ if foo }} Hiii {{ end }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast, @r#"
         TemplateAst {
@@ -850,35 +885,39 @@ mod tests {
 
     #[test]
     fn check_invalid_action() {
-        let input = r#"{{ value }}
+        let input = NomoInput::from(
+            r#"{{ value }}
         {{ value }}
         {{ value }}
         {{ value }}
-        {{ value }}"#;
+        {{ value }}"#,
+        );
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = parse(parsed.tokens()).unwrap_err();
+        let ast = parse(input, parsed.tokens()).unwrap_err();
 
-        insta::assert_snapshot!(ast.to_report(input));
+        insta::assert_snapshot!(ast);
     }
 
     #[test]
     fn check_nested_simple_if() {
-        let input = r#"{{ if foo }}
+        let input = NomoInput::from(
+            r#"{{ if foo }}
             {{ if bar }}
                 Hiii
             {{ end }}
         {{ end }}
 
         {{= value }}
-        "#;
+        "#,
+        );
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
         insta::assert_debug_snapshot!("simple_if_tokens", parsed);
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!("simple_if_ast", ast);
     }
@@ -887,9 +926,9 @@ mod tests {
     fn check_parsing_block() {
         use winnow::RecoverableParser;
 
-        let input = "{{ foo }}";
+        let input = NomoInput::from("{{ foo }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input).unwrap();
 
         let result = alt((
             parse_end,
@@ -924,6 +963,7 @@ mod tests {
                             range: 2..6,
                         },
                     ),
+                    replacement: None,
                     is_fatal: false,
                 },
             ],
@@ -933,11 +973,11 @@ mod tests {
 
     #[test]
     fn check_empty_if_output() {
-        let input = "{{ if foo }}{{ end }}";
+        let input = NomoInput::from("{{ if foo }}{{ end }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast, @r#"
         TemplateAst {
@@ -970,99 +1010,101 @@ mod tests {
 
     #[test]
     fn check_if_else() {
-        let input = "{{ if foo }} foo {{ else }} bar {{ end }}";
+        let input = NomoInput::from("{{ if foo }} foo {{ else }} bar {{ end }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast);
     }
 
     #[test]
     fn check_if_else_if() {
-        let input = "{{ if foo }} foo {{ else if bar }} bar {{ end }}";
+        let input = NomoInput::from("{{ if foo }} foo {{ else if bar }} bar {{ end }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast);
     }
 
     #[test]
     fn check_trim_whitespace() {
-        let input = "{{ if foo -}} foo {{- else if bar -}} bar {{- end }}";
+        let input = NomoInput::from("{{ if foo -}} foo {{- else if bar -}} bar {{- end }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast);
     }
 
     #[test]
     fn check_for_loop() {
-        let input = "{{ for value in array }} Hi: {{= value }} {{ else }} No Content :C {{ end }}";
+        let input = NomoInput::from(
+            "{{ for value in array }} Hi: {{= value }} {{ else }} No Content :C {{ end }}",
+        );
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast);
     }
 
     #[test]
     fn check_math_expression() {
-        let input = "{{= 5 * 3 + 2 / 3 }}";
+        let input = NomoInput::from("{{= 5 * 3 + 2 / 3 }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast);
     }
 
     #[test]
     fn check_logical_expression() {
-        let input = "{{= true && false || 3 >= 2 && 5 == 2 }}";
+        let input = NomoInput::from("{{= true && false || 3 >= 2 && 5 == 2 }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast);
     }
 
     #[test]
     fn check_function_call() {
-        let input = "{{= foo(2 * 3, bar(2 + baz)) }}";
+        let input = NomoInput::from("{{= foo(2 * 3, bar(2 + baz)) }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast);
     }
 
     #[test]
     fn check_conditional_access() {
-        let input = "{{= foo? }}";
+        let input = NomoInput::from("{{= foo? }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast);
     }
 
     #[test]
     fn check_access_operator() {
-        let input = "{{= foo?.bar }}";
+        let input = NomoInput::from("{{= foo?.bar }}");
 
-        let parsed = crate::lexer::parse(input.into()).unwrap();
+        let parsed = crate::lexer::parse(input.clone()).unwrap();
 
-        let ast = panic_pretty(input, parse(parsed.tokens()));
+        let ast = panic_pretty(parse(input, parsed.tokens()));
 
         insta::assert_debug_snapshot!(ast);
     }
diff --git a/src/parser/snapshots/nomo__parser__tests__check_invalid_action.snap b/src/parser/snapshots/nomo__parser__tests__check_invalid_action.snap
index 0587130..d4b9694 100644
--- a/src/parser/snapshots/nomo__parser__tests__check_invalid_action.snap
+++ b/src/parser/snapshots/nomo__parser__tests__check_invalid_action.snap
@@ -1,6 +1,6 @@
 ---
 source: src/parser/mod.rs
-expression: ast.to_report(input)
+expression: ast
 ---
 error: Standlone action block
   ╭▸ 
@@ -8,23 +8,31 @@ expression: ast.to_report(input)
   │    ━━━━━
   │
   ╰ help: If you want to output this expression, add a '=' to the block
+
 error: Standlone action block
   ╭▸ 
 2 │         {{ value }}
   │            ━━━━━
+  │
   ╰ help: If you want to output this expression, add a '=' to the block
+
 error: Standlone action block
   ╭▸ 
 3 │         {{ value }}
   │            ━━━━━
+  │
   ╰ help: If you want to output this expression, add a '=' to the block
+
 error: Standlone action block
   ╭▸ 
 4 │         {{ value }}
   │            ━━━━━
+  │
   ╰ help: If you want to output this expression, add a '=' to the block
+
 error: Standlone action block
   ╭▸ 
 5 │         {{ value }}
   │            ━━━━━
+  │
   ╰ help: If you want to output this expression, add a '=' to the block
diff --git a/src/winnow_ext.rs b/src/winnow_ext.rs
new file mode 100644
index 0000000..4d7450a
--- /dev/null
+++ b/src/winnow_ext.rs
@@ -0,0 +1,51 @@
+use std::marker::PhantomData;
+
+use winnow::Parser;
+use winnow::error::AddContext;
+use winnow::error::ParserError;
+use winnow::stream::Stream;
+
+pub trait ParserExt<I, O, E>: Parser<I, O, E> {
+    fn with_context<F, C>(self, context: F) -> WithContext<Self, I, O, E, F, C>
+    where
+        Self: Sized,
+        I: Stream,
+        E: AddContext<I, C> + ParserError<I>,
+        F: Fn(&<I as Stream>::Slice, &E) -> C,
+    {
+        WithContext {
+            parser: self,
+            context,
+            _pd: PhantomData,
+        }
+    }
+}
+
+pub struct WithContext<P, I, O, E, F, C> {
+    parser: P,
+    context: F,
+    _pd: PhantomData<(I, O, E, C)>,
+}
+
+impl<P, I, O, E, F, C> Parser<I, O, E> for WithContext<P, I, O, E, F, C>
+where
+    P: Parser<I, O, E>,
+    I: Stream,
+    E: AddContext<I, C> + ParserError<I>,
+    F: Fn(&<I as Stream>::Slice, &E) -> C,
+{
+    fn parse_next(&mut self, input: &mut I) -> winnow::Result<O, E> {
+        let start = input.checkpoint();
+
+        let res = self.parser.parse_next(input);
+
+        res.map_err(|err| {
+            let offset = input.offset_from(&start);
+            input.reset(&start);
+            let taken = input.next_slice(offset);
+
+            let context = (self.context)(&taken, &err);
+            err.add_context(input, &start, context)
+        })
+    }
+}
diff --git a/tests/checks.rs b/tests/checks.rs
index 35e0976..29278b1 100644
--- a/tests/checks.rs
+++ b/tests/checks.rs
@@ -7,11 +7,13 @@ fn check_files() {
     for file in files {
         let input = std::fs::read_to_string(file.unwrap().path()).unwrap();
 
-        let Ok(parsed) = nomo::lexer::parse(input.into()) else {
+        let input = nomo::input::NomoInput::from(input);
+
+        let Ok(parsed) = nomo::lexer::parse(input.clone()) else {
             continue;
         };
 
-        let Ok(ast) = nomo::parser::parse(parsed.tokens()) else {
+        let Ok(ast) = nomo::parser::parse(input, parsed.tokens()) else {
             continue;
         };
 
diff --git a/tests/errors/invalid_controls.1-parsed.snap b/tests/errors/invalid_controls.1-parsed.snap
index 13fc397..1ae070d 100644
--- a/tests/errors/invalid_controls.1-parsed.snap
+++ b/tests/errors/invalid_controls.1-parsed.snap
@@ -2,7 +2,7 @@
 source: tests/file_tests.rs
 expression: parsed
 info:
-  input: "{{ if }} {{ end }}\n{{ if if }} {{ end }}\n{{ if if }} {{ for foo in bar }} {{ end }} {{ end }}\n{{ for in bar }} {{ end }}\n{{ for blah bar }} {{ end }}\n{{ else }}"
+  input: "{{ if }} {{ end }}\n{{ if if }} {{ end }}\n{{ if if }} {{ for foo in bar }} {{ end }} {{ end }}\n{{ for in bar }} {{ end }}\n{{ for blah bar }} {{ end }}\n{{ for blah}} {{ end }}\n{{ else }}"
   context: {}
 ---
 ParsedTemplate {
@@ -99,8 +99,21 @@ ParsedTemplate {
         [Whitespace]"\n" (149..150),
         [LeftDelim]"{{" (150..152),
         [Whitespace]" " (152..153),
-        [ConditionalElse]"else" (153..157),
-        [Whitespace]" " (157..158),
-        [RightDelim]"}}" (158..160),
+        [For]"for" (153..156),
+        [Whitespace]" " (156..157),
+        [Ident]"blah" (157..161),
+        [RightDelim]"}}" (161..163),
+        [Whitespace]" " (163..164),
+        [LeftDelim]"{{" (164..166),
+        [Whitespace]" " (166..167),
+        [End]"end" (167..170),
+        [Whitespace]" " (170..171),
+        [RightDelim]"}}" (171..173),
+        [Whitespace]"\n" (173..174),
+        [LeftDelim]"{{" (174..176),
+        [Whitespace]" " (176..177),
+        [ConditionalElse]"else" (177..181),
+        [Whitespace]" " (181..182),
+        [RightDelim]"}}" (182..184),
     ],
 }
diff --git a/tests/errors/invalid_controls.2-ast.snap b/tests/errors/invalid_controls.2-ast.snap
index 265172d..e3d5494 100644
--- a/tests/errors/invalid_controls.2-ast.snap
+++ b/tests/errors/invalid_controls.2-ast.snap
@@ -2,31 +2,50 @@
 source: tests/file_tests.rs
 expression: ast
 info:
-  input: "{{ if }} {{ end }}\n{{ if if }} {{ end }}\n{{ if if }} {{ for foo in bar }} {{ end }} {{ end }}\n{{ for in bar }} {{ end }}\n{{ for blah bar }} {{ end }}\n{{ else }}"
+  input: "{{ if }} {{ end }}\n{{ if if }} {{ end }}\n{{ if if }} {{ for foo in bar }} {{ end }} {{ end }}\n{{ for in bar }} {{ end }}\n{{ for blah bar }} {{ end }}\n{{ for blah}} {{ end }}\n{{ else }}"
   context: {}
 ---
 error: Expected an expression after 'if'
   ╭▸ 
 1 │ {{ if }} {{ end }}
-  │    ━━
-  ╰╴
+  ╰╴   ━━
+
 error: Expected an expression after 'if'
   ╭▸ 
 2 │ {{ if if }} {{ end }}
   ╰╴      ━━
+
 error: Expected an expression after 'if'
   ╭▸ 
 3 │ {{ if if }} {{ for foo in bar }} {{ end }} {{ end }}
   ╰╴      ━━
-error: Expected identifier here
+
+error: Missing ident here
   ╭▸ 
 4 │ {{ for in bar }} {{ end }}
   ╰╴       ━━━━━━
-error: Missing `in` in `for _ in <expr>`
+
+error: Missing `in` in `for .. in ` loop
   ╭▸ 
 5 │ {{ for blah bar }} {{ end }}
-  ╰╴       ━━━━━━━━
-error: An error occurred while producing an Ast
+  │        ━━━━━━━━
+  ╰╴
+note: Try adding it
+  ╭╴
+5 │ {{ for blah in bar }} {{ end }}
+  ╰╴            ++
+
+error: Missing `in` in `for .. in ` loop
   ╭▸ 
-6 │ {{ else }}
-  ╰╴   ━━━━━━━
+6 │ {{ for blah}} {{ end }}
+  │        ━━━━
+  ╰╴
+note: Try adding it
+  ╭╴
+6 │ {{ for blah in}} {{ end }}
+  ╰╴            ++
+
+error: Unexpected keyword
+  ╭▸ 
+7 │ {{ else }}
+  ╰╴   ━━━━
diff --git a/tests/errors/invalid_controls.nomo b/tests/errors/invalid_controls.nomo
index f64a1b2..a9d079f 100644
--- a/tests/errors/invalid_controls.nomo
+++ b/tests/errors/invalid_controls.nomo
@@ -5,4 +5,5 @@
 {{ if if }} {{ for foo in bar }} {{ end }} {{ end }}
 {{ for in bar }} {{ end }}
 {{ for blah bar }} {{ end }}
+{{ for blah}} {{ end }}
 {{ else }}
\ No newline at end of file
diff --git a/tests/file_tests.rs b/tests/file_tests.rs
index 5f5b325..bd3334e 100644
--- a/tests/file_tests.rs
+++ b/tests/file_tests.rs
@@ -5,6 +5,7 @@ use std::path::Path;
 
 use nomo::Context;
 use nomo::functions::FunctionMap;
+use nomo::input::NomoInput;
 
 test_each_file::test_each_path! { for ["nomo"] in "./tests/cases/" as cases => check_for_input }
 
@@ -43,16 +44,18 @@ fn check_for_input([path]: [&Path; 1]) {
         context.try_insert(k, v).unwrap();
     }
 
-    let parsed = nomo::lexer::parse(input.into()).unwrap();
+    let input = NomoInput::from(input);
+
+    let parsed = nomo::lexer::parse(input.clone()).unwrap();
 
     let _guard = settings.bind_to_scope();
 
     insta::assert_debug_snapshot!(format!("{basename}.1-parsed"), parsed);
 
-    let ast = match nomo::parser::parse(parsed.tokens()) {
+    let ast = match nomo::parser::parse(input, parsed.tokens()) {
         Ok(ast) => ast,
         Err(err) => {
-            eprintln!("{}", err.to_report(input));
+            eprintln!("{}", err);
             panic!("Could not evaluate ast");
         }
     };
@@ -97,7 +100,9 @@ fn check_errors([path]: [&Path; 1]) {
 
     let _guard = settings.bind_to_scope();
 
-    let parsed = nomo::lexer::parse(input.into()).map_err(|err| err.to_report());
+    let input = NomoInput::from(input);
+
+    let parsed = nomo::lexer::parse(input.clone()).map_err(|err| err.to_report());
 
     match &parsed {
         Ok(parsed) => insta::assert_debug_snapshot!(format!("{basename}.1-parsed"), parsed),
@@ -108,7 +113,7 @@ fn check_errors([path]: [&Path; 1]) {
         return;
     };
 
-    let ast = nomo::parser::parse(parsed.tokens()).map_err(|err| err.to_report(input));
+    let ast = nomo::parser::parse(input, parsed.tokens()).map_err(|err| err.to_string());
 
     match &ast {
         Ok(ast) => insta::assert_debug_snapshot!(format!("{basename}.2-ast"), ast),