Queries

tree-sitter-haskell/examples/QueryDemo.hs

@@ -0,0 +1,50 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module Main (main) where
+
+import           TreeSitter.Parser
+import           TreeSitter.Tree
+import           TreeSitter.Query
+import           TreeSitter.Haskell
+import           TreeSitter.Node
+import           Foreign.C.String
+import           Foreign.Ptr                    ( nullPtr)
+import           Foreign.Marshal.Alloc          ( malloc)
+import           Foreign.Storable               ( peek
+                                                , poke
+                                                )
+
+main :: IO ()
+main = do
+  parser <- ts_parser_new
+  _      <- ts_parser_set_language parser tree_sitter_haskell
+  let source = "module Test (main) where\nimport Lib\nf1 = undefined"
+
+  (str, len)  <- newCStringLen source
+  tree        <- ts_parser_parse_string parser nullPtr str len
+  nodePointer <- malloc
+  ts_tree_root_node_p tree nodePointer
+
+  let source = "(function_body) @f"
+  (str, len)  <- newCStringLen source
+  errorOffset <- malloc
+  errorType   <- malloc
+  cursor      <- ts_query_cursor_new
+  query       <- ts_query_new tree_sitter_haskell str len errorOffset errorType
+  ts_query_cursor_exec_p cursor query nodePointer
+  matchPointer <- malloc
+
+  success      <- ts_query_cursor_next_match cursor matchPointer
+  print $ "success: " ++ show success
+  match   <- (peek matchPointer)
+  capture <- peek $ matchCaptures match
+  print capture
+
+  -- print the match
+  ts_node <- malloc
+  poke ts_node (captureTSNode capture)
+  ts_node_string_p ts_node >>= peekCString >>= print
+
+  success <- ts_query_cursor_next_match cursor matchPointer
+  print $ "success: " ++ show success
+

tree-sitter-haskell/tree-sitter-haskell.cabal

@@ -38,6 +38,18 @@ flag build-examples
   description: Build tree-sitter-haskell examples.
   default:     False
 
+executable query-demo
+  import: common
+  main-is: QueryDemo.hs
+  if !flag(build-examples)
+    buildable: False
+  hs-source-dirs: examples
+  ghc-options: -threaded -rtsopts -with-rtsopts=-N
+  build-depends:
+      base
+    , tree-sitter
+    , tree-sitter-haskell
+
 executable demo
   import: common
   main-is: Demo.hs

tree-sitter/src/TreeSitter/Node.hs

@@ -4,17 +4,20 @@ module TreeSitter.Node
 ( Node(..)
 , nodeStartPoint
 , nodeStartByte
+, nodeString
 , TSPoint(..)
 , TSNode(..)
 , FieldId(..)
 , ts_node_copy_child_nodes
 , ts_node_poke_p
+, ts_node_string_p
 ) where
 
 import Foreign
 import Foreign.C
 import GHC.Generics
 import TreeSitter.Symbol (TSSymbol)
+import TreeSitter.Struct
 
 data Node = Node
   { nodeTSNode     :: !TSNode
@@ -29,6 +32,12 @@ data Node = Node
   }
   deriving (Show, Eq, Generic)
 
+nodeString :: Node -> IO String
+nodeString node = do
+  ts_node <- malloc
+  poke ts_node $ nodeTSNode node
+  ts_node_string_p ts_node >>= peekCString
+
 nodeStartPoint :: Node -> TSPoint
 nodeStartPoint node = let TSNode _ p _ _ _ = nodeTSNode node in p
 
@@ -44,29 +53,6 @@ data TSNode = TSNode !Word32 !TSPoint !Word32 !(Ptr ()) !(Ptr ())
 newtype FieldId = FieldId { getFieldId :: Word16 }
   deriving (Eq, Ord, Show, Storable)
 
