From 40080c2aafbb31b45446c0bd4cf56c6d3f72201a Mon Sep 17 00:00:00 2001
From: runcong zhang <39239990+ddzrc@users.noreply.github.com>
Date: Thu, 13 Jul 2023 16:14:57 +0800
Subject: [PATCH] feat:add map and array filter code

---
 internal/encoder/code.go  | 25 ++++++++++++++--
 internal/encoder/vm/vm.go |  3 ++
 query_test.go             | 62 +++++++++++++++++++++++++++++++++++++--
 3 files changed, 85 insertions(+), 5 deletions(-)

diff --git a/internal/encoder/code.go b/internal/encoder/code.go
index 5b08fae..e707ecd 100644
--- a/internal/encoder/code.go
+++ b/internal/encoder/code.go
@@ -271,7 +271,13 @@ func (c *SliceCode) ToOpcode(ctx *compileContext) Opcodes {
 	return Opcodes{header}.Add(codes...).Add(elemCode).Add(end)
 }
 
-func (c *SliceCode) Filter(_ *FieldQuery) Code {
+func (c *SliceCode) Filter(fieldQuery *FieldQuery) Code {
+	if len(fieldQuery.Fields) > 0 && fieldQuery.Fields[0].Name == "#" {
+		return &SliceCode{
+			value: c.value.Filter(fieldQuery.Fields[0]),
+			typ:   c.typ,
+		}
+	}
 	return c
 }
 
@@ -316,7 +322,13 @@ func (c *ArrayCode) ToOpcode(ctx *compileContext) Opcodes {
 	return Opcodes{header}.Add(codes...).Add(elemCode).Add(end)
 }
 
-func (c *ArrayCode) Filter(_ *FieldQuery) Code {
+func (c *ArrayCode) Filter(fieldQuery *FieldQuery) Code {
+	if len(fieldQuery.Fields) > 0 && fieldQuery.Fields[0].Name == "#" {
+		return &ArrayCode{
+			value: c.value.Filter(fieldQuery.Fields[0]),
+			typ:   c.typ,
+		}
+	}
 	return c
 }
 
@@ -366,7 +378,14 @@ func (c *MapCode) ToOpcode(ctx *compileContext) Opcodes {
 	return Opcodes{header}.Add(keyCodes...).Add(value).Add(valueCodes...).Add(key).Add(end)
 }
 
-func (c *MapCode) Filter(_ *FieldQuery) Code {
+func (c *MapCode) Filter(fieldQuery *FieldQuery) Code {
+	if len(fieldQuery.Fields) > 0 && fieldQuery.Fields[0].Name == "#" {
+		return &MapCode{
+			value: c.value.Filter(fieldQuery.Fields[0]),
+			typ:   c.typ,
+			key:   c.key,
+		}
+	}
 	return c
 }
 
diff --git a/internal/encoder/vm/vm.go b/internal/encoder/vm/vm.go
index 645d20f..7567b29 100644
--- a/internal/encoder/vm/vm.go
+++ b/internal/encoder/vm/vm.go
@@ -203,6 +203,9 @@ func Run(ctx *encoder.RuntimeContext, b []byte, codeSet *encoder.OpcodeSet) ([]b
 				}
 			}
 			ctx.KeepRefs = append(ctx.KeepRefs, up)
+			if code.FieldQuery != nil {
+				ctx.Option.Context = encoder.SetFieldQueryToContext(ctx.Option.Context, code.FieldQuery)
+			}
 			ifaceCodeSet, err := encoder.CompileToGetCodeSet(ctx, uintptr(unsafe.Pointer(typ)))
 			if err != nil {
 				return nil, err
diff --git a/query_test.go b/query_test.go
index d8cac61..24ee64b 100644
--- a/query_test.go
+++ b/query_test.go
@@ -28,6 +28,20 @@ type queryTestZ struct {
 	ZA string
 	ZB bool
 	ZC int
+	ZD []queryTestW
+	ZE map[string]queryTestV
+}
+
+type queryTestW struct {
+	WA string
+	WB bool
+	WC int
+}
+
+type queryTestV struct {
+	VA string
+	VB bool
+	VC int
 }
 
 func (z *queryTestZ) MarshalJSON(ctx context.Context) ([]byte, error) {
@@ -45,6 +59,10 @@ func TestFieldQuery(t *testing.T) {
 			json.BuildSubFieldQuery("YC").Fields(
 				"ZA",
 				"ZB",
+				json.BuildSubFieldQuery("ZD").Fields(json.BuildSubFieldQuery("#").Fields(
+					"WA", "WC")),
+				json.BuildSubFieldQuery("ZE").Fields(json.BuildSubFieldQuery("#").Fields(
+					"VA", "VC")),
 			),
 		),
 	)
@@ -77,6 +95,38 @@ func TestFieldQuery(t *testing.T) {
 							{
 								Name: "ZB",
 							},
+							{
+								Name: "ZD",
+								Fields: []*json.FieldQuery{
+									{
+										Name: "#",
+										Fields: []*json.FieldQuery{
+											{
+												Name: "WA",
+											},
+											{
+												Name: "WC",
+											},
+										},
+									},
+								},
+							},
+							{
+								Name: "ZE",
+								Fields: []*json.FieldQuery{
+									{
+										Name: "#",
+										Fields: []*json.FieldQuery{
+											{
+												Name: "VA",
+											},
+											{
+												Name: "VC",
+											},
+										},
+									},
+								},
+							},
 						},
 					},
 				},
@@ -89,7 +139,7 @@ func TestFieldQuery(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if queryStr != `["XA","XB",{"XC":["YA","YB",{"YC":["ZA","ZB"]}]}]` {
+	if queryStr != `["XA","XB",{"XC":["YA","YB",{"YC":["ZA","ZB",{"ZD":[{"#":["WA","WC"]}]},{"ZE":[{"#":["VA","VC"]}]}]}]}]` {
 		t.Fatalf("failed to create query string. %s", queryStr)
 	}
 	ctx := json.SetFieldQueryToContext(context.Background(), query)
@@ -103,6 +153,14 @@ func TestFieldQuery(t *testing.T) {
 				ZA: "za",
 				ZB: true,
 				ZC: 3,
+				ZD: []queryTestW{
+					{WA: "wa1", WB: true, WC: 1},
+					{WA: "wa2", WB: true, WC: 1},
+				},
+				ZE: map[string]queryTestV{
+					"key1": {VA: "va1", VB: true, VC: 1},
+					"key2": {VA: "va2", VB: true, VC: 1},
+				},
 			},
 			YD: true,
 			YE: 4,
@@ -113,7 +171,7 @@ func TestFieldQuery(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	expected := `{"XA":1,"XB":"xb","XC":{"YA":2,"YB":"yb","YC":{"ZA":"za","ZB":true}}}`
+	expected := `{"XA":1,"XB":"xb","XC":{"YA":2,"YB":"yb","YC":{"ZA":"za","ZB":true,"ZD":[{"WA":"wa1","WC":1},{"WA":"wa2","WC":1}],"ZE":{"key1":{"VA":"va1","VC":1},"key2":{"VA":"va2","VC":1}}}}}`
 	got := string(b)
 	if expected != got {
 		t.Fatalf("failed to encode with field query: expected %q but got %q", expected, got)