From c1eeb6939f8b8a2000cb9909b4e73212f33f73be Mon Sep 17 00:00:00 2001
From: James Prior <jamesgr.prior@gmail.com>
Date: Sun, 17 Sep 2023 13:33:50 +0100
Subject: [PATCH] More well-typedness test cases

---
 src/path/environment.ts                       |  9 ++-
 src/path/lex.ts                               |  4 ++
 src/path/parse.ts                             |  6 --
 tests/path/ietf_compare.test.ts               |  2 +-
 .../path/ietf_function_well_typedness.test.ts | 64 +++++++++++++++++--
 5 files changed, 72 insertions(+), 13 deletions(-)

diff --git a/src/path/environment.ts b/src/path/environment.ts
index 00b9a5c..ca2835b 100644
--- a/src/path/environment.ts
+++ b/src/path/environment.ts
@@ -160,7 +160,14 @@ export class JSONPathEnvironment {
           }
           break;
         case FunctionExpressionType.NodesType:
-          if (!(arg instanceof JSONPathQuery)) {
+          if (
+            !(
+              arg instanceof JSONPathQuery ||
+              (arg instanceof FunctionExtension &&
+                this.functionRegister.get(arg.name)?.returnType ===
+                  FunctionExpressionType.NodesType)
+            )
+          ) {
             throw new JSONPathTypeError(
               `${token.value}() argument ${idx} must be of NodesType`,
               arg.token,
diff --git a/src/path/lex.ts b/src/path/lex.ts
index d9590f9..96e045b 100644
--- a/src/path/lex.ts
+++ b/src/path/lex.ts
@@ -341,6 +341,10 @@ function lexInsideFilter(l: Lexer): StateFn | null {
       case "":
       case "]":
         l.filterLevel -= 1;
+        if (l.parenStack.length === 1) {
+          l.error("unbalanced parentheses");
+          return null;
+        }
         l.backup();
         return lexInsideBracketedSelection;
       case ",":
diff --git a/src/path/parse.ts b/src/path/parse.ts
index 6df4910..0f433a8 100644
--- a/src/path/parse.ts
+++ b/src/path/parse.ts
@@ -382,12 +382,6 @@ export class Parser {
 
       args.push(func.bind(this)(stream));
 
-      // Without this, that closing parenthesis of the inner function
-      // can incorrectly close the outer function.
-      if (args.at(-1) instanceof FunctionExtension) {
-        stream.next();
-      }
-
       if (stream.peek.kind !== TokenKind.RPAREN) {
         if (stream.peek.kind === TokenKind.RBRACKET) break;
         stream.expectPeek(TokenKind.COMMA);
diff --git a/tests/path/ietf_compare.test.ts b/tests/path/ietf_compare.test.ts
index b80f409..10e8d71 100644
--- a/tests/path/ietf_compare.test.ts
+++ b/tests/path/ietf_compare.test.ts
@@ -1,5 +1,5 @@
 /**
- * Filter expression comparison examples.
+ * Filter expression comparison tests from IETF spec examples.
  *
  * The test cases defined here are taken from version 20 of the JSONPath
  * internet draft, draft-ietf-jsonpath-base-20. In accordance with
diff --git a/tests/path/ietf_function_well_typedness.test.ts b/tests/path/ietf_function_well_typedness.test.ts
index abc3361..46bdc1f 100644
--- a/tests/path/ietf_function_well_typedness.test.ts
+++ b/tests/path/ietf_function_well_typedness.test.ts
@@ -1,5 +1,5 @@
 /**
- * Function well-typedness examples.
+ * Function well-typedness tests from IETF spec examples.
  *
  * The test cases defined here are taken from version 20 of the JSONPath
  * internet draft, draft-ietf-jsonpath-base-20. In accordance with
@@ -37,7 +37,12 @@
  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
-import { JSONPathTypeError, compile } from "../../src";
+import {
+  FilterFunction,
+  FunctionExpressionType,
+  JSONPathEnvironment,
+  JSONPathTypeError,
+} from "../../src";
 
 type TestCase = {
   description: string;
@@ -66,18 +71,67 @@ const TEST_CASES: TestCase[] = [
     path: "$[?count(1) == 1]",
     valid: false,
   },
+  {
+    description: "nested function, LogicalType -> NodesType",
+    path: "$[?count(foo(@.*)) == 1]",
+    valid: true,
+  },
+  {
+    description: "match, singular query, string literal",
+    path: "$[?match(@.timezone, 'Europe/.*')]",
+    valid: true,
+  },
+  {
+    description: "match, singular query, string literal, compared",
+    path: "$[?match(@.timezone, 'Europe/.*') == true]",
+    valid: false,
+  },
+  {
+    description: "value, comparison",
+    path: "$[?value(@..color) == 'red']",
+    valid: true,
+  },
+  {
+    description: "value, existence",
+    path: "$[?value(@..color)]",
+    valid: false,
+  },
+  {
+    description: "function, logical return type, existence",
+    path: "$[?bar(@.a)]",
+    valid: true,
+  },
 ];
 
-// TODO: tests with mock functions
+class MockFooFunction implements FilterFunction {
+  argTypes = [FunctionExpressionType.NodesType];
+  returnType = FunctionExpressionType.NodesType;
+
+  call(nodes: NodeList): NodeList {
+    return nodes;
+  }
+}
+
+class MockBarFunction implements FilterFunction {
+  argTypes = [FunctionExpressionType.NodesType];
+  returnType = FunctionExpressionType.LogicalType;
+
+  call(): boolean {
+    return false;
+  }
+}
 
 describe("IETF function well-typedness examples", () => {
+  const env = new JSONPathEnvironment();
+  env.functionRegister.set("foo", new MockFooFunction());
+  env.functionRegister.set("bar", new MockBarFunction());
   test.each<TestCase>(TEST_CASES)(
     "$description",
     ({ path, valid }: TestCase) => {
       if (!valid) {
-        expect(() => compile(path)).toThrow(JSONPathTypeError);
+        expect(() => env.compile(path)).toThrow(JSONPathTypeError);
       } else {
-        compile(path);
+        env.compile(path);
       }
     },
   );