Support for 'inline def' inside 'query()' blocks

Tinkoff · Jun 28, 2022 · e72a21e · e72a21e
1 parent 26f3c2f
commit e72a21e
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 1 deletion.
diff --git a/oolong-core/src/main/scala/ru/tinkoff/oolong/AstParser.scala b/oolong-core/src/main/scala/ru/tinkoff/oolong/AstParser.scala
@@ -146,6 +146,9 @@ private[oolong] class DefaultAstParser(using quotes: Quotes) extends AstParser {
       case AsTerm(Literal(BooleanConstant(c))) =>
         makeConst(c)
 
+      case InlinedSubquery(term) =>
+        parse(term.asExpr)
+
       case _ =>
         report.errorAndAbort("Unexpected expr while parsing AST: " + input.show + s"; term: ${showTerm(input.asTerm)}")
     }

diff --git a/oolong-core/src/main/scala/ru/tinkoff/oolong/Utils.scala b/oolong-core/src/main/scala/ru/tinkoff/oolong/Utils.scala
@@ -37,6 +37,24 @@ private[oolong] object Utils {
     }
   }
 
+  object InlinedSubquery {
+    def unapply(using quotes: Quotes)(
+        term: quotes.reflect.Term
+    ): Option[quotes.reflect.Term] = {
+      import quotes.reflect.*
+      term match {
+        case Inlined(_, _, expansion)          => unapply(expansion)
+        case Typed(term, TypeIdent("Boolean")) => Some(term)
+        case _                                 => None
+      }
+    }
+
+    def unapply(expr: Expr[Any])(using quotes: Quotes): Option[quotes.reflect.Term] = {
+      import quotes.reflect.*
+      unapply(expr.asTerm)
+    }
+  }
+
   object AsTerm {
     def unapply(using quotes: Quotes)(
         expr: Expr[Any]
@@ -60,6 +78,11 @@ private[oolong] object Utils {
             loop(next, acc)
           case Ident(name) =>
             Some((name, acc))
+
+          // Ident() will be inlined if queries are composed via "inline def", e.g.
+          case Inlined(_, _, expansion) =>
+            loop(expansion, acc)
+
           case _ =>
             None
         }

diff --git a/oolong-mongo/src/test/scala/ru/tinkoff/oolong/mongo/QuerySpec.scala b/oolong-mongo/src/test/scala/ru/tinkoff/oolong/mongo/QuerySpec.scala
@@ -20,6 +20,10 @@ import ru.tinkoff.oolong.dsl.*
 
 class QuerySpec extends AnyFunSuite {
 
+  trait TestClassAncestor {
+    def intField: Int
+  }
+
   case class TestClass(
       intField: Int,
       stringField: String,
@@ -28,7 +32,7 @@ class QuerySpec extends AnyFunSuite {
       optionField: Option[Long],
       optionInnerClassField: Option[InnerClass],
       listField: List[Double]
-  )
+  ) extends TestClassAncestor
 
   case class InnerClass(
       fieldOne: String,
@@ -303,4 +307,74 @@ class QuerySpec extends AnyFunSuite {
     )
   }
 
+  inline def mySubquery1(doc: TestClass): Boolean = doc.intField == 123
+
+  test("calling an 'inline def' within a query() block #1") {
+
+    val q = query[TestClass](mySubquery1(_))
+
+    assert(
+      q == BsonDocument(
+        "intField" -> BsonDocument("$eq" -> BsonInt32(123))
+      )
+    )
+  }
+
+  test("calling an 'inline def' within a query() block #2") {
+
+    val q = query[TestClass](x => mySubquery1(x))
+
+    assert(
+      q == BsonDocument(
+        "intField" -> BsonDocument("$eq" -> BsonInt32(123))
+      )
+    )
+  }
+
+  test("calling a generic 'inline def' within a query() block") {
+
+    inline def genericSubquery[A <: TestClassAncestor](doc: A): Boolean = doc.intField == 123
+
+    val q = query[TestClass](x => genericSubquery(x))
+
+    assert(
+      q == BsonDocument(
+        "intField" -> BsonDocument("$eq" -> BsonInt32(123))
+      )
+    )
+  }
+
+  test("composing queries via 'inline def' #1") {
+
+    val q = query[TestClass](x => x.stringField == "qqq" && mySubquery1(x))
+
+    assert(
+      q == BsonDocument(
+        "$and" -> BsonArray.fromIterable(
+          List(
+            BsonDocument("stringField" -> BsonDocument("$eq" -> BsonString("qqq"))),
+            BsonDocument("intField"    -> BsonDocument("$eq" -> BsonInt32(123)))
+          )
+        )
+      )
+    )
+  }
+
+  test("composing queries via 'inline def' #2") {
+
+    inline def mySubquery2(tc: TestClass): Boolean = mySubquery1(tc) || tc.intField == 456
+
+    val q = query[TestClass](mySubquery2(_))
+
+    assert(
+      q == BsonDocument(
+        "$or" -> BsonArray.fromIterable(
+          List(
+            BsonDocument("intField" -> BsonDocument("$eq" -> BsonInt32(123))),
+            BsonDocument("intField" -> BsonDocument("$eq" -> BsonInt32(456)))
+          )
+        )
+      )
+    )
+  }
 }