-
--- | 'Struct' is a strict 'Monad' with automatic alignment & advancing, & inferred type.
-newtype Struct a = Struct { runStruct :: forall b . Ptr b -> IO (a, Ptr a) }
-
-evalStruct :: Struct a -> Ptr b -> IO a
-evalStruct s p = fmap fst $! runStruct s p
-{-# INLINE evalStruct #-}
-
-peekStruct :: forall a . Storable a => Struct a
-peekStruct = Struct (\ p -> do
-  let aligned = alignPtr (castPtr p) (alignment (undefined :: a))
-  a <- peek aligned
-  pure (a, aligned `plusPtr` sizeOf a))
-{-# INLINE peekStruct #-}
-
-pokeStruct :: Storable a => a -> Struct ()
-pokeStruct a = Struct (\ p -> do
-  let aligned = alignPtr (castPtr p) (alignment a)
-  poke aligned a
-  pure ((), castPtr aligned `plusPtr` sizeOf a))
-{-# INLINE pokeStruct #-}
-
-
 instance Storable Node where
   alignment _ = alignment (undefined :: TSNode)
   {-# INLINE alignment #-}
@@ -126,41 +112,7 @@ instance Storable TSNode where
     pokeStruct p2
   {-# INLINE poke #-}
 
-instance Functor Struct where
-  fmap f a = Struct go where
-    go p = do
-      (a', p') <- runStruct a p
-      let fa = f a'
-      fa `seq` p' `seq` pure (fa, castPtr p')
-    {-# INLINE go #-}
-  {-# INLINE fmap #-}
-
-instance Applicative Struct where
-  pure a = Struct (\ p -> pure (a, castPtr p))
-  {-# INLINE pure #-}
-
-  f <*> a = Struct go where
-    go p = do
-      (f', p')  <- runStruct f          p
-      (a', p'') <- p' `seq` runStruct a (castPtr p')
-      let fa = f' a'
-      fa `seq` p'' `seq` pure (fa, castPtr p'')
-    {-# INLINE go #-}
-  {-# INLINE (<*>) #-}
-
-instance Monad Struct where
-  return = pure
-  {-# INLINE return #-}
-
-  a >>= f = Struct go where
-    go p = do
-      (a', p')   <- runStruct a               p
-      (fa', p'') <- p' `seq` runStruct (f a') (castPtr p')
-      fa' `seq` p'' `seq` pure (fa', p'')
-    {-# INLINE go #-}
-  {-# INLINE (>>=) #-}
-
-
 foreign import ccall unsafe "src/bridge.c ts_node_copy_child_nodes" ts_node_copy_child_nodes :: Ptr TSNode -> Ptr Node -> IO ()
 -- NB: this leaves the field name as NULL.
 foreign import ccall unsafe "src/bridge.c ts_node_poke_p" ts_node_poke_p :: Ptr TSNode -> Ptr Node -> IO ()
+foreign import ccall unsafe "src/bridge.c ts_node_string_p" ts_node_string_p :: Ptr TSNode ->  IO CString

tree-sitter/src/TreeSitter/Query.hs

@@ -0,0 +1,104 @@
+{-# LANGUAGE DeriveGeneric #-}
+
+module TreeSitter.Query
+  ( TSQuery
+  , TSQueryCursor
+  , TSQueryError(..)
+  , captureTSNode
+  , captureIndex
+  , matchId
+  , matchPatternindex
+  , matchCaptureCount
+  , matchCaptures
+  , ts_query_new
+  , ts_query_cursor_exec_p
+  , ts_query_cursor_new
+  , ts_query_cursor_next_match
+  )
+where
+
+import           Foreign
+import           Foreign.C                      ( CString )
+import           GHC.Generics
+import           TreeSitter.Language
+import           TreeSitter.Node
+import           TreeSitter.Struct
+
+-- | A tree-sitter query.
+--
+--   This type is uninhabited and used only for type safety within 'Ptr' values.
+data TSQuery
+
+-- | A tree-sitter query cursor.
+--
+--   This type is uninhabited and used only for type safety within 'Ptr' values.
+data TSQueryCursor
+
+data TSQueryError
+  = TSQueryErrorNone
+  | TSQueryErrorSyntax
+  | TSQueryErrorNodeType
+  | TSQueryErrorField
+  | TSQueryErrorCapture
+  deriving (Show, Enum, Eq)
+
+data TSQueryCapture = TSQueryCapture
+  { captureTSNode :: !TSNode
+  , captureIndex :: !Word32
+  } deriving (Show, Eq, Generic)
+
+data TSQueryMatch = TSQueryMatch
+  { matchId :: !Word32
+  , matchPatternindex :: !Word16
+  , matchCaptureCount :: !Word16
+  , matchCaptures :: Ptr TSQueryCapture
+  } deriving (Show, Eq, Generic)
+
+instance Storable TSQueryError where
+  alignment _ = alignment (0 :: Int8)
+  {-# INLINE alignment #-}
+  sizeOf _ = 1
+  {-# INLINE sizeOf #-}
+  peek = evalStruct $ toEnum <$> peekStruct
+  {-# INLINE peek #-}
+  poke ptr e = evalStruct (pokeStruct $ fromEnum e) ptr 
+  {-# INLINE poke #-}
+
+instance Storable TSQueryCapture where
+  alignment _ = alignment (0 :: Int32)
+  {-# INLINE alignment #-}
+  sizeOf _ = 10
+  {-# INLINE sizeOf #-}
+  peek = evalStruct $ TSQueryCapture <$> peekStruct <*> peekStruct
+  {-# INLINE peek #-}
+  poke ptr (TSQueryCapture a b) = flip evalStruct ptr $ do
+    pokeStruct a
+    pokeStruct b
+  {-# INLINE poke #-}
+
+instance Storable TSQueryMatch where
+  alignment _ = alignment (0 :: Int32)
+  {-# INLINE alignment #-}
+  sizeOf _ = 12
+  {-# INLINE sizeOf #-}
+  peek =
+    evalStruct
+      $   TSQueryMatch
+      <$> peekStruct
+      <*> peekStruct
+      <*> peekStruct
+      <*> peekStruct
+  {-# INLINE peek #-}
+  poke ptr (TSQueryMatch a b c d) = flip evalStruct ptr $ do
+    pokeStruct a
+    pokeStruct b
+    pokeStruct c
+    pokeStruct d
+  {-# INLINE poke #-}
+
+
+foreign import ccall safe "ts_query_new" ts_query_new :: Ptr Language -> CString -> Int -> Ptr Word32 -> Ptr TSQueryError -> IO (Ptr TSQuery)
+foreign import ccall safe "ts_query_cursor_new" ts_query_cursor_new :: IO (Ptr TSQueryCursor)
+foreign import ccall safe "src/bridge.c ts_query_cursor_exec_p" ts_query_cursor_exec_p :: Ptr TSQueryCursor -> Ptr TSQuery -> Ptr Node -> IO ()
+foreign import ccall safe "ts_query_cursor_next_match" ts_query_cursor_next_match :: Ptr TSQueryCursor -> Ptr TSQueryMatch -> IO Bool
+

tree-sitter/src/TreeSitter/Struct.hs

@@ -0,0 +1,64 @@
+{-# LANGUAGE DeriveGeneric, GeneralizedNewtypeDeriving, RankNTypes, ScopedTypeVariables,TypeOperators , DeriveAnyClass#-}
+{-# OPTIONS_GHC -funbox-strict-fields #-}
+
+module TreeSitter.Struct (Struct, evalStruct, peekStruct, pokeStruct) where
+
+import Foreign
+
+-- | 'Struct' is a strict 'Monad' with automatic alignment & advancing, & inferred type.
+newtype Struct a = Struct { runStruct :: forall b . Ptr b -> IO (a, Ptr a) }
+
+instance Functor Struct where
+  fmap f a = Struct go where
+    go p = do
+      (a', p') <- runStruct a p
+      let fa = f a'
+      fa `seq` p' `seq` pure (fa, castPtr p')
+    {-# INLINE go #-}
+  {-# INLINE fmap #-}
+
+instance Applicative Struct where
+  pure a = Struct (\ p -> pure (a, castPtr p))
+  {-# INLINE pure #-}
+
+  f <*> a = Struct go where
+    go p = do
+      (f', p')  <- runStruct f          p
+      (a', p'') <- p' `seq` runStruct a (castPtr p')
+      let fa = f' a'
+      fa `seq` p'' `seq` pure (fa, castPtr p'')
+    {-# INLINE go #-}
+  {-# INLINE (<*>) #-}
+
+instance Monad Struct where
+  return = pure
+  {-# INLINE return #-}
+
+  a >>= f = Struct go where
+    go p = do
+      (a', p')   <- runStruct a               p
+      (fa', p'') <- p' `seq` runStruct (f a') (castPtr p')
+      fa' `seq` p'' `seq` pure (fa', p'')
+    {-# INLINE go #-}
+  {-# INLINE (>>=) #-}
+
+evalStruct :: Struct a -> Ptr b -> IO a
+evalStruct s p = fmap fst $! runStruct s p
+{-# INLINE evalStruct #-}
+
+peekStruct :: forall a . Storable a => Struct a
+peekStruct = Struct
+  ( \p -> do
+    let aligned = alignPtr (castPtr p) (alignment (undefined :: a))
+    a <- peek aligned
+    pure (a, aligned `plusPtr` sizeOf a)
+  )
+{-# INLINE peekStruct #-}
+
+pokeStruct :: Storable a => a -> Struct ()
+pokeStruct a = Struct
+  ( \p -> do
+    let aligned = alignPtr (castPtr p) (alignment a)
+    poke aligned a
+    pure ((), castPtr aligned `plusPtr` sizeOf a)
+  )

tree-sitter/src/bridge.c

@@ -23,6 +23,18 @@ void ts_parser_log_to_stderr(TSParser *parser) {
   ts_parser_set_logger(parser, (TSLogger) {.log = log_to_stdout, .payload = NULL});
 }
 
+char* ts_node_string_p(TSNode* node) {
+  assert(node != NULL);
+  return ts_node_string(*node);
+}
+
+void ts_query_cursor_exec_p(TSQueryCursor* cursor, const TSQuery* query, Node* node) {
+  assert(cursor != NULL);
+  assert(query != NULL);
+  assert(node != NULL);
+  ts_query_cursor_exec(cursor, query, node->node);
+}
+
 static inline void ts_node_poke(const char *fieldName, TSNode node, Node *out) {
   out->node = node;
   out->symbol = ts_node_symbol(node);

tree-sitter/tree-sitter.cabal

@@ -65,6 +65,8 @@ library
                      , TreeSitter.Symbol
                      , TreeSitter.Tree
                      , TreeSitter.Cursor
+                     , TreeSitter.Query
+  other-modules:       TreeSitter.Struct
   include-dirs:        vendor/tree-sitter/lib/include
                      , vendor/tree-sitter/lib/src
   install-includes:    tree_sitter/api.h