yara_x_fmt/

lib.rs

/*! A code formatter for YARA rules

This crate implements a code format for YARA in the spirit of other tools like
`rustfmt` and `gofmt`.

# Usage

```no_run
# use std::fs::File;
use yara_x_fmt::Formatter;

let input = File::open("original.yar").unwrap();
let output = File::create("formatted.yar").unwrap();

Formatter::new().format(input, output).unwrap();
```
*/
use std::io;
use std::io::Cursor;

use thiserror::Error;

use tokens::Token::*;
use tokens::TokenStream;
use yara_x_parser::cst::{CSTStream, Event, SyntaxKind};
use yara_x_parser::{Parser, Span};

use crate::align::Align;
use crate::format_hex_patterns::FormatHexPatterns;
use crate::indentation::AddIndentation;
use crate::tokens::categories::*;
use crate::tokens::*;
use crate::trailing_spaces::RemoveTrailingSpaces;

mod align;
mod bubble;
mod comments;
mod format_hex_patterns;
mod indentation;
mod processor;
mod tokens;
mod trailing_spaces;

#[cfg(test)]
mod tests;

/// Errors returned by [`Formatter::format`].
#[derive(Error, Debug)]
#[allow(clippy::large_enum_variant)]
pub enum Error {
    /// Error while reading from input.
    #[error("read error: {0}")]
    ReadError(io::Error),

    /// Error while writing to output.
    #[error("write error: {0}")]
    WriteError(io::Error),

    /// The input file contained invalid UTF-8.
    #[error("invalid UTF-8 at {0}")]
    InvalidUTF8(Span),
}

/// Specifies how to indent the formatted code.
#[derive(Copy, Clone)]
pub enum Indentation {
    /// Use a given number of spaces per indentation level.
    Spaces(usize),
    /// Use one tab per indentation level.
    Tabs,
}

/// Formats YARA source code automatically.
pub struct Formatter {
    align_metadata: bool,
    align_patterns: bool,
    indent_section_headers: bool,
    indent_section_contents: bool,
    newline_before_curly_brace: bool,
    empty_line_before_section_header: bool,
    empty_line_after_section_header: bool,
    tab_size: usize,
    indentation: Indentation,
}

impl Default for Formatter {
    fn default() -> Self {
        Self::new()
    }
}

// Formatter public API.
impl Formatter {
    /// Creates a new formatter.
    pub fn new() -> Self {
        Formatter {
            align_metadata: true,
            align_patterns: true,
            indent_section_headers: true,
            indent_section_contents: true,
            newline_before_curly_brace: false,
            empty_line_before_section_header: true,
            empty_line_after_section_header: false,
            tab_size: 4,
            indentation: Indentation::Spaces(2),
        }
    }

    /// Specify if the metadata block must be aligned.
    ///
    /// If true, the metadata block will be converted from this...
    ///
    /// ```text
    /// rule test {
    ///   meta:
    ///     short = "foo"
    ///     very_long = "bar"
    ///     even_longer = "baz"
    ///   condition:
    ///     ...
    /// }
    /// ```
    ///
    /// ... to this ...
    ///
    /// ```text
    /// rule test {
    ///   meta:
    ///     short       = "foo"
    ///     very_long   = "bar"
    ///     even_longer = "baz"
    ///   condition:
    ///     ...
    /// }
    /// ```
    ///
    /// The default value is `true`.
    pub fn align_metadata(mut self, yes: bool) -> Self {
        self.align_metadata = yes;
        self
    }

    /// Specify if the patterns definitions must be aligned.
    ///
    /// If true, the strings block will be converted from this...
    ///
    /// ```text
    /// rule test {
    ///   strings:
    ///     $short = "foo"
    ///     $very_long = "bar"
    ///     $even_longer = "baz"
    ///   condition:
    ///     ...
    /// }
    /// ```
    ///
    /// ... to this ...
    ///
    /// ```text
    /// rule test {
    ///   strings:
    ///     $short       = "foo"
    ///     $very_long   = "bar"
    ///     $even_longer = "baz"
    ///   condition:
    ///     ...
    /// }
    /// ```
    ///
    /// The default value is `true`.
    pub fn align_patterns(mut self, yes: bool) -> Self {
        self.align_patterns = yes;
        self
    }

    /// Specify if the section definitions must be aligned.
    ///
    /// If true, the section headers look like this...
    ///
    /// ```text
    /// rule test {
    ///   strings:
    ///     $short = "foo"
    ///     $very_long = "bar"
    ///     $even_longer = "baz"
    ///   condition:
    ///     ...
    /// }
    /// ```
    ///
    /// And if false, the section headers look like this...
    ///
    /// ```text
    /// rule test {
    /// strings:
    ///   $short       = "foo"
    ///   $very_long   = "bar"
    ///   $even_longer = "baz"
    /// condition:
    ///   ...
    /// }
    /// ```
    ///
    /// The default value is `true`.
    pub fn indent_section_headers(mut self, yes: bool) -> Self {
        self.indent_section_headers = yes;
        self
    }

