Skip to main content

ast_grep_core/
source.rs

1//! This module defines the `Doc` and `Content` traits to abstract away source code encoding issues.
2//!
3//! ast-grep supports three kinds of encoding: utf-8 for CLI, utf-16 for nodeJS napi and `Vec<char>` for wasm.
4//! Different encoding will produce different tree-sitter Node's range and position.
5//!
6//! The `Content` trait is defined to abstract different encoding.
7//! It is used as associated type bound `Source` in the `Doc` trait.
8//! Its associated type `Underlying`  represents the underlying type of the content, e.g. `Vec<u8>`, `Vec<u16>`.
9//!
10//! `Doc` is a trait that defines a document that can be parsed by Tree-sitter.
11//! It has a `Source` associated type bounded by `Content` that represents the source code of the document,
12//! and a `Lang` associated type that represents the language of the document.
13
14use crate::{Position, language::Language, node::KindId};
15use std::borrow::Cow;
16use std::ops::Range;
17
18// https://github.com/tree-sitter/tree-sitter/blob/e4e5ffe517ca2c668689b24cb17c51b8c6db0790/cli/src/parse.rs
19#[derive(Debug)]
20pub struct Edit<S: Content> {
21  pub position: usize,
22  pub deleted_length: usize,
23  pub inserted_text: Vec<S::Underlying>,
24}
25
26/// NOTE: Some method names are the same as tree-sitter's methods.
27/// Fully Qualified Syntax may needed https://stackoverflow.com/a/44445976/2198656
28pub trait SgNode<'r>: Clone {
29  fn parent(&self) -> Option<Self>;
30  fn children(&self) -> impl ExactSizeIterator<Item = Self>;
31  // named_children is a subset of children, we only supply default impl here
32  // since underlying tree-sitter is using the same is_named check
33  // the only optimization we missed here is the iteractor size via named_child_count
34  fn named_children(&self) -> impl Iterator<Item = Self> {
35    self.children().filter(|n| n.is_named())
36  }
37  fn kind(&self) -> Cow<'_, str>;
38  fn kind_id(&self) -> KindId;
39  fn node_id(&self) -> usize;
40  fn range(&self) -> std::ops::Range<usize>;
41  fn start_pos(&self) -> Position;
42  fn end_pos(&self) -> Position;
43
44  // default implentation
45  fn ancestors(&self, _root: Self) -> impl Iterator<Item = Self> {
46    let mut ancestors = vec![];
47    let mut current = self.clone();
48    while let Some(parent) = current.parent() {
49      ancestors.push(parent.clone());
50      current = parent;
51    }
52    ancestors.reverse();
53    ancestors.into_iter()
54  }
55  fn dfs(&self) -> impl Iterator<Item = Self> {
56    let mut stack = vec![self.clone()];
57    std::iter::from_fn(move || {
58      if let Some(node) = stack.pop() {
59        let children: Vec<_> = node.children().collect();
60        stack.extend(children.into_iter().rev());
61        Some(node)
62      } else {
63        None
64      }
65    })
66  }
67  fn child(&self, nth: usize) -> Option<Self> {
68    self.children().nth(nth)
69  }
70  fn next(&self) -> Option<Self> {
71    let parent = self.parent()?;
72    let mut children = parent.children();
73    while let Some(child) = children.next() {
74      if child.node_id() == self.node_id() {
75        return children.next();
76      }
77    }
78    None
79  }
80  fn prev(&self) -> Option<Self> {
81    let parent = self.parent()?;
82    let children = parent.children();
83    let mut prev = None;
84    for child in children {
85      if child.node_id() == self.node_id() {
86        return prev;
87      }
88      prev = Some(child);
89    }
90    None
91  }
92  fn next_all(&self) -> impl Iterator<Item = Self> {
93    let mut next = self.next();
94    std::iter::from_fn(move || {
95      let n = next.clone()?;
96      next = n.next();
97      Some(n)
98    })
99  }
100  fn prev_all(&self) -> impl Iterator<Item = Self> {
101    let mut prev = self.prev();
102    std::iter::from_fn(move || {
103      let n = prev.clone()?;
104      prev = n.prev();
105      Some(n)
106    })
107  }
108  fn is_named(&self) -> bool {
109    true
110  }
111  /// N.B. it is different from is_named && is_leaf
112  /// if a node has no named children.
113  fn is_named_leaf(&self) -> bool {
114    self.is_leaf()
115  }
116  fn is_leaf(&self) -> bool {
117    self.children().count() == 0
118  }
119
120  // missing node is a tree-sitter specific concept
121  fn is_missing(&self) -> bool {
122    false
123  }
124  fn is_error(&self) -> bool {
125    false
126  }
127  fn is_extra(&self) -> bool {
128    false
129  }
130
131  fn field(&self, name: &str) -> Option<Self>;
132  fn field_children(&self, field_id: Option<u16>) -> impl Iterator<Item = Self>;
133  fn child_by_field_id(&self, field_id: u16) -> Option<Self>;
134}
135
136pub trait Doc: Clone + 'static {
137  type Source: Content;
138  type Lang: Language;
139  type Node<'r>: SgNode<'r>;
140  fn get_lang(&self) -> &Self::Lang;
141  fn get_source(&self) -> &Self::Source;
142  fn do_edit(&mut self, edit: &Edit<Self::Source>) -> Result<(), String>;
143  fn root_node(&self) -> Self::Node<'_>;
144  fn get_node_text<'a>(&'a self, node: &Self::Node<'a>) -> Cow<'a, str>;
145}
146
147pub trait Content: Sized {
148  type Underlying: Clone + PartialEq;
149  fn get_range(&self, range: Range<usize>) -> &[Self::Underlying];
150  /// Used for string replacement. We need this for
151  /// indentation and deindentation.
152  fn decode_str(src: &str) -> Cow<'_, [Self::Underlying]>;
153  /// Used for string replacement. We need this for
154  /// transformation.
155  fn encode_bytes(bytes: &[Self::Underlying]) -> Cow<'_, str>;
156  /// Get the character column at the given position
157  fn get_char_column(&self, column: usize, offset: usize) -> usize;
158}
159
160impl Content for String {
161  type Underlying = u8;
162  fn get_range(&self, range: Range<usize>) -> &[Self::Underlying] {
163    &self.as_bytes()[range]
164  }
165  fn decode_str(src: &str) -> Cow<'_, [Self::Underlying]> {
166    Cow::Borrowed(src.as_bytes())
167  }
168  fn encode_bytes(bytes: &[Self::Underlying]) -> Cow<'_, str> {
169    String::from_utf8_lossy(bytes)
170  }
171
172  /// This is an O(n) operation. We assume the col will not be a
173  /// huge number in reality. This may be problematic for special
174  /// files like compressed js
175  fn get_char_column(&self, _col: usize, offset: usize) -> usize {
176    let src = self.as_bytes();
177    let mut col = 0;
178    // TODO: is it possible to use SIMD here???
179    for &b in src[..offset].iter().rev() {
180      if b == b'\n' {
181        break;
182      }
183      // https://en.wikipedia.org/wiki/UTF-8#Description
184      if b & 0b1100_0000 != 0b1000_0000 {
185        col += 1;
186      }
187    }
188    col
189  }
190}