From b877bd43b1d96874785aef827d9c8bbb6132f43e Mon Sep 17 00:00:00 2001
From: tidwall <joshbaker77@gmail.com>
Date: Thu, 27 Jun 2019 17:51:42 -0700
Subject: [PATCH] Allow for chaining syntax in array subselects

---
 gjson.go      | 76 +++++++++++++++++++++++++++++++++++++++++++++++++++
 gjson_test.go | 56 ++++++++++++++++++++++++++++++-------
 2 files changed, 122 insertions(+), 10 deletions(-)

diff --git a/gjson.go b/gjson.go
index 0ac15fb..518edfa 100644
--- a/gjson.go
+++ b/gjson.go
@@ -891,6 +891,11 @@ func parseObjectPath(path string) (r objectPathResult) {
 						r.path = path[i+1:]
 						r.more = true
 						return
+					} else if path[i] == '|' {
+						r.part = string(epart)
+						r.pipe = path[i+1:]
+						r.piped = true
+						return
 					} else if path[i] == '*' || path[i] == '?' {
 						r.wild = true
 					}
@@ -1321,6 +1326,12 @@ func parseArray(c *parseContext, i int, path string) (int, bool) {
 			case ']':
 				if rp.arrch && rp.part == "#" {
 					if rp.alogok {
+						left, right, ok := splitPossiblePipe(rp.alogkey)
+						if ok {
+							rp.alogkey = left
+							c.pipe = right
+							c.piped = true
+						}
 						var jsons = make([]byte, 0, 64)
 						jsons = append(jsons, '[')
 						for j, k := 0, 0; j < len(alog); j++ {
@@ -1368,6 +1379,71 @@ func parseArray(c *parseContext, i int, path string) (int, bool) {
 	return i, false
 }
 
+func splitPossiblePipe(path string) (left, right string, ok bool) {
+	// take a quick peek for the pipe character. If found we'll split the piped
+	// part of the path into the c.pipe field and shorten the rp.
+	var possible bool
+	for i := 0; i < len(path); i++ {
+		if path[i] == '|' {
+			possible = true
+			break
+		}
+	}
+	if !possible {
+		return
+	}
+
+	// split the left and right side of the path with the pipe character as
+	// the delimiter. This is a little tricky because we'll need to basically
+	// parse the entire path.
+
+	for i := 0; i < len(path); i++ {
+		if path[i] == '\\' {
+			i++
+		} else if path[i] == '.' {
+			if i == len(path)-1 {
+				return
+			}
+			if path[i+1] == '#' {
+				i += 2
+				if i == len(path) {
+					return
+				}
+				if path[i] == '[' {
+					// inside selector, balance brackets
+					i++
+					depth := 1
+					for ; i < len(path); i++ {
+						if path[i] == '\\' {
+							i++
+						} else if path[i] == '[' {
+							depth++
+						} else if path[i] == ']' {
+							depth--
+							if depth == 0 {
+								break
+							}
+						} else if path[i] == '"' {
+							// inside selector string, balance quotes
+							i++
+							for ; i < len(path); i++ {
+								if path[i] == '\\' {
+									i++
+								} else if path[i] == '"' {
+									break
+								}
+							}
+						}
+					}
+				}
+			}
+		} else if path[i] == '|' {
+			return path[:i], path[i+1:], true
+		}
+	}
+	return
+}
+
 // ForEachLine iterates through lines of JSON as specified by the JSON Lines
 // format (http://jsonlines.org/).
 // Each line is returned as a GJSON Result.
diff --git a/gjson_test.go b/gjson_test.go
index 606375d..98bd525 100644
--- a/gjson_test.go
+++ b/gjson_test.go
@@ -1488,25 +1488,53 @@ func TestModifier(t *testing.T) {
 
 func TestChaining(t *testing.T) {
 	json := `{
-		"friends": [
-		  {"first": "Dale", "last": "Murphy", "age": 44},
-		  {"first": "Roger", "last": "Craig", "age": 68},
-		  {"first": "Jane", "last": "Murphy", "age": 47}
-		]
+		"info": {
+			"friends": [
+				{"first": "Dale", "last": "Murphy", "age": 44},
+				{"first": "Roger", "last": "Craig", "age": 68},
+				{"first": "Jane", "last": "Murphy", "age": 47}
+			]
+		}
 	  }`
-	res := Get(json, "friends|0|first").String()
+	res := Get(json, "info.friends|0|first").String()
 	if res != "Dale" {
 		t.Fatalf("expected '%v', got '%v'", "Dale", res)
 	}
-	res = Get(json, "friends|@reverse|0|age").String()
+	res = Get(json, "info.friends|@reverse|0|age").String()
 	if res != "47" {
 		t.Fatalf("expected '%v', got '%v'", "47", res)
 	}
+	res = Get(json, "@ugly|i\\nfo|friends.0.first").String()
+	if res != "Dale" {
+		t.Fatalf("expected '%v', got '%v'", "Dale", res)
+	}
+}
+
+func TestSplitPipe(t *testing.T) {
+	split := func(t *testing.T, path, el, er string, eo bool) {
+		t.Helper()
+		left, right, ok := splitPossiblePipe(path)
+		// fmt.Printf("%-40s [%v] [%v] [%v]\n", path, left, right, ok)
+		if left != el || right != er || ok != eo {
+			t.Fatalf("expected '%v/%v/%v', got '%v/%v/%v",
+				el, er, eo, left, right, ok)
+		}
+	}
+
+	split(t, "hello", "", "", false)
+	split(t, "hello.world", "", "", false)
+	split(t, "hello|world", "hello", "world", true)
+	split(t, "hello\\|world", "", "", false)
+	split(t, "hello.#", "", "", false)
+	split(t, `hello.#[a|1="asdf\"|1324"]#\|that`, "", "", false)
+	split(t, `hello.#[a|1="asdf\"|1324"]#|that.more|yikes`,
+		`hello.#[a|1="asdf\"|1324"]#`, "that.more|yikes", true)
+	split(t, `a.#[]#\|b`, "", "", false)
 
 }
 
 func TestArrayEx(t *testing.T) {
-	s := `
+	json := `
 	[
 		{
 			"c":[
@@ -1518,12 +1546,20 @@ func TestArrayEx(t *testing.T) {
 			]
 		}
 	]`
-	res := Get(s, "@ugly|#.c.#[a=10.11]").String()
+	res := Get(json, "@ugly|#.c.#[a=10.11]").String()
 	if res != `[{"a":10.11}]` {
 		t.Fatalf("expected '%v', got '%v'", `[{"a":10.11}]`, res)
 	}
-	res = Get(s, "@ugly|#.c.#").String()
+	res = Get(json, "@ugly|#.c.#").String()
 	if res != `[1,1]` {
 		t.Fatalf("expected '%v', got '%v'", `[1,1]`, res)
 	}
+	res = Get(json, "@reverse|0|c|0|a").String()
+	if res != "11.11" {
+		t.Fatalf("expected '%v', got '%v'", "11.11", res)
+	}
+	res = Get(json, "#.c|#").String()
+	if res != "2" {
+		t.Fatalf("expected '%v', got '%v'", "2", res)
+	}
 }