    /// Specify if the section contents must be aligned.
    ///
    /// If true, the section contents look like this...
    ///
    /// ```text
    /// rule test {
    ///   strings:
    ///     $short = "foo"
    ///     $very_long = "bar"
    ///     $even_longer = "baz"
    ///   condition:
    ///     ...
    /// }
    /// ```
    ///
    /// And if false, the section contents look like this...
    ///
    /// ```text
    /// rule test {
    ///   strings:
    ///   $short       = "foo"
    ///   $very_long   = "bar"
    ///   $even_longer = "baz"
    ///   condition:
    ///   ...
    /// }
    /// ```
    ///
    /// The default value is `true`.
    pub fn indent_section_contents(mut self, yes: bool) -> Self {
        self.indent_section_contents = yes;
        self
    }

    /// Number of spaces to indent, if indenting at all. Set to 0 to use tabs.
    ///
    /// The default is `2`.
    #[deprecated(
        since = "1.7.0",
        note = "use `.indentation(Indentation::Spaces(n))` or `.indentation(Indentation::Tabs)` instead"
    )]
    pub fn indent_spaces(mut self, n: u8) -> Self {
        if n == 0 {
            self.indentation = Indentation::Tabs
        } else {
            self.indentation = Indentation::Spaces(n as usize)
        }
        self
    }

    /// Specifies how to indent the formatted source code.
    ///
    /// ```
    /// # use yara_x_fmt::{Formatter, Indentation};
    /// let indent_with_two_spaces = Formatter::new().indentation(Indentation::Spaces(2));
    /// let indent_with_tabs = Formatter::new().indentation(Indentation::Tabs);
    /// ```
    pub fn indentation(mut self, indentation: Indentation) -> Self {
        self.indentation = indentation;
        self
    }

    /// Specifies the tab size (in spaces) expected in the unformatted source
    /// code.
    ///
    /// If the input contains tab characters, the formatter uses this value to
    /// determine how many spaces each tab represents. Setting this incorrectly
    /// can lead to misaligned formatting when the code mixes tabs and spaces.
    ///
    /// Defaults to `4`.
    pub fn input_tab_size(mut self, tab_size: usize) -> Self {
        self.tab_size = tab_size;
        self
    }

    /// Specify if newline should be added before the opening curly brace in a
    /// rule declaration. If false the rule will look like this:
    ///
    /// ```text
    /// rule test {
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    /// And if true, the rule will look like this:
    ///
    /// ```text
    /// rule test
    /// {
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    /// The default value is `false`.
    pub fn newline_before_curly_brace(mut self, yes: bool) -> Self {
        self.newline_before_curly_brace = yes;
        self
    }

    /// Specify if an empty line should be added before the section header in a
    /// rule. If false the rule will look like this:
    ///
    /// ```text
    /// rule test {
    ///   meta:
    ///     foo = "bar"
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    /// And if true, the rule will look like this:
    ///
    /// ```text
    /// rule test {
    ///
    ///   meta:
    ///     foo = "bar"
    ///
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    /// The default value is `false`.
    pub fn empty_line_before_section_header(mut self, yes: bool) -> Self {
        self.empty_line_before_section_header = yes;
        self
    }

    /// Specify if an empty line should be added after the section header in a
    /// rule. If false the rule will look like this:
    ///
    /// ```text
    /// rule test {
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    /// And if true, the rule will look like this:
    ///
    /// ```text
    /// rule test {
    ///   condition:
    ///
    ///     true
    /// }
    /// ```
    ///
    /// The default value is `false`.
    pub fn empty_line_after_section_header(mut self, yes: bool) -> Self {
        self.empty_line_after_section_header = yes;
        self
    }

    /// Reads YARA source code from `input` and write it into `output` after
    /// formatting.
    ///
    /// Returns `true` if the output differs from the input.
    ///
    /// This function will fail if it can't read from the input, write to the
    /// output, or when the input contains invalid UTF-8 characters.
    pub fn format<R, W>(
        &self,
        mut input: R,
        mut output: W,
    ) -> Result<bool, Error>
    where
        R: io::Read,
        W: io::Write,
    {
        let mut invalid_utf8 = Option::None;
        let mut in_buf = Vec::with_capacity(256);

        input.read_to_end(&mut in_buf).map_err(Error::ReadError)?;

        let cst_stream = CSTStream::from(Parser::new(in_buf.as_slice()));

        // Inspect the CST stream looking for events indicating the presence of
        // invalid UTF-8 sequences.
        let events = cst_stream.into_iter().inspect(|evt| {
            if let Event::Token { kind: SyntaxKind::INVALID_UTF8, span } = evt
            {
                invalid_utf8.get_or_insert(span.clone());
            }
        });

        let tokens = Tokens::new(in_buf.as_slice(), events);
        let mut out_buf = Cursor::new(Vec::new());

        self.format_impl(tokens)
            .write_to(&mut out_buf)
            .map_err(Error::WriteError)?;

        if let Some(span) = invalid_utf8 {
            return Err(Error::InvalidUTF8(span));
        }

        let modified = in_buf.ne(out_buf.get_ref());

        output.write_all(out_buf.get_ref()).map_err(Error::WriteError)?;

        Ok(modified)
    }
}

