[rust] Support attach_comments in Rust bindings

CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) a
 
 ### Added
 
+- The Rust bindings now expose the `attach_comments` parse option, along with a `Node#comments` accessor for reading the comments attached to a node.
 - `pm_options_stack_budget_set` was added to the C API. It sets the number of bytes of C stack the parser may use before reporting a `stack_exhausted` error, instead of overflowing the stack. `SIZE_MAX` means no budget, which is the default. The nesting depth limit alone assumes a stack of roughly 8MB, so embedders that parse on a smaller stack, such as a thread's or a fiber's, should pass a budget.
 - A new `stack` error level, for errors that should raise `SystemStackError`. Its only error is `stack_exhausted`, reported once per parse when the stack budget is used up. It is separate from `nesting_too_deep`, which stays a syntax error, because the same source may parse with more stack. When errors are formatted, a `stack_exhausted` error takes precedence over the syntax errors around it, as argument and load errors do. The level is `:stack` in Ruby, `"stack"` in JavaScript, `ERROR_STACK` in Java, and `PM_ERROR_LEVEL_STACK` in C and Rust.
 

rust/ruby-prism-sys/build/main.rs

@@ -199,7 +199,10 @@ fn generate_bindings(ruby_include_path: &Path, clang_args: &[String], layout_tes
         .allowlist_function("pm_magic_comment_key")
         .allowlist_function("pm_magic_comment_next")
         .allowlist_function("pm_magic_comment_value")
+        .allowlist_function("pm_node_comment")
+        .allowlist_function("pm_node_comments_size")
         .allowlist_function("pm_node_list_append")
+        .allowlist_function("pm_options_attach_comments_set")
         .allowlist_function("pm_options_command_line_set")
         .allowlist_function("pm_options_encoding_locked_set")
         .allowlist_function("pm_options_encoding_set")

rust/ruby-prism/build.rs

@@ -542,7 +542,7 @@ use std::marker::PhantomData;
 
 #[allow(clippy::wildcard_imports)]
 use ruby_prism_sys::*;
-use crate::{{ConstantId, ConstantList, Integer, Location, NodeList}};
+use crate::{{Comments, ConstantId, ConstantList, Integer, Location, NodeList}};
 "
     )?;
 
