From afa8e96e82d2c4ce67f3e57b305a1c29e5846f6e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Koz=C5=82owski?= <kubukoz@gmail.com>
Date: Fri, 1 Nov 2024 01:58:20 +0100
Subject: [PATCH] Add visitors and folds

---
 .scalafmt.conf                                | 10 +++++--
 build.sbt                                     | 10 ++++---
 .../treesitter4s/TreeSitterAPI.scala          | 18 +++++++++++++
 .../TreeSitterTest.scala}                     |  7 ++---
 .../treesitter4s/tests/BindingTests.scala     | 26 +++++++++++++++++++
 5 files changed, 62 insertions(+), 9 deletions(-)
 rename core/src/test/scala/org/polyvariant/treesitter4s/{TreeSitterAPITest.scala => lowlevel/TreeSitterTest.scala} (86%)

diff --git a/.scalafmt.conf b/.scalafmt.conf
index b32145d..f224fde 100644
--- a/.scalafmt.conf
+++ b/.scalafmt.conf
@@ -1,5 +1,9 @@
-runner.dialect = "scala3"
-version = 3.5.8
+version = 3.8.3
+
+runner.dialect=scala3
+runner.dialectOverride.allowSignificantIndentation = false
+runner.dialectOverride.allowQuietSyntax = true
+
 maxColumn = 100
 align.preset = some
 
@@ -30,3 +34,5 @@ danglingParentheses.exclude = [
 ]
 verticalMultiline.newlineAfterOpenParen = true
 verticalMultiline.atDefnSite = true
+
+assumeStandardLibraryStripMargin = true
diff --git a/build.sbt b/build.sbt
index c0912e3..88b345b 100644
--- a/build.sbt
+++ b/build.sbt
@@ -51,10 +51,12 @@ val commonJVMSettings = Seq(
   Compile / doc / javacOptions -= "-Xlint:all",
   Test / fork := true,
   scalacOptions ++= {
-    if (scalaVersion.value.startsWith("2.13"))
-      Seq("-Wnonunit-statement")
-    else
-      Nil
+    Seq("-Wnonunit-statement") ++ {
+      if (scalaVersion.value.startsWith("3"))
+        Seq("-no-indent")
+      else
+        Nil
+    }
   },
 )
 
diff --git a/core/src/main/scala/org/polyvariant/treesitter4s/TreeSitterAPI.scala b/core/src/main/scala/org/polyvariant/treesitter4s/TreeSitterAPI.scala
index 3aac73f..5214299 100644
--- a/core/src/main/scala/org/polyvariant/treesitter4s/TreeSitterAPI.scala
+++ b/core/src/main/scala/org/polyvariant/treesitter4s/TreeSitterAPI.scala
@@ -49,4 +49,22 @@ trait Node {
   def fields: Map[String, Node]
   def startByte: Int
   def endByte: Int
+
+  def visit[A](visitor: Node.Visitor[A]): A = visitor.onNode(this, this.children)
+
+  // A specialized form of visitor, where every node is provided already visited.
+  def fold[A](folder: Node.Folder[A]): A = folder.onNode(this, this.children.map(_.fold(folder)))
+
+}
+
+object Node {
+
+  trait Folder[A] {
+    def onNode(node: Node, children: List[A]): A
+  }
+
+  trait Visitor[A] {
+    def onNode(node: Node, children: List[Node]): A
+  }
+
 }
diff --git a/core/src/test/scala/org/polyvariant/treesitter4s/TreeSitterAPITest.scala b/core/src/test/scala/org/polyvariant/treesitter4s/lowlevel/TreeSitterTest.scala
similarity index 86%
rename from core/src/test/scala/org/polyvariant/treesitter4s/TreeSitterAPITest.scala
rename to core/src/test/scala/org/polyvariant/treesitter4s/lowlevel/TreeSitterTest.scala
index 14be50b..97aa51a 100644
--- a/core/src/test/scala/org/polyvariant/treesitter4s/TreeSitterAPITest.scala
+++ b/core/src/test/scala/org/polyvariant/treesitter4s/lowlevel/TreeSitterTest.scala
@@ -14,14 +14,14 @@
  * limitations under the License.
  */
 
-package org.polyvariant.treesitter4s
+package org.polyvariant.treesitter4s.lowlevel
 
 import weaver.*
 
-object TreeSitterAPITest extends FunSuite {
+object TreeSitterTest extends FunSuite {
   test("Tree Sitter loads") {
     try {
-      println(TreeSitterAPI)
+      println(TreeSitter.instance)
       success
     } catch {
       case e: ExceptionInInitializerError =>
@@ -30,4 +30,5 @@ object TreeSitterAPITest extends FunSuite {
     }
 
   }
+
 }
diff --git a/tests/src/test/scala/org/polyvariant/treesitter4s/tests/BindingTests.scala b/tests/src/test/scala/org/polyvariant/treesitter4s/tests/BindingTests.scala
index 10b0834..40ff680 100644
--- a/tests/src/test/scala/org/polyvariant/treesitter4s/tests/BindingTests.scala
+++ b/tests/src/test/scala/org/polyvariant/treesitter4s/tests/BindingTests.scala
@@ -20,6 +20,7 @@ import cats.implicits._
 import org.polyvariant.treesitter4s.Tree
 import weaver._
 import org.polyvariant.treesitter4s.TreeSitterAPI
+import org.polyvariant.treesitter4s.Node
 
 object BindingTests extends FunSuite {
   val ts = TreeSitterAPI.make("python")
@@ -47,6 +48,31 @@ object BindingTests extends FunSuite {
     assert.eql(rootNode.children.lift(0).isDefined, true)
   }
 
+  test("Node.fold") {
+    val rootNode = parseExample("""
+                                  |def foo():
+                                  | return 1
+                                  |
+                                  |def bar():
+                                  |  def baz():
+                                  |    return 2
+                                  |  return baz()
+                                  |""".stripMargin)
+      .rootNode
+      .getOrElse(sys.error("missing root node"))
+
+    val functions = rootNode.fold[List[Node]] { (node, children) =>
+      if (node.tpe == "function_definition")
+        node :: children.flatten
+      else
+        children.flatten
+    }
+
+    val functionNames = functions.map(_.fields("name").source)
+
+    assert.eql(functionNames, "foo" :: "bar" :: "baz" :: Nil)
+  }
+
   // test("root node child by index (out of range)") {
   //   val tree = parseExample("class Hello {}")
   //   val rootNode = tree.rootNode.getOrElse(sys.error("missing root node"))