// Private API for formatter.
impl Formatter {
    fn format_impl<'a, I>(&self, input: I) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        // The first step is inserting newlines between top-level statements
        // (rules, imports and includes) if they are in the same line.
        let tokens = processor::Processor::new(input)
            //
            // Insert newline in front of import and include statements, making
            // sure that they start at a new line. The newline is not inserted if
            // the statement is at the start of the file.
            //
            // Example:
            //
            // import "foo" import "bar"
            //
            // Inserts newline before import "bar".
            //
            .add_rule(
                |ctx| {
                    let next_token = ctx.token(1);
                    let prev_token = ctx.token(-1);

                    matches!(
                        next_token,
                        Begin(SyntaxKind::IMPORT_STMT)
                            | Begin(SyntaxKind::INCLUDE_STMT)
                    ) && prev_token.neq(&Begin(SyntaxKind::SOURCE_FILE))
                        && prev_token.is_not(*NEWLINE)
                },
                processor::actions::newline,
            )
            //
            // Insert newline in front of rule declarations, making sure that
            // rule declarations starts at a new line. The newline is not
            // inserted if the rule is at the start of the file.
            //
            // Example:
            //
            // rule foo { ... } rule bar { ... }
            //
            // Inserts newline before "rule bar".
            //
            .add_rule(
                |ctx| {
                    let next_token = ctx.token(1);
                    let prev_token = ctx.token(-1);

                    next_token.eq(&Begin(SyntaxKind::RULE_DECL))
                        && prev_token.neq(&Begin(SyntaxKind::SOURCE_FILE))
                        && prev_token.is_not(*NEWLINE)
                },
                processor::actions::newline,
            );

        // Process comments before removing whitespaces. Whitespaces must
        // be the original ones while the comments are being processed in
        // order to maintain comment indentation unaltered.
        let tokens =
            comments::CommentProcessor::new(tokens).tab_size(self.tab_size);

        // Remove all whitespaces from the original source.
        let tokens = processor::Processor::new(tokens).add_rule(
            |ctx| ctx.token(1).is(*WHITESPACE),
            processor::actions::drop,
        );

        // Displace tail comments and newlines up with respect to tokens that
        // indicate the start and end of a grammar rule. This effectively moves
        // such comments and newlines to the innermost grammar rule preceding
        // them. For example, suppose that we have the following rule:
        //
        //   import "test"  // Comment
        //
        //   rule test {
        //     strings:
        //       $a = "foo"
        //     condition:
        //       true
        //   }
        //
        // The sequence of tokens produced for that rule looks like:
        //
        //   Begin(SOURCE_FILE)
        //   Begin(IMPORT_STMT)
        //   Keyword("import")
        //   Literal("tests")
        //   End(IMPORT_STMT)
        //   Begin(RULE_DECL)
        //   TailComment("// Comment")
        //   Newline
        //   Keyword("rule")
        //   .... more
        //
        // Notice how TailComment("// Comment") and the Newline that follows
        // are placed just before the "rule" keyword, and inside the rule_decl
        // grammar rule. This is not the most natural place for this tail comment,
        // its natural place is just after the "tests" literal, like this:
        //
        //   Begin(SOURCE_FILE)
        //   Begin(IMPORT_STMT)
        //   Keyword("import")
        //   Literal("tests")
        //   TailComment("// Comment")
        //   Newline
        //   End(IMPORT_STMT)
        //   Begin(RULE_DECL)
        //   Keyword("rule")
        //   .... more
        //
        // That's exactly what the Bubble pipeline does. See the documentation for
        // the `bubble` module for more details.
        let tokens = bubble::Bubble::new(
            tokens,
            |token| {
                matches!(token, Token::TailComment(_)) || token.is(*NEWLINE)
            },
            |token| matches!(token, Token::Begin(_) | Token::End(_)),
        );

        // Displace newlines down with respect to tokens that indicate end of a
        // grammar rule. This effectively moves them to the outermost grammar
        // rule. See the documentation for the `bubble` module for more details.
        let tokens = bubble::Bubble::new(
            tokens,
            |token| matches!(token, Token::End(_)),
            |token| token.is(*NEWLINE),
        );

