From 930cf1c7b17c39f7d5ff871548a4232473559dc0 Mon Sep 17 00:00:00 2001
From: Bernhard <bernhard@shiftleft.io>
Date: Mon, 25 Sep 2023 16:41:02 +0200
Subject: [PATCH] add codegen for nodestarters with integrated filters

---
 .../src/main/scala/overflowdb/codegen/CodeGen.scala   | 11 +++++++++++
 codegen/src/main/scala/overflowdb/schema/Schema.scala |  9 +++++++++
 .../schemas/src/main/scala/TestSchema01.scala         |  2 +-
 .../tests/src/test/scala/Schema01Test.scala           | 11 +++++++++++
 4 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/codegen/src/main/scala/overflowdb/codegen/CodeGen.scala b/codegen/src/main/scala/overflowdb/codegen/CodeGen.scala
index bce1e31f..292bc7dd 100644
--- a/codegen/src/main/scala/overflowdb/codegen/CodeGen.scala
+++ b/codegen/src/main/scala/overflowdb/codegen/CodeGen.scala
@@ -85,6 +85,17 @@ class CodeGen(schema: Schema) {
       starters.append(
         s"""@overflowdb.traversal.help.Doc(info = "All nodes of type ${typ.className}, i.e. with label ${typ.name}")
            |def ${typ.starterName.get}: Iterator[nodes.${typ.className}] = overflowdb.traversal.InitialTraversal.from[nodes.${typ.className}](wrapper.graph, "${typ.name}")""".stripMargin)
+      typ.primaryKey match {
+        case Some(property:Property[_]) =>
+          //fixme
+          val nameCamelCase = camelCase(property.name)
+          val baseType = typeFor(property)
+          starters.append(
+            s"""@overflowdb.traversal.help.Doc(info = "Shorthand for ${typ.starterName.get}.${nameCamelCase}")
+               |def ${typ.starterName.get}(key: ${baseType}): Iterator[nodes.${typ.className}] =
+               |  new  ${traversalsPackage}.${typ.className}TraversalExtGen(${typ.starterName.get}).${nameCamelCase}(key)""".stripMargin)
+        case _ =>
+      }
     }
     for (typ <- schema.nodeBaseTypes if typ.starterName.isDefined) {
       val types = schema.nodeTypes.filter{_.extendzRecursively.contains(typ)}
diff --git a/codegen/src/main/scala/overflowdb/schema/Schema.scala b/codegen/src/main/scala/overflowdb/schema/Schema.scala
index 833acea9..78012a0a 100644
--- a/codegen/src/main/scala/overflowdb/schema/Schema.scala
+++ b/codegen/src/main/scala/overflowdb/schema/Schema.scala
@@ -139,6 +139,15 @@ class NodeType(name: String, comment: Option[String], schemaInfo: SchemaInfo)
   extends AbstractNodeType(name, comment, schemaInfo) with HasOptionalProtoId  {
   protected val _containedNodes: mutable.Set[ContainedNode] = mutable.Set.empty
 
+  private var _primaryKey: Option[Property[_]] = None
+
+  def primaryKey(p: Property[_]): this.type = {
+    this._primaryKey = Some(p)
+    this
+  }
+
+  def primaryKey: Option[Property[_]] = this._primaryKey
+
   lazy val classNameDb = s"${className}Db"
 
   /** all node types that extend this node */
diff --git a/integration-tests/schemas/src/main/scala/TestSchema01.scala b/integration-tests/schemas/src/main/scala/TestSchema01.scala
index e768b37b..f4f8411a 100644
--- a/integration-tests/schemas/src/main/scala/TestSchema01.scala
+++ b/integration-tests/schemas/src/main/scala/TestSchema01.scala
@@ -13,7 +13,7 @@ class TestSchema01 extends TestSchema {
   val node1 = builder.addNodeType(
     name = "NODE1",
     comment = "sample node 1"
-  ).addProperties(name, order)
+  ).addProperties(name, order).primaryKey(name)
 
   val node2 = builder.addNodeType(
     name = "NODE2",
diff --git a/integration-tests/tests/src/test/scala/Schema01Test.scala b/integration-tests/tests/src/test/scala/Schema01Test.scala
index 58b13973..d67668c3 100644
--- a/integration-tests/tests/src/test/scala/Schema01Test.scala
+++ b/integration-tests/tests/src/test/scala/Schema01Test.scala
@@ -81,6 +81,17 @@ class Schema01Test extends AnyWordSpec with Matchers {
       // TODO generate domain-specific setters in codegen
     }
 
+    "use generated node starters" in {
+      val domainspecific = new testschema01.TestSchema(graph)
+      object TmpImplicit {
+        implicit def toStarter(wrapped: testschema01.TestSchema): testschema01.GeneratedNodeStarterExt = new testschema01.GeneratedNodeStarterExt(wrapped)
+      }
+      import TmpImplicit._
+      domainspecific.node1.toList shouldBe List(node1a, node1b)
+      domainspecific.node1(".*1b").toList shouldBe List(node1b)
+
+    }
+
     "generate NewNodes" in {
       val newNode2 = NewNode2()
         .name("name1")