nff: improve diagnostics and vmap expression handling in parser and formatter

src/ast.rs

@@ -127,7 +127,7 @@ pub enum Expression {
 
     // Vmap expressions (value maps)
     Vmap {
-        expr: Box<Expression>,
+        expr: Option<Box<Expression>>,
         map: Vec<(Expression, Expression)>,
     },
 

src/parser.rs

@@ -311,9 +311,8 @@ impl Parser {
                         self.advance(); // consume 'policy'
                         let policy = self.parse_policy()?;
                         chain = chain.with_policy(policy);
+                        self.consume(TokenKind::Semicolon, "Expected ';' after policy")?;
                     }
-
-                    self.consume(TokenKind::Semicolon, "Expected ';' after policy")?;
                 }
                 Some(TokenKind::CommentLine(_)) => {
                     self.advance();
@@ -511,7 +510,7 @@ impl Parser {
 
                         // Return a vmap expression with the previous expression as the mapping target
                         return Ok(Expression::Vmap {
-                            expr: Box::new(expr),
+                            expr: Some(Box::new(expr)),
                             map,
                         });
                     }
@@ -828,8 +827,8 @@ impl Parser {
 
                 self.consume(TokenKind::RightBrace, "Expected '}' to close vmap")?;
 
-                // The expression that came before "vmap" is the expr being mapped
+                // No expression available at parse time, will be filled by post-processing if needed
-                let expr = Box::new(Expression::Identifier("dummy".to_string())); // This will be replaced in post-processing
+                let expr = None;
 
                 Ok(Expression::Vmap { expr, map })
             }
@@ -854,6 +853,23 @@ impl Parser {
                 self.consume(TokenKind::RightBrace, "Expected '}' to close set")?;
                 Ok(Expression::Set(elements))
             }
+            Some(TokenKind::Accept) => {
+                self.advance();
+                Ok(Expression::Identifier("accept".to_string()))
+            }
+            Some(TokenKind::Drop) => {
+                self.advance();
+                Ok(Expression::Identifier("drop".to_string()))
+            }
+            Some(TokenKind::Reject) => {
+                self.advance();
+                Ok(Expression::Identifier("reject".to_string()))
+            }
+            Some(TokenKind::Protocol) => {
+                self.advance(); // consume 'protocol'
+                let protocol = self.parse_identifier_or_keyword()?;
+                Ok(Expression::Protocol(protocol))
+            }
             _ => Err(ParseError::InvalidExpression {
                 message: format!(
                     "Unexpected token in expression: {}",

src/syntax.rs

@@ -175,10 +175,11 @@ impl NftablesFormatter {
             // Add policy on the same line if present
             if let Some(policy) = &chain.policy {
                 write!(output, " policy {}", policy).unwrap();
+                output.push_str(";\n");
+            } else {
+                output.push_str("\n");
             }
 
-            output.push_str(";\n");
-
             if !chain.rules.is_empty() && !self.config.optimize {
                 output.push('\n');
             }
@@ -294,8 +295,11 @@ impl NftablesFormatter {
             }
 
             Expression::Vmap { expr, map } => {
-                self.format_expression(output, expr);
+                if let Some(expr) = expr {
-                output.push_str(" vmap { ");
+                    self.format_expression(output, expr);
+                    output.push(' ');
+                }
+                output.push_str("vmap { ");
                 for (i, (key, value)) in map.iter().enumerate() {
                     if i > 0 {
                         output.push_str(", ");