From e67e8dcba439c09816e9eddcf2bbf77da30c09b8 Mon Sep 17 00:00:00 2001
From: lonerapier <lonerapier@proton.me>
Date: Thu, 22 Aug 2024 00:38:58 +0530
Subject: [PATCH] add nested structure example circuit

---
 circuits/fetcher.circom                    | 178 +++++++++++++++++++--
 circuits/test/fetcher.test.ts              |  29 +++-
 json_examples/test/value_array_object.json |   2 +-
 3 files changed, 189 insertions(+), 20 deletions(-)

diff --git a/circuits/fetcher.circom b/circuits/fetcher.circom
index d898614..a3947a2 100644
--- a/circuits/fetcher.circom
+++ b/circuits/fetcher.circom
@@ -5,6 +5,7 @@ include "parser.circom";
 include "language.circom";
 include "utils.circom";
 include "circomlib/circuits/mux1.circom";
+include "circomlib/circuits/gates.circom";
 include "@zk-email/circuits/utils/functions.circom";
 include "@zk-email/circuits/utils/array.circom";
 
@@ -46,6 +47,18 @@ template InsideValue(n) {
     out <== if_parsing_value * parsing_string_or_number;
 }
 
+template InsideObjectAtDepth(n, depth) {
+    signal input stack[n][2];
+    signal input parsing_string;
+    signal input parsing_number;
+
+    signal output out;
+
+    signal if_parsing_value <== stack[depth][0] * stack[depth][1];
+
+    out <== if_parsing_value * (parsing_string + parsing_number);
+}
+
 template InsideArrayIndex(n, index) {
     signal input stack[n][2];
     signal input parsing_string;
@@ -70,16 +83,9 @@ template InsideArrayIndexAtDepth(n, index, depth) {
 
     signal output out;
 
-    // component topOfStack = GetTopOfStack(n);
-    // topOfStack.stack <== stack;
-    // signal current_val[2] <== topOfStack.value;
-    // signal pointer <== topOfStack.pointer;
-
     signal inside_array <== IsEqual()([stack[depth][0], 2]);
     signal inside_index <== IsEqual()([stack[depth][1], index]);
     signal inside_array_index <== inside_array * inside_index;
-    // signal at_depth <== IsEqual()([pointer, depth]);
-    // signal inside_array_index_at_depth <== inside_array_index * at_depth;
     out <== inside_array_index * (parsing_string + parsing_number);
 }
 
@@ -176,8 +182,9 @@ template KeyMatchAtDepth(dataLen, n, keyLen, depth) {
     signal substring_match <== IsSubstringMatchWithIndex(dataLen, keyLen)(data, key, 100, index);
 
     signal is_key_between_quotes <== is_start_of_key_equal_to_quote * is_end_of_key_equal_to_quote;
+    log("key pointer", pointer, depth);
     signal is_parsing_correct_key <== is_key_between_quotes * parsing_key;
-    signal is_key_at_depth <== IsEqual()([pointer, depth]);
+    signal is_key_at_depth <== IsEqual()([pointer-1, depth]);
     signal is_parsing_correct_key_at_depth <== is_parsing_correct_key * is_key_at_depth;
     // log("key match", index, end_of_key, is_end_of_key_equal_to_quote, substring_match);
 
@@ -228,13 +235,13 @@ template ExtractValue(DATA_BYTES, MAX_STACK_HEIGHT, keyLen, maxValueLen) {
         State[data_idx].parsing_number <== State[data_idx - 1].next_parsing_number;
 
         parsing_key[data_idx-1] <== InsideKey(MAX_STACK_HEIGHT)(State[data_idx].stack, State[data_idx].parsing_string, State[data_idx].parsing_number);
-        // log("parsing key:", parsing_key[data_idx]);
+        log("parsing key:", parsing_key[data_idx-1]);
 
         parsing_value[data_idx-1] <== InsideValue(MAX_STACK_HEIGHT)(State[data_idx].stack, State[data_idx].parsing_string, State[data_idx].parsing_number);
-        // log("parsing value:", parsing_value[data_idx]);
+        log("parsing value:", parsing_value[data_idx-1]);
 
         is_key_match[data_idx-1] <== KeyMatch(DATA_BYTES, keyLen)(data, key, 100, data_idx-1, parsing_key[data_idx-1]);
-        // log("is_key_match", is_key_match[data_idx]);
+        log("is_key_match", is_key_match[data_idx-1]);
 
         // is the value getting parsed has a matched key?
         // use mux1 to carry parse_key forward to value
@@ -246,7 +253,7 @@ template ExtractValue(DATA_BYTES, MAX_STACK_HEIGHT, keyLen, maxValueLen) {
         // log("is_new_kv_pair:", is_next_pair[data_idx-1]);
 
         is_key_match_for_value[data_idx] <== Mux1()([is_key_match_for_value[data_idx-1] * (1-is_next_pair[data_idx-1]), is_key_match[data_idx-1] * (1-is_next_pair[data_idx-1])], is_key_match[data_idx-1]);
-        // log("is_key_match_for_value:", is_key_match_for_value[data_idx]);
+        log("is_key_match_for_value:", is_key_match_for_value[data_idx]);
 
         // mask[i] = data[i] * parsing_value[i] * is_key_match_for_value[i]
         value_mask[data_idx-1] <== data[data_idx-1] * parsing_value[data_idx-1];
@@ -556,6 +563,8 @@ template ExtractStringMultiDepth(DATA_BYTES, MAX_STACK_HEIGHT, keyLen1, depth1,
 
     value_starting_index <== ExtractMultiDepth(DATA_BYTES, MAX_STACK_HEIGHT, keyLen1, depth1, keyLen2, depth2, maxValueLen)(data, key1, key2);
 
+    log(value_starting_index[DATA_BYTES-2]);
+
     value <== SelectSubArray(DATA_BYTES, maxValueLen)(data, value_starting_index[DATA_BYTES-2]+1, maxValueLen);
 
     for (var i=0 ; i<maxValueLen; i++) {
@@ -678,4 +687,149 @@ template ExtractNestedArray(DATA_BYTES, MAX_STACK_HEIGHT, keyLen, index1, depth1
     // }
 
     // value <== number_value[maxValueLen-1];
+}
+
+template ExtractMultiDepthNestedObject(DATA_BYTES, MAX_STACK_HEIGHT, keyLen1, depth1, keyLen2, depth2, index3, depth3, index4, depth4, maxValueLen) {
+    signal input data[DATA_BYTES];
+
+    signal input key1[keyLen1];
+    signal input key2[keyLen2];
+
+    signal output value_starting_index[DATA_BYTES];
+
+    signal mask[DATA_BYTES];
+    // mask[0] <== 0;
+
+    var logDataLen = log2Ceil(DATA_BYTES);
+
+    component State[DATA_BYTES];
+    State[0] = StateUpdate(MAX_STACK_HEIGHT);
+    State[0].byte           <== data[0];
+    for(var i = 0; i < MAX_STACK_HEIGHT; i++) {
+        State[0].stack[i]   <== [0,0];
+    }
+    State[0].parsing_string <== 0;
+    State[0].parsing_number <== 0;
+
+    signal parsing_key[DATA_BYTES];
+    signal parsing_value[DATA_BYTES];
+    signal parsing_object1_value[DATA_BYTES];
+    signal parsing_object2_value[DATA_BYTES];
+    signal parsing_array1[DATA_BYTES];
+    signal parsing_array2[DATA_BYTES];
+    signal is_key1_match[DATA_BYTES];
+    signal is_key2_match[DATA_BYTES];
+    signal is_key1_match_for_value[DATA_BYTES];
+    is_key1_match_for_value[0] <== 0; // TODO: this might not be correct way to initialise
+    signal is_key2_match_for_value[DATA_BYTES];
+    is_key2_match_for_value[0] <== 0; // TODO: this might not be correct way to initialise
+    signal is_value_match[DATA_BYTES];
+    is_value_match[0] <== 0;
+    signal value_mask[DATA_BYTES];
+    signal is_next_pair_at_depth1[DATA_BYTES];
+    signal is_next_pair_at_depth2[DATA_BYTES];
+    for(var data_idx = 1; data_idx < DATA_BYTES; data_idx++) {
+        // Debugging
+        for(var i = 0; i<MAX_STACK_HEIGHT; i++) {
+            log("State[", data_idx-1, "].stack[", i,"]    ", "= [",State[data_idx-1].next_stack[i][0], "][", State[data_idx-1].next_stack[i][1],"]" );
+        }
+        log("State[", data_idx-1, "].byte", "= ", data[data_idx-1]);
+        log("State[", data_idx-1, "].parsing_string", "= ", State[data_idx-1].next_parsing_string);
+        log("State[", data_idx-1, "].parsing_number", "= ", State[data_idx-1].next_parsing_number);
+
+        State[data_idx]                  = StateUpdate(MAX_STACK_HEIGHT);
+        State[data_idx].byte           <== data[data_idx];
+        State[data_idx].stack          <== State[data_idx - 1].next_stack;
+        State[data_idx].parsing_string <== State[data_idx - 1].next_parsing_string;
+        State[data_idx].parsing_number <== State[data_idx - 1].next_parsing_number;
+
+        // - parsing key
+        // - parsing value (different for string/numbers and array)
+        // - key match (key 1, key 2)
+        // - is next pair
+        // - is key match for value
+        // - value_mask
+        // - mask
+
+        parsing_key[data_idx-1] <== InsideKey(MAX_STACK_HEIGHT)(State[data_idx].stack, State[data_idx].parsing_string, State[data_idx].parsing_number);
+        // log("parsing key:", parsing_key[data_idx]);
+
+        parsing_array1[data_idx-1] <== InsideArrayIndexAtDepth(MAX_STACK_HEIGHT, index3, depth3)(State[data_idx].stack, State[data_idx].parsing_string, State[data_idx].parsing_number);
+        parsing_array2[data_idx-1] <== InsideArrayIndexAtDepth(MAX_STACK_HEIGHT, index4, depth4)(State[data_idx].stack, State[data_idx].parsing_string, State[data_idx].parsing_number);
+        // parsing_value[data_idx-1] <== parsing_array1[data_idx-1] * parsing_array2[data_idx-1];
+
+        parsing_object1_value[data_idx-1] <== InsideObjectAtDepth(MAX_STACK_HEIGHT, depth1)(State[data_idx].stack, State[data_idx].parsing_string, State[data_idx].parsing_number);
+        parsing_object2_value[data_idx-1] <== InsideObjectAtDepth(MAX_STACK_HEIGHT, depth2)(State[data_idx].stack, State[data_idx].parsing_string, State[data_idx].parsing_number);
+        parsing_value[data_idx-1] <== MultiAND(4)([parsing_array1[data_idx-1], parsing_array2[data_idx-1], parsing_object1_value[data_idx-1], parsing_object2_value[data_idx-1]]);
+        log("parsing value:", parsing_array1[data_idx-1], parsing_array2[data_idx-1], parsing_object1_value[data_idx-1], parsing_object2_value[data_idx-1], parsing_value[data_idx-1]);
+
+        is_key1_match[data_idx-1] <== KeyMatchAtDepth(DATA_BYTES, MAX_STACK_HEIGHT, keyLen1, depth1)(data, key1, 100, data_idx-1, parsing_key[data_idx-1], State[data_idx].stack);
+        is_key2_match[data_idx-1] <== KeyMatchAtDepth(DATA_BYTES, MAX_STACK_HEIGHT, keyLen2, depth2)(data, key2, 100, data_idx-1, parsing_key[data_idx-1], State[data_idx].stack);
+        log("is_key_match", is_key1_match[data_idx-1], is_key2_match[data_idx-1]);
+
+        // is_next_pair represents if we are currently parsing kv pair of depth greater than key's depth
+        // eg: `{ "a": { "d" : "e", "e": "c" }, "e": { "f": "a", "e": "2" } }`
+        is_next_pair_at_depth1[data_idx-1] <== NextKVPairAtDepth(MAX_STACK_HEIGHT, depth1)(State[data_idx].stack, data[data_idx-1]);
+        is_next_pair_at_depth2[data_idx-1] <== NextKVPairAtDepth(MAX_STACK_HEIGHT, depth2)(State[data_idx].stack, data[data_idx-1]);
+        log("is_new_kv_pair:", is_next_pair_at_depth1[data_idx-1], is_next_pair_at_depth2[data_idx-1]);
+
+        // is the value getting parsed has a matched key?
+        // use mux1 to carry parse_key forward to value
+        // is_key_match_for_value should reset when moving to next kv pair
+        // `is_key_match = 0` -> 0
+        // `is_key_match = 1` -> 1 until new kv pair
+        // `new kv pair = 1`  -> 0
+        // all the keys should match for the correct value
+        is_key1_match_for_value[data_idx] <== Mux1()([is_key1_match_for_value[data_idx-1] * (1-is_next_pair_at_depth1[data_idx-1]), is_key1_match[data_idx-1] * (1-is_next_pair_at_depth1[data_idx-1])], is_key1_match[data_idx-1]);
+        is_key2_match_for_value[data_idx] <== Mux1()([is_key2_match_for_value[data_idx-1] * (1-is_next_pair_at_depth2[data_idx-1]), is_key2_match[data_idx-1] * (1-is_next_pair_at_depth2[data_idx-1])], is_key2_match[data_idx-1]);
+        log("is_key_match_for_value:", is_key1_match_for_value[data_idx], is_key2_match_for_value[data_idx]);
+
+        is_value_match[data_idx] <== is_key1_match_for_value[data_idx] * is_key2_match_for_value[data_idx];
+        // log("is_value_match", is_value_match[data_idx]);
+
+        // mask[i] = data[i] * parsing_value[i] * is_key_match_for_value[i]
+        value_mask[data_idx-1] <== data[data_idx-1] * parsing_value[data_idx-1];
+        mask[data_idx-1] <== value_mask[data_idx-1] * is_value_match[data_idx];
+        log("mask", mask[data_idx-1]);
+        log("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx");
+    }
+
+    // Debugging
+    for(var i = 0; i < MAX_STACK_HEIGHT; i++) {
+        log("State[", DATA_BYTES-1, "].stack[", i,"]    ", "= [",State[DATA_BYTES -1].next_stack[i][0], "][", State[DATA_BYTES - 1].next_stack[i][1],"]" );
+    }
+    log("State[", DATA_BYTES-1, "].parsing_string", "= ", State[DATA_BYTES-1].next_parsing_string);
+    log("State[", DATA_BYTES-1, "].parsing_number", "= ", State[DATA_BYTES-1].next_parsing_number);
+    log("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx");
+
+    // signal value_starting_index[DATA_BYTES];
+    signal is_zero_mask[DATA_BYTES];
+    signal is_prev_starting_index[DATA_BYTES];
+    value_starting_index[0] <== 0;
+    is_zero_mask[0] <== IsZero()(mask[0]);
+    for (var i=1 ; i<DATA_BYTES-1 ; i++) {
+        is_zero_mask[i] <== IsZero()(mask[i]);
+        is_prev_starting_index[i] <== IsZero()(value_starting_index[i-1]);
+        value_starting_index[i] <== value_starting_index[i-1] + i * (1-is_zero_mask[i]) * is_prev_starting_index[i];
+    }
+}
+
+template ExtractStringMultiDepthNested(DATA_BYTES, MAX_STACK_HEIGHT, keyLen1, depth1, keyLen2, depth2, index3, depth3, index4, depth4, maxValueLen) {
+    signal input data[DATA_BYTES];
+    signal input key1[keyLen1];
+    signal input key2[keyLen2];
+
+    signal output value[maxValueLen];
+
+    signal value_starting_index[DATA_BYTES];
+
+    value_starting_index <== ExtractMultiDepthNestedObject(DATA_BYTES, MAX_STACK_HEIGHT, keyLen1, depth1, keyLen2, depth2, index3, depth3, index4, depth4, maxValueLen)(data, key1, key2);
+
+    log("value_starting_index", value_starting_index[DATA_BYTES-2]);
+    // TODO: why +1 not required here,when required on all other string implss?
+    value <== SelectSubArray(DATA_BYTES, maxValueLen)(data, value_starting_index[DATA_BYTES-2], maxValueLen);
+
+    for (var i=0 ; i<maxValueLen; i++) {
+        log("value[",i,"]=", value[i]);
+    }
 }
\ No newline at end of file
diff --git a/circuits/test/fetcher.test.ts b/circuits/test/fetcher.test.ts
index 0965f08..2280a4e 100644
--- a/circuits/test/fetcher.test.ts
+++ b/circuits/test/fetcher.test.ts
@@ -13,10 +13,9 @@ function readInputFile(filename: string, key: any[]): [number[], number[][], num
     let keyUnicode: number[][] = [];
     for (let i = 0; i < key.length; i++) {
         keyUnicode[i] = [];
-        let key_string = key.toString();
-        for (let j = 0; j < key_string[i].length; j++) {
-
-            keyUnicode[i].push(key_string[i].charCodeAt(j));
+        let key_string = key[i].toString();
+        for (let j = 0; j < key_string.length; j++) {
+            keyUnicode[i].push(key_string.charCodeAt(j));
         }
     }
 
@@ -95,7 +94,7 @@ describe("ExtractValue", () => {
             });
             console.log("#constraints:", await circuit.getConstraintCount());
 
-            await circuit.expectPass({ data: input, key: keyUnicode }, { value: output });
+            await circuit.expectPass({ data: input, key: keyUnicode[0] }, { value: output });
         }
     });
 });
@@ -109,7 +108,7 @@ describe("ExtractValueMultiDepth", () => {
         circuit = await circomkit.WitnessTester(`Extract`, {
             file: "circuits/fetcher",
             template: "ExtractStringMultiDepth",
-            params: [input.length, 3, 1, 1, 1, 2, 1],
+            params: [input.length, 3, 1, 0, 1, 1, 1],
         });
         console.log("#constraints:", await circuit.getConstraintCount());
 
@@ -118,6 +117,22 @@ describe("ExtractValueMultiDepth", () => {
         let [input1, keyUnicode1, output1] = readInputFile("value_object.json", ["e", "f"]);
         await circuit.expectPass({ data: input1, key1: keyUnicode1[0], key2: keyUnicode1[1] }, { value: output1 });
     });
+
+    it("value_array_object: {\"a\":[{\"b\":[1,4]},{\"c\":\"b\"}]}", async () => {
+        let index_0 = 0;
+        let index_1 = 0;
+        let [input, keyUnicode, output] = readInputFile("value_array_object.json", ["a", index_0, "b", index_1]);
+        // console.log(`input: ${input}, key: ${keyUnicode}, output: ${output}`);
+
+        circuit = await circomkit.WitnessTester(`Extract`, {
+            file: "circuits/fetcher",
+            template: "ExtractStringMultiDepthNested",
+            params: [input.length, 4, 1, 0, 1, 2, index_0, 1, index_1, 3, 1],
+        });
+        console.log("#constraints:", await circuit.getConstraintCount());
+
+        await circuit.expectPass({ data: input, key1: keyUnicode[0], key2: keyUnicode[2] }, { value: output });
+    });
 });
 
 describe("ExtractValueMultiDepth", () => {
@@ -127,7 +142,7 @@ describe("ExtractValueMultiDepth", () => {
         let index_0 = 1;
         let index_1 = 0;
         let [input, keyUnicode, output] = readInputFile("value_array_nested.json", ["a", index_0, index_1]);
-        console.log(input, keyUnicode, output);
+
         circuit = await circomkit.WitnessTester(`Extract`, {
             file: "circuits/fetcher",
             template: "ExtractNestedArray",
diff --git a/json_examples/test/value_array_object.json b/json_examples/test/value_array_object.json
index 2bbaec9..3a0a544 100644
--- a/json_examples/test/value_array_object.json
+++ b/json_examples/test/value_array_object.json
@@ -1 +1 @@
-{"a":[{"b":5},{"c":"b"}]}
\ No newline at end of file
+{"a":[{"b":[1,4]},{"c":"b"}]}
\ No newline at end of file