1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103

use crate::{
    ast::Ast,
    tokenizer::{Position, Token, TokenKind, TokenReference},
    visitors::VisitorMut,
};

#[derive(Default)]
struct UpdatePositionsRewriter {
    start_position: Position,
    next_is_new_line: bool,
}

impl UpdatePositionsRewriter {
    fn update_token<'ast>(&mut self, token: &Token<'ast>) -> Token<'ast> {
        let display = token.to_string();

        let mut end_position = self.start_position;

        if token.token_kind() != TokenKind::Eof {
            for character in display.chars() {
                if self.next_is_new_line {
                    self.next_is_new_line = false;
                    end_position.line += 1;
                    end_position.character = 1;
                }

                if character == '\n' {
                    self.next_is_new_line = true;
                } else {
                    end_position.character += 1;
                }

                end_position.bytes += character.len_utf8();
            }
        }

        let result = Token {
            start_position: self.start_position,
            end_position,
            token_type: token.token_type.to_owned(),
        };

        self.start_position = end_position;

        result
    }
}

impl<'ast> VisitorMut<'ast> for UpdatePositionsRewriter {
    fn visit_token_reference(&mut self, token: TokenReference<'ast>) -> TokenReference<'ast> {
        TokenReference::new(
            token
                .leading_trivia()
                .map(|token| self.update_token(token))
                .collect(),
            self.update_token(token.token()),
            token
                .trailing_trivia()
                .map(|token| self.update_token(token))
                .collect(),
        )
    }
}

impl Ast<'_> {
    /// Will update the positions of all the tokens in the tree
    /// Necessary if you are both mutating the tree and need the positions of the tokens
    pub fn update_positions(self) -> Self {
        let mut rewriter = UpdatePositionsRewriter {
            start_position: Position {
                bytes: 0,
                character: 1,
                line: 1,
            },

            ..Default::default()
        };

        rewriter.visit_ast(self)
    }
}

#[cfg(test)]
mod tests {
    use crate::{node::Node, parse};
    use pretty_assertions::assert_eq;

    #[test]
    fn test_update_positions_validity() {
        let ast = parse("local foo = 1\nlocal bar = 2").unwrap();
        let old_positions: Vec<_> = ast
            .tokens()
            .map(|token| (token.start_position(), token.end_position()))
            .collect();
        let ast = ast.update_positions();
        assert_eq!(
            old_positions,
            ast.tokens()
                .map(|token| (token.start_position(), token.end_position()))
                .collect::<Vec<_>>(),
        );
    }
}