dprint_swc_ext/common/
tokens.rs

1use rustc_hash::FxHashMap;
2
3use super::pos::*;
4use crate::swc::parser::token::TokenAndSpan;
5
6pub struct TokenContainer<'a> {
7  pub tokens: &'a [TokenAndSpan],
8  // Uses an FxHashMap because it has faster lookups for u32 keys than the default hasher.
9  start_to_index: FxHashMap<SourcePos, usize>,
10  end_to_index: FxHashMap<SourcePos, usize>,
11}
12
13impl<'a> TokenContainer<'a> {
14  pub fn new(tokens: &'a [TokenAndSpan]) -> Self {
15    TokenContainer {
16      tokens,
17      start_to_index: tokens.iter().enumerate().map(|(i, token)| (token.start(), i)).collect(),
18      end_to_index: tokens.iter().enumerate().map(|(i, token)| (token.end(), i)).collect(),
19    }
20  }
21
22  pub fn get_token_index_at_start(&self, start: SourcePos) -> Option<usize> {
23    self.start_to_index.get(&start).copied()
24  }
25
26  pub fn get_token_index_at_end(&self, end: SourcePos) -> Option<usize> {
27    self.end_to_index.get(&end).copied()
28  }
29
30  pub fn get_token_at_index(&self, index: usize) -> Option<&TokenAndSpan> {
31    self.tokens.get(index)
32  }
33
34  pub fn get_tokens_in_range(&self, start: SourcePos, end: SourcePos) -> &'a [TokenAndSpan] {
35    let start_index = self.get_leftmost_token_index(start);
36    let end_index = self.get_rightmost_token_index(end);
37
38    let start_index = start_index.unwrap_or_else(|| end_index.unwrap_or(0));
39    let end_index = end_index.map(|i| i + 1).unwrap_or(start_index);
40
41    &self.tokens[start_index..end_index]
42  }
43
44  fn get_leftmost_token_index(&self, start: SourcePos) -> Option<usize> {
45    if let Some(&start_index) = self.start_to_index.get(&start) {
46      Some(start_index)
47    // fallback
48    } else if let Some(&start_index) = self.end_to_index.get(&start) {
49      Some(start_index + 1)
50    } else {
51      // todo: binary search leftmost
52      for (i, token) in self.tokens.iter().enumerate() {
53        if token.start() >= start {
54          return Some(i);
55        }
56      }
57
58      None
59    }
60  }
61
62  fn get_rightmost_token_index(&self, end: SourcePos) -> Option<usize> {
63    if let Some(&end_index) = self.end_to_index.get(&end) {
64      Some(end_index)
65    // fallback
66    } else if let Some(&end_index) = self.start_to_index.get(&end) {
67      if end_index > 0 { Some(end_index - 1) } else { None }
68    } else {
69      // todo: binary search rightmost
70      for (i, token) in self.tokens.iter().enumerate().rev() {
71        if token.end() <= end {
72          return Some(i);
73        }
74      }
75
76      None
77    }
78  }
79
80  pub fn get_previous_token(&self, start: SourcePos) -> Option<&TokenAndSpan> {
81    let index = self.start_to_index.get(&start);
82    if let Some(&index) = index {
83      if index == 0 { None } else { Some(&self.tokens[index - 1]) }
84    } else {
85      // todo: binary search leftmost
86      let mut last_token = None;
87      for token in self.tokens {
88        if token.end() > start {
89          return last_token;
90        } else {
91          last_token = Some(token);
92        }
93      }
94
95      None
96    }
97  }
98
99  pub fn get_next_token(&self, end: SourcePos) -> Option<&TokenAndSpan> {
100    if let Some(index) = self.end_to_index.get(&end) {
101      self.tokens.get(index + 1)
102    } else {
103      // todo: binary search rightmost
104      for token in self.tokens {
105        if token.start() > end {
106          return Some(token);
107        }
108      }
109
110      None
111    }
112  }
113}
114
115#[cfg(test)]
116mod test {
117  use std::path::PathBuf;
118
119  use super::super::pos::SourcePos;
120  use super::TokenContainer;
121  use crate::common::SourceRangedForSpanned;
122  use crate::test_helpers::*;
123
124  #[test]
125  fn get_next_token() {
126    let (_, tokens, _, _) = get_swc_module(&PathBuf::from("path.js"), r#"let /* a */ a = 5;"#);
127    let token_container = TokenContainer::new(&tokens);
128    // low token of previous token
129    assert_eq!(token_container.get_next_token(SourcePos::new(0)).unwrap().start(), SourcePos::new(12));
130    // hi of previous token
131    assert_eq!(token_container.get_next_token(SourcePos::new(3)).unwrap().start(), SourcePos::new(12));
132    // in comment before token
133    assert_eq!(token_container.get_next_token(SourcePos::new(5)).unwrap().start(), SourcePos::new(12));
134    // in whitespace before token
135    assert_eq!(token_container.get_next_token(SourcePos::new(11)).unwrap().start(), SourcePos::new(12));
136    // at hi of last token
137    assert_eq!(token_container.get_next_token(SourcePos::new(18)), None);
138  }
139}