        // Remove newlines in multiple cases.
        let tokens = processor::Processor::new(tokens)
            // Remove all newlines at the beginning of the file. When the
            // processor is at the beginning of the file token(-1) is None.
            // Notice that this works because all these newlines have been
            // moved up and placed before Token::Begin(source_file) by the
            // first bubble pipeline.
            .add_rule(
                |ctx| ctx.token(-1).is(*NONE) && ctx.token(1).is(*NEWLINE),
                processor::actions::drop,
            )
            // Remove excess of consecutive newlines, only two consecutive
            // newlines are allowed.
            .add_rule(
                |ctx| {
                    ctx.token(1).is(*NEWLINE)
                        && ctx.token(2).is(*NEWLINE)
                        && ctx.token(3).is(*NEWLINE)
                },
                processor::actions::drop,
            )
            // Remove excess of newlines at the end of the file.
            .add_rule(
                |ctx| {
                    ctx.token(-1).is(*NEWLINE)
                        && ctx.token(1).is(*NEWLINE)
                        && ctx.token(2).is(*NONE)
                },
                processor::actions::drop,
            )
            // Remove newlines between rule tags and between rule modifiers.
            .add_rule(
                |ctx| {
                    (ctx.in_rule(SyntaxKind::RULE_MODS, false)
                        || ctx.in_rule(SyntaxKind::RULE_TAGS, false))
                        && ctx.token(-1).is_not(*COMMENT)
                        && ctx.token(1).is(*NEWLINE)
                },
                processor::actions::drop,
            )
            // Remove newlines after rule modifiers.
            .add_rule(
                |ctx| {
                    ctx.token(-1).eq(&End(SyntaxKind::RULE_MODS))
                        && ctx.token(1).is(*NEWLINE)
                },
                processor::actions::drop,
            )
            // Remove newlines after rule tags.
            .add_rule(
                |ctx| {
                    ctx.token(-1).eq(&End(SyntaxKind::RULE_TAGS))
                        && ctx.token(1).is(*NEWLINE)
                },
                processor::actions::drop,
            );

