From e82adf92b74ae039e4bbc56cb484b4ddb1db85d0 Mon Sep 17 00:00:00 2001
From: PoiScript <poiscript@gmail.com>
Date: Mon, 1 Apr 2024 11:41:40 +0800
Subject: [PATCH] feat: Org::node_at_offset api

---
 src/org.rs | 39 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/org.rs b/src/org.rs
index 1a3a144..f0e736b 100644
--- a/src/org.rs
+++ b/src/org.rs
@@ -1,5 +1,5 @@
 use rowan::ast::AstNode;
-use rowan::GreenNode;
+use rowan::{GreenNode, TextSize};
 
 use crate::ast::Document;
 use crate::config::ParseConfig;
@@ -66,4 +66,41 @@ impl Org {
         }
         find(SyntaxNode::new_root(self.green.clone()))
     }
+
+    /// Returns node in given offset
+    ///
+    /// ```rust
+    /// use orgize::{Org, ast::Headline};
+    ///
+    /// let org = Org::parse("\n\n* foo\n* bar");
+    ///
+    /// assert!(org.node_at_offset::<Headline>(0).is_none());
+    ///
+    /// let hdl = org.node_at_offset::<Headline>(2).unwrap();
+    /// assert_eq!(hdl.title_raw(), "foo");
+    ///
+    /// let hdl = org.node_at_offset::<Headline>(9).unwrap();
+    /// assert_eq!(hdl.title_raw(), "bar");
+    ///
+    /// assert!(org.node_at_offset::<Headline>(999).is_none());
+    /// ```
+    pub fn node_at_offset<N: AstNode<Language = OrgLanguage>>(
+        &self,
+        offset: impl Into<TextSize>,
+    ) -> Option<N> {
+        let offset = offset.into();
+        fn find<N: AstNode<Language = OrgLanguage>>(
+            node: SyntaxNode,
+            offset: TextSize,
+        ) -> Option<N> {
+            if !node.text_range().contains(offset) {
+                None
+            } else if N::can_cast(node.kind()) {
+                N::cast(node)
+            } else {
+                node.children().find_map(|node| find(node, offset))
+            }
+        }
+        find(SyntaxNode::new_root(self.green.clone()), offset)
+    }
 }