@@ -614,6 +614,24 @@ impl<'pr> Node<'pr> {{
     writeln!(file, "    }}")?;
     writeln!(file)?;
 
+    writeln!(file, "    /// Returns an iterator over the comments attached to this node, ordered by")?;
+    writeln!(file, "    /// start offset.")?;
+    writeln!(file, "    ///")?;
+    writeln!(file, "    /// Comments are only attached when the source is parsed with the")?;
+    writeln!(file, "    /// [`attach_comments`](crate::Options::attach_comments) option.")?;
+    writeln!(file, "    #[allow(clippy::match_same_arms)]")?;
+    writeln!(file, "    #[must_use]")?;
+    writeln!(file, "    pub fn comments(&self) -> Comments<'pr> {{")?;
+    writeln!(file, "        let (parser, node) = match *self {{")?;
+    for node in &config.nodes {
+        writeln!(file, "            Self::{} {{ parser, pointer, .. }} => (parser, pointer.cast::<pm_node_t>()),", node.name)?;
+    }
+    writeln!(file, "        }};")?;
+    writeln!(file, "        let count = unsafe {{ pm_node_comments_size(parser, node) }};")?;
+    writeln!(file, "        Comments::for_node(parser, node, count)")?;
+    writeln!(file, "    }}")?;
+    writeln!(file)?;
+
     for node in &config.nodes {
         writeln!(file, "    /// Returns the node as a `{}`.", node.name)?;
         writeln!(file, "    #[must_use]")?;

rust/ruby-prism/src/lib.rs

@@ -30,8 +30,8 @@ pub use self::parse_result::{Comment, CommentType, Comments, Diagnostic, Diagnos
 pub use self::string_query::{InvalidEncodingError, StringQuery};
 
 use ruby_prism_sys::{
-    pm_arena_new, pm_options_command_line_set, pm_options_encoding_locked_set, pm_options_encoding_set, pm_options_filepath_set, pm_options_free, pm_options_frozen_string_literal_set, pm_options_line_set, pm_options_main_script_set, pm_options_new, pm_options_partial_script_set,
-    pm_options_scope_forwarding_set, pm_options_scope_init, pm_options_scope_local_mut, pm_options_scope_mut, pm_options_scopes_init, pm_options_t, pm_options_version_set, pm_parse, pm_parser_new, pm_string_init,
+    pm_arena_new, pm_options_attach_comments_set, pm_options_command_line_set, pm_options_encoding_locked_set, pm_options_encoding_set, pm_options_filepath_set, pm_options_free, pm_options_frozen_string_literal_set, pm_options_line_set, pm_options_main_script_set, pm_options_new,
+    pm_options_partial_script_set, pm_options_scope_forwarding_set, pm_options_scope_init, pm_options_scope_local_mut, pm_options_scope_mut, pm_options_scopes_init, pm_options_t, pm_options_version_set, pm_parse, pm_parser_new, pm_string_init,
 };
 
 /// The version of Ruby syntax to parse with.
@@ -155,6 +155,7 @@ impl Scope {
 }
 
 /// Options that can be passed to the parser.
+#[allow(clippy::struct_excessive_bools)]
 #[derive(Debug, Clone, Default)]
 pub struct Options {
     filepath: Option<String>,
@@ -166,6 +167,7 @@ pub struct Options {
     version: Option<Version>,
     main_script: bool,
     partial_script: bool,
+    attach_comments: bool,
     scopes: Vec<Scope>,
 }
 
@@ -235,6 +237,14 @@ impl Options {
         self
     }
 
+    /// Sets the attach comments option. When enabled, comments are attached to
+    /// the nodes they belong to and can be read with [`Node::comments`].
+    #[must_use]
+    pub const fn attach_comments(mut self, attach_comments: bool) -> Self {
+        self.attach_comments = attach_comments;
+        self
+    }
+
     /// Adds a scope to the options.
     #[must_use]
     pub fn scope(mut self, scope: Scope) -> Self {
@@ -300,6 +310,10 @@ impl Options {
             unsafe { pm_options_partial_script_set(opts, true) };
         }
 
+        if self.attach_comments {
+            unsafe { pm_options_attach_comments_set(opts, true) };
+        }
+
         if !self.scopes.is_empty() {
             unsafe { pm_options_scopes_init(opts, self.scopes.len()) };
 
@@ -384,7 +398,35 @@ pub fn parse_with_options<'a>(source: &'a [u8], options: &'a ParseOptions) -> Pa
 
 #[cfg(test)]
 mod tests {
-    use super::parse;
+    use super::{parse, parse_with_options, Options};
+
+    #[test]
+    fn attach_comments_test() {
+        let source = "foo # comment 1\nbar # comment 2\n";
+
+        let options = Options::default().attach_comments(true).build();
+        let result = parse_with_options(source.as_ref(), &options);
+        let statements = result.node().as_program_node().unwrap().statements();
+        let body: Vec<_> = statements.body().iter().collect();
+
+        assert_eq!(body.len(), 2);
+        assert_eq!(body[0].comments().len(), 1);
+
+        let comments: Vec<_> = body[0].comments().collect();
+        assert_eq!(comments.len(), 1);
+        assert_eq!(comments[0].text(), b"# comment 1");
+
+        let comments: Vec<_> = body[1].comments().collect();
+        assert_eq!(comments.len(), 1);
+        assert_eq!(comments[0].text(), b"# comment 2");
+
+        // Without the option, no comments are attached to the nodes.
+        let result = parse(source.as_ref());
+        let statements = result.node().as_program_node().unwrap().statements();
+        let body: Vec<_> = statements.body().iter().collect();
+        assert_eq!(body[0].comments().len(), 0);
+        assert_eq!(body[1].comments().len(), 0);
+    }
 
     #[test]
     fn comments_test() {

rust/ruby-prism/src/parse_result/comments.rs

@@ -2,7 +2,7 @@
 
 use std::marker::PhantomData;
 
-use ruby_prism_sys::{pm_comment_location, pm_comment_t, pm_comment_type, pm_comment_type_t, pm_magic_comment_key, pm_magic_comment_t, pm_magic_comment_value, pm_parser_start, pm_parser_t};
+use ruby_prism_sys::{pm_comment_location, pm_comment_t, pm_comment_type, pm_comment_type_t, pm_magic_comment_key, pm_magic_comment_t, pm_magic_comment_value, pm_node_comment, pm_node_t, pm_parser_start, pm_parser_t};
 
 use super::Location;
 
@@ -54,34 +54,73 @@ impl<'pr> Comment<'pr> {
     }
 }
 
-/// An iterator over comments collected from the parse result.
+enum CommentsSource {
+    Slice(Vec<*const pm_comment_t>),
+    Node { node: *const pm_node_t, count: u32 },
+}
+
+/// An iterator over comments collected from the parse result or attached to a node.
 pub struct Comments<'pr> {
-    ptrs: Vec<*const pm_comment_t>,
+    source: CommentsSource,
     index: usize,
     parser: *const pm_parser_t,
     marker: PhantomData<&'pr pm_comment_t>,
 }
 
 impl Comments<'_> {
     pub(crate) const fn new(ptrs: Vec<*const pm_comment_t>, parser: *const pm_parser_t) -> Self {
-        Comments { ptrs, index: 0, parser, marker: PhantomData }
+        Comments {
+            source: CommentsSource::Slice(ptrs),
+            index: 0,
+            parser,
+            marker: PhantomData,
+        }
+    }
+
+    pub(crate) const fn for_node(parser: *const pm_parser_t, node: *const pm_node_t, count: u32) -> Self {
+        Comments {
+            source: CommentsSource::Node { node, count },
+            index: 0,
+            parser,
+            marker: PhantomData,
+        }
+    }
+
+    const fn remaining(&self) -> usize {
+        match &self.source {
+            CommentsSource::Slice(ptrs) => ptrs.len().saturating_sub(self.index),
+            CommentsSource::Node { count, .. } => (*count as usize).saturating_sub(self.index),
+        }
     }
 }
 
 impl<'pr> Iterator for Comments<'pr> {
     type Item = Comment<'pr>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.index < self.ptrs.len() {
-            let comment = self.ptrs[self.index];
-            self.index += 1;
-            Some(Comment { raw: comment, parser: self.parser, marker: PhantomData })
-        } else {
-            None
-        }
+        let comment = match &self.source {
+            CommentsSource::Slice(ptrs) => *ptrs.get(self.index)?,
+            CommentsSource::Node { node, count } => {
+                if self.index >= *count as usize {
+                    return None;
+                }
+
+                unsafe { pm_node_comment(self.parser, *node, u32::try_from(self.index).ok()?) }
+            },
+        };
+
+        self.index += 1;
+        Some(Comment { raw: comment, parser: self.parser, marker: PhantomData })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.remaining();
+        (remaining, Some(remaining))
     }
 }
 
+impl ExactSizeIterator for Comments<'_> {}
+
 /// A magic comment that was found during parsing.
 #[derive(Debug)]
 pub struct MagicComment<'pr> {