        let tokens = if self.newline_before_curly_brace {
            // Ensure we have a newline before the opening "{" in a rule
            // declaration. Be careful to only insert it if one does not already
            // exist.
            Box::new(processor::Processor::new(tokens).add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::RULE_DECL, false)
                        && ctx.token(1).eq(&LBRACE)
                        && ctx.token(-1).is_not(*NEWLINE)
                },
                processor::actions::newline,
            ))
        } else {
            // Remove newlines before the opening "{" in a rule declaration.
            // It takes into account that we can have one or two newlines
            // before the "{" character. In a previous step we have removed
            // consecutive newlines leaving two at most.
            Box::new(
                    processor::Processor::new(tokens)
                .add_rule(
                    |ctx| {
                        ctx.in_rule(SyntaxKind::RULE_DECL, false)
                            && ctx.token(1).is(*NEWLINE)
                            && (
                                // Newline followed by "{"  or ...
                                ctx.token(2).eq(&LBRACE) ||
                                // ... two newlines followed by "{"
                                ctx.token(2).is(*NEWLINE) && ctx.token(3).eq(&LBRACE)
                            )
                    },
                    processor::actions::drop,
                ))
        };

        let tokens = processor::Processor::new(tokens)
            //
            // Insert additional newline in front of a rule declaration that
            // already starts at a newline, but only if not preceded by a
            // comment. In other words, this adds empty lines in between rule
            // declarations, but don't do it if the rule is preceded by a
            // comment.
            //
            // Example:
            //
            //  rule foo {
            //    ...
            //  }
            //  rule bar {
            //    ...
            //  }
            //
            // Inserts newline before "rule bar".
            //
            .add_rule(
                |ctx| {
                    ctx.token(1).eq(&Begin(SyntaxKind::RULE_DECL))
                        && ctx.token(-1).is(*NEWLINE)
                        && ctx.token(-2).is_not(*NEWLINE | *COMMENT)
                },
                processor::actions::newline,
            );

        let tokens = if self.empty_line_before_section_header {
            Box::new(
                processor::Processor::new(tokens)
                    .set_passthrough(*CONTROL)
                    .add_rule(
                        |ctx| {
                            matches!(
                                ctx.token(1),
                                Keyword(b"meta")
                                    | Keyword(b"strings")
                                    | Keyword(b"condition")
                            ) && ctx.token(-1).is_not(*NEWLINE)
                        },
                        processor::actions::emptyline,
                    )
                    .add_rule(
                        |ctx| {
                            matches!(
                                ctx.token(1),
                                Keyword(b"meta")
                                    | Keyword(b"strings")
                                    | Keyword(b"condition")
                            ) && ctx.token(-1).is(*NEWLINE)
                                && ctx.token(-2).is_not(*NEWLINE)
                        },
                        processor::actions::newline,
                    ),
            )
        } else {
            Box::new(
                processor::Processor::new(tokens)
                    .set_passthrough(*CONTROL)
                    .add_rule(
                        |ctx| {
                            matches!(
                                ctx.token(1),
                                Keyword(b"meta")
                                    | Keyword(b"strings")
                                    | Keyword(b"condition")
                            ) && ctx.token(-1).is_not(*NEWLINE)
                        },
                        processor::actions::newline,
                    )
                    .add_rule(
                        |ctx| {
                            ctx.token(1).is(*NEWLINE)
                                && matches!(
                                    ctx.token(2),
                                    Keyword(b"meta")
                                        | Keyword(b"strings")
                                        | Keyword(b"condition")
                                )
                                && ctx.token(-1).is(*NEWLINE)
                        },
                        processor::actions::drop,
                    ),
            )
        };

        // Always remove empty line before the first section header.
        let tokens = processor::Processor::new(tokens)
            .set_passthrough(*CONTROL)
            .add_rule(
                |ctx| {
                    ctx.token(1).is(*NEWLINE)
                        && matches!(
                            ctx.token(2),
                            Keyword(b"meta")
                                | Keyword(b"strings")
                                | Keyword(b"condition")
                        )
                        && ctx.token(-1).is(*NEWLINE)
                        && ctx.token(-2).eq(&LBRACE)
                },
                processor::actions::drop,
            );

        let tokens = if self.empty_line_after_section_header {
            Box::new(
                processor::Processor::new(tokens)
                    .add_rule(
                        |ctx| {
                            ctx.token(-1).eq(&COLON)
                                && matches!(
                                    ctx.token(-2),
                                    Keyword(b"meta")
                                        | Keyword(b"strings")
                                        | Keyword(b"condition")
                                )
                                && ctx.token(1).is_not(*NEWLINE)
                        },
                        processor::actions::emptyline,
                    )
                    .add_rule(
                        |ctx| {
                            ctx.token(-1).eq(&COLON)
                                && matches!(
                                    ctx.token(-2),
                                    Keyword(b"meta")
                                        | Keyword(b"strings")
                                        | Keyword(b"condition")
                                )
                                && ctx.token(1).is(*NEWLINE)
                                && ctx.token(2).is_not(*NEWLINE)
                        },
                        processor::actions::newline,
                    ),
            )
        } else {
            Box::new(
                processor::Processor::new(tokens)
                    .add_rule(
                        |ctx| {
                            ctx.token(-1).eq(&COLON)
                                && matches!(
                                    ctx.token(-2),
                                    Keyword(b"meta")
                                        | Keyword(b"strings")
                                        | Keyword(b"condition")
                                )
                                && ctx.token(1).is_not(*NEWLINE)
                        },
                        processor::actions::newline,
                    )
                    .add_rule(
                        |ctx| {
                            ctx.token(-1).eq(&COLON)
                                && matches!(
                                    ctx.token(-2),
                                    Keyword(b"meta")
                                        | Keyword(b"strings")
                                        | Keyword(b"condition")
                                )
                                && ctx.token(1).is(*NEWLINE)
                                && ctx.token(2).is(*NEWLINE)
                        },
                        processor::actions::drop,
                    ),
            )
        };

        // Add newline at multiple places.
        let tokens = processor::Processor::new(tokens)
            .set_passthrough(*CONTROL)
            // Add a newline in front of meta definitions in the "meta" section.
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::META_DEF, false)
                        && ctx.token(1).is(*IDENTIFIER)
                        && ctx.token(-1).is_not(*NEWLINE)
                },
                processor::actions::newline,
            )
            // Add newline in front of pattern identifiers in the "strings"
            // section.
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::PATTERN_DEF, false)
                        && ctx.token(1).is(*IDENTIFIER)
                        && ctx.token(-1).is_not(*NEWLINE)
                },
                processor::actions::newline,
            )
            // Add newline before each identifier in a `with` statement.
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::WITH_DECL, false)
                        && ctx.token(1).is(*IDENTIFIER)
                        && ctx.token(-1).is_not(*NEWLINE)
                },
                processor::actions::newline,
            )
            // Add newline before the closing brace at the end of rule.
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::RULE_DECL, false)
                        && ctx.token(1).eq(&RBRACE)
                        && ctx.token(-1).is_not(*NEWLINE)
                },
                processor::actions::newline,
            )
            // Add empty line at the end of the file
            .add_rule(
                |ctx| ctx.token(1).is(*NONE) && ctx.token(-1).is_not(*NEWLINE),
                processor::actions::newline,
            );

        let tokens = FormatHexPatterns::new(tokens);

        let tokens: Box<dyn Iterator<Item = Token<'a>>> =
            if self.indent_section_headers {
                Box::new(Self::indent_body(tokens))
            } else {
                Box::new(tokens)
            };

        let tokens: Box<dyn Iterator<Item = Token<'a>>> =
            if self.indent_section_contents {
                Box::new(Self::indent_sections(tokens))
            } else {
                Box::new(tokens)
            };

        let tokens = Self::indent_hex_patterns(tokens);
        let tokens = Self::indent_parenthesized_exprs(tokens);
        let tokens = Self::indent_with_expr(tokens);

        // indent_body and indent_sections will insert Indentation tokens, but
        // won't take into account that those tokens must appear before the
        // newline they are expected to affect. This fixes the issue by moving
        // indentation tokens in front of newline tokens if they appear in
        // reverse order.
        let tokens = bubble::Bubble::new(
            tokens,
            |token| matches!(token, Indentation(_)),
            |token| token.is(*NEWLINE),
        );

        // Make sure that comments (except inline comments) are followed by
        // newline. In most cases this is already the case, but some of the rules
        // that remove newlines may remove those appearing after the comment.
        let tokens = processor::Processor::new(tokens)
            .set_passthrough(*CONTROL)
            .add_rule(
                |ctx| {
                    matches!(
                        ctx.token(-1),
                        HeadComment(_) | TailComment(_) | BlockComment(_)
                    ) && ctx.token(1).is_not(*NEWLINE)
                },
                processor::actions::newline,
            );

        let tokens = Self::add_spacing(tokens);
        let tokens = Self::align_comments_in_hex_patterns(tokens);

        let tokens: Box<dyn Iterator<Item = Token<'a>>> =
            if self.align_metadata {
                Box::new(Self::align_meta_section(tokens))
            } else {
                Box::new(tokens)
            };

        let tokens: Box<dyn Iterator<Item = Token<'a>>> =
            if self.align_patterns {
                Box::new(Self::align_patterns_section(tokens))
            } else {
                Box::new(tokens)
            };

        let tokens = AddIndentation::new(tokens, self.indentation);

        RemoveTrailingSpaces::new(tokens)
    }

    /// Indents the sections (meta, strings, condition) of a rule one level up.
    /// For example, for this input:
    ///
    /// ```text
    /// rule foo {
    /// strings:
    /// $a = "foo"
    /// condition:
    /// true
    /// }
    /// ```
    ///
    /// ... the result is ...
    ///
    /// ```text
    /// rule foo {
    /// strings:
    ///   $a = "foo"
    /// condition:
    ///   true
    /// }
    /// ```
    fn indent_sections<'a, I>(input: I) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        processor::Processor::new(input)
            // Ignore all comments
            .set_passthrough(*COMMENT)
            // Increase indentation after "condition:"
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::CONDITION_BLK, false)
                        && ctx.token(-1).eq(&COLON)
                },
                processor::actions::insert(Indentation(1)),
            )
            // Decrease indentation after the condition.
            .add_rule(
                |ctx| {
                    ctx.token(1).eq(&End(SyntaxKind::CONDITION_BLK))
                        && ctx.token(-1).neq(&Indentation(-1))
                },
                processor::actions::insert(Indentation(-1)),
            )
            // Increase indentation after "meta:"
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::META_BLK, false)
                        && ctx.token(-1).eq(&COLON)
                },
                processor::actions::insert(Indentation(1)),
            )
            // Decrease indentation after meta definitions
            .add_rule(
                |ctx| {
                    ctx.token(1).eq(&End(SyntaxKind::META_BLK))
                        && ctx.token(-1).neq(&Indentation(-1))
                },
                processor::actions::insert(Indentation(-1)),
            )
            // Increase indentation after "strings:"
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::PATTERNS_BLK, false)
                        && ctx.token(-1).eq(&COLON)
                },
                processor::actions::insert(Indentation(1)),
            )
            // Decrease indentation after pattern definitions.
            .add_rule(
                |ctx| {
                    ctx.token(1).eq(&End(SyntaxKind::PATTERNS_BLK))
                        && ctx.token(-1).neq(&Indentation(-1))
                },
                processor::actions::insert(Indentation(-1)),
            )
    }

    /// Indents the body of a rule. For this input...
    ///
    /// ```text
    /// rule foo {
    /// strings:
    /// $a = "foo"
    /// condition:
    /// true
    /// }
    /// ```
    /// ... the result is ...
    ///
    /// ```text
    /// rule foo {
    ///   strings:
    ///   $a = "foo"
    ///   condition:
    ///   true
    /// }
    /// ```
    ///
    fn indent_body<'a, I>(input: I) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        processor::Processor::new(input)
            // Ignore all comments.
            .set_passthrough(*COMMENT)
            // Increase indentation after the opening brace in a rule
            // declaration.
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::RULE_DECL, false)
                        && ctx.token(-1).eq(&LBRACE)
                },
                processor::actions::insert(Indentation(1)),
            )
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::RULE_DECL, false)
                        && ctx.token(1).eq(&RBRACE)
                        && ctx.token(-1).neq(&Indentation(-1))
                },
                processor::actions::insert(Indentation(-1)),
            )
    }

    /// Indent parenthesized expressions in rule conditions. For this input...
    ///
    /// ```text
    /// rule foo {
    /// strings:
    ///   $a = "foo"
    ///   $b = "bar"
    /// condition:
    ///    (
    ///    $a and $b
    ///    )
    /// }
    /// ```
    ///
    /// ... the result is ...
    ///
    /// ```text
    /// rule foo {
    /// strings:
    ///   $a = "foo"
    ///   $b = "bar"
    /// condition:
    ///    (
    ///      $a and $b
    ///    )
    /// }
    /// ```
    fn indent_parenthesized_exprs<'a, I>(input: I) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        processor::Processor::new(input)
            .set_passthrough(*COMMENT)
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::BOOLEAN_EXPR, true)
                        && ctx.token(-1).eq(&LPAREN)
                },
                processor::actions::insert(Indentation(1)),
            )
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::BOOLEAN_EXPR, true)
                        && ctx.token(1).eq(&RPAREN)
                        && ctx.token(-1).neq(&Indentation(-1))
                },
                processor::actions::insert(Indentation(-1)),
            )
    }

    /// Indent `with` expressions. For this input...
    ///
    /// ```text
    /// rule foo {
    /// condition:
    ///   with
    ///   foo = "foo"
    ///   bar = "bar": (...)
    /// }
    /// ```
    ///
    /// ... the result is ...
    ///
    /// ```text
    /// rule foo {
    /// condition:
    ///   with
    ///     foo = "foo"
    ///     bar = "bar": (...)
    /// }
    /// ```
    fn indent_with_expr<'a, I>(input: I) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        processor::Processor::new(input)
            // Ignore all comments.
            .set_passthrough(*COMMENT)
            // Increase indentation after the `with` keyword.
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::WITH_EXPR, false)
                        && ctx.token(-1).eq(&Keyword(b"with"))
                },
                processor::actions::insert(Indentation(1)),
            )
            // Decrease indentation after the `with` expression.
            .add_rule(
                |ctx| {
                    ctx.token(1).eq(&End(SyntaxKind::WITH_EXPR))
                        && ctx.token(-1).neq(&Indentation(-1))
                },
                processor::actions::insert(Indentation(-1)),
            )
    }

    /// Indent hex patterns. For this input...
    ///
    /// ```text
    ///   $a = {
    ///   00 ..
    ///   01 ..
    ///   02
    ///   }
    /// }
    /// ```
    ///
    /// ... the result is ...
    ///
    /// ```text
    ///   $a = {
    ///     00 ..
    ///     01 ..
    ///     02
    ///   }
    /// }
    /// ```
    fn indent_hex_patterns<'a, I>(input: I) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        processor::Processor::new(input)
            .set_passthrough(*COMMENT)
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::HEX_PATTERN, true)
                        && ctx.token(-1).eq(&LBRACE)
                },
                processor::actions::insert(Indentation(1)),
            )
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::HEX_PATTERN, true)
                        && ctx.token(1).eq(&RBRACE)
                        && ctx.token(-1).neq(&Indentation(-1))
                },
                processor::actions::insert(Indentation(-1)),
            )
    }

    /// Aligns the equals signs in pattern definitions. For example, for this
    /// input:
    ///
    /// ```text
    /// rule foo {
    ///   strings:
    ///     $short = "foo"
    ///     $very_long = "bar"
    ///     $even_longer = "baz"
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    /// ... the result is ...
    ///
    /// ```text
    /// rule foo {
    ///   strings:
    ///     $short       = "foo"
    ///     $very_long   = "bar"
    ///     $even_longer = "baz"
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    /// Pattern groups separated by empty lines are handled independently, for
    /// example:
    ///
    /// ```text
    /// rule foo {
    ///   strings:
    ///     $short     = "foo"
    ///     $very_long = "bar"
    ///
    ///     $even_longer    = "baz"
    ///     $longest_of_all = "qux"
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    /// The patterns in the first block are aligned together, but they are not
    /// influenced by the longer lines in the second block.
    ///
    /// The input must contain at least one newline character after each
    /// pattern definition.
    fn align_patterns_section<'a, I>(input: I) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        Self::align(input, SyntaxKind::PATTERNS_BLK, SyntaxKind::PATTERN_DEF)
    }

    /// Aligns the equals signs in metadata definitions.
    ///
    /// This is similar to [`Formatter::align_patterns_section`] but for metadata.
    fn align_meta_section<'a, I>(input: I) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        Self::align(input, SyntaxKind::META_BLK, SyntaxKind::META_DEF)
    }

    fn align<'a, I>(
        input: I,
        block_kind: SyntaxKind,
        item_kind: SyntaxKind,
    ) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        // First insert the alignment markers at the appropriate places...
        let input_with_markers = processor::Processor::new(input)
            // Insert `AlignmentBlockBegin` after the start of each block.
            .add_rule(
                move |ctx| ctx.token(-1).eq(&Begin(block_kind)),
                processor::actions::insert(AlignmentBlockBegin),
            )
            // Insert `AlignmentBlockEnd` just before the end each block.
            .add_rule(
                move |ctx| {
                    ctx.token(1).eq(&End(block_kind))
                        && ctx.token(-1).neq(&AlignmentBlockEnd)
                },
                processor::actions::insert(AlignmentBlockEnd),
            )
            .add_rule(
                move |ctx| {
                    ctx.in_rule(block_kind, false)
                        && ctx.token(-2).eq(&Newline)
                        && ctx.token(-1).eq(&Newline)
                },
                |ctx| {
                    ctx.push_output_token(Some(AlignmentBlockEnd));
                    ctx.push_output_token(Some(AlignmentBlockBegin));
                },
            )
            // Insert `AlignmentMarker` before each equal sign.
            .add_rule(
                move |ctx| {
                    ctx.in_rule(item_kind, false)
                        && ctx.token(1).eq(&EQUAL)
                        && ctx.token(-1).neq(&AlignmentMarker)
                },
                processor::actions::insert(AlignmentMarker),
            );

        // ... then pass the token stream with the markers to Aligner, which
        // returns a token stream that replaces the markers with the
        // appropriate number of spaces.
        Align::new(input_with_markers)
    }

    /// Aligns tail comments inside hex patterns. For this input...
    ///
    /// ```text
    /// rule foo {
    ///   strings:
    ///     $hex = {
    ///        00 01  // Lorem
    ///        00 01 02  // ipsum
    ///     }
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    /// ... the result is ...
    ///
    /// ```text
    /// rule foo {
    ///   strings:
    ///     $hex = {
    ///        00 01     // Lorem
    ///        00 01 02  // ipsum
    ///     }
    ///   condition:
    ///     true
    /// }
    /// ```
    ///
    fn align_comments_in_hex_patterns<'a, I>(
        input: I,
    ) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        // First insert the alignment markers at the appropriate places...
        let input_with_markers = processor::Processor::new(input)
            .add_rule(
                |ctx| ctx.token(-1).eq(&Begin(SyntaxKind::HEX_PATTERN)),
                processor::actions::insert(AlignmentBlockBegin),
            )
            .add_rule(
                |ctx| {
                    ctx.token(1).eq(&End(SyntaxKind::HEX_PATTERN))
                        && ctx.token(-1).neq(&AlignmentBlockEnd)
                },
                processor::actions::insert(AlignmentBlockEnd),
            )
            .add_rule(
                |ctx| {
                    ctx.in_rule(SyntaxKind::HEX_PATTERN, true)
                        && matches!(ctx.token(1), Token::TailComment(_))
                        && ctx.token(-1).neq(&AlignmentMarker)
                },
                processor::actions::insert(AlignmentMarker),
            );

        // ... then pass the token stream with the markers to Aligner, which
        // returns a token stream that replaces the markers with the
        // appropriate number of spaces.
        Align::new(input_with_markers)
    }

    fn add_spacing<'a, I>(input: I) -> impl TokenStream<'a> + 'a
    where
        I: TokenStream<'a> + 'a,
    {
        processor::Processor::new(input)
            // Ignore all control tokens.
            .set_passthrough(*CONTROL)
            // Insert spaces in-between all tokens, except in the following
            // cases:
            // - No space after "(" and "["
            // - No space before ")" and "]"
            // - No space before ":"
            // - No space before ","
            // - No space before or after ".." (e.g: (0..10))
            // - No space before or after "." (e.g: foo.bar)
            // - No space in-between identifiers and "(" or "[" (e.g: array[0],
            //   func("foo")).
            // - No space before or after "-" in pattern modifiers and hex jumps
            //   (e.g: xor(0-255), [0-10]).
            .add_rule(
                |ctx| {
                    let prev_token = ctx.token(-1);
                    let next_token = ctx.token(1);

                    // Insert space if previous token is anything except ( or [,
                    // and next token is anything except ) or ].
                    let add_space = prev_token.is(*TEXT ^ *LGROUPING)
                        && next_token.is(*TEXT ^ *RGROUPING);

                    let drop_space =
                        // Don't insert space if next token is ":"
                        next_token.eq(&COLON)
                        || next_token.eq(&COMMA)
                        // Don't insert space after "-"
                        || prev_token.eq(&HYPHEN)
                        // Don't insert spaces around "."
                        || prev_token.eq(&DOT)
                        || next_token.eq(&DOT)
                        // don't insert space in-between some identifier and "("
                        // or "[".
                        || prev_token.is(*IDENTIFIER)
                            && next_token.is(*LGROUPING)
                        // don't insert space in-between some identifier and "*"
                        // like in $a*
                        || prev_token.is(*IDENTIFIER)
                            && next_token.eq(&ASTERISK)
                        // don't insert spaces before "(" in pattern modifiers.
                        || ctx.in_rule(SyntaxKind::PATTERN_MOD, false)
                            && next_token.is(*LGROUPING)
                        // don't insert spaces before or after "-" in pattern
                        // modifiers and hex jumps.
                        || (ctx.in_rule(SyntaxKind::PATTERN_MOD, false) ||
                            ctx.in_rule(SyntaxKind::HEX_JUMP, false))
                            && (next_token.eq(&HYPHEN) || prev_token.eq(&HYPHEN));

                    add_space && !drop_space
                },
                processor::actions::space,
            )
            // Insert space before after inline comment.
            .add_rule(
                |ctx| {
                    matches!(
                        ctx.token(-1),
                        InlineComment(_)
                    )
                },
                processor::actions::space
            )
            // Insert two spaces before trailing comments in a line.
            .add_rule(
                |ctx| {
                    ctx.token(-1).is(*TEXT) &&
                    ctx.token(1).is(*COMMENT)
                },
                |ctx| {
                    ctx.push_output_token(Some(Whitespace));
                    ctx.push_output_token(Some(Whitespace));
                }
            